From eee007c4b2a4bb83e1943f68a7090feafde8ce48 Mon Sep 17 00:00:00 2001 From: yenshiw Date: Mon, 18 May 2026 22:09:16 -0700 Subject: [PATCH 01/11] Enable Qwen3.5 TRT-RTX shared-buffer inference --- examples/python/common.py | 1 + src/models/model_type.h | 2 +- src/models/recurrent_state.cpp | 59 +++++++++++++++++++++------ src/models/recurrent_state.h | 2 + src/python/py/models/builders/base.py | 9 +++- 5 files changed, 57 insertions(+), 16 deletions(-) diff --git a/examples/python/common.py b/examples/python/common.py index a9002a3522..1b88c87ade 100644 --- a/examples/python/common.py +++ b/examples/python/common.py @@ -557,6 +557,7 @@ def get_ep_args(parser: argparse.ArgumentParser) -> None: "cpu", # CPU EP "cuda", # GenAI canonical name for CUDA EP "CUDAExecutionProvider", # CUDA EP + "NvTensorRtRtx", # GenAI canonical name for Nvidia IHV EP "NvTensorRTRTXExecutionProvider", # Nvidia IHV EP "OpenVINOExecutionProvider", # Intel IHV EP "QNNExecutionProvider", # Qualcomm IHV EP diff --git a/src/models/model_type.h b/src/models/model_type.h index 51efb901c3..3c5063f665 100644 --- a/src/models/model_type.h +++ b/src/models/model_type.h @@ -65,4 +65,4 @@ struct ModelType { } }; -} // namespace Generators \ No newline at end of file +} // namespace Generators diff --git a/src/models/recurrent_state.cpp b/src/models/recurrent_state.cpp index 7d2f4104c2..7cf1c9a63e 100644 --- a/src/models/recurrent_state.cpp +++ b/src/models/recurrent_state.cpp @@ -88,22 +88,42 @@ RecurrentState::RecurrentState(State& state) const int num_layers = static_cast(layer_indices_.size()); - pasts_.resize(num_layers * 2); + past_present_share_buffer_ = state_.params_->IsPastPresentShareBufferEnabled(model_.config_->model.type); + if (g_log.enabled && past_present_share_buffer_) { + Log("info", "RecurrentState: using shared past/present buffers"); + } + presents_.reserve(num_layers * 2); auto& allocator = model_.p_device_kvcache_->GetAllocator(); - 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_); + if (past_present_share_buffer_) { + // 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) { + presents_.push_back(OrtValue::CreateTensor(allocator, conv_shape_, conv_type_)); + presents_.push_back(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_)); - } + ZeroStates(presents_); + } else { + pasts_.resize(num_layers * 2); - // Zero-initialize past and present states - ZeroStates(pasts_); - ZeroStates(presents_); + 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() { @@ -114,7 +134,10 @@ void RecurrentState::Add() { 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()); + // In shared-buffer mode the same OrtValue is intentionally registered as + // input and output. Non-shared mode keeps the older ping-pong buffers. + auto* past = past_present_share_buffer_ ? presents_[i].get() : pasts_[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()); @@ -128,7 +151,8 @@ void RecurrentState::Add() { 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])); + auto& past = past_present_share_buffer_ ? presents_[i] : pasts_[i]; + past_byte_spans_.push_back(ByteWrapTensor(device, *past)); present_byte_spans_.push_back(ByteWrapTensor(device, *presents_[i])); } } @@ -136,6 +160,9 @@ void RecurrentState::Add() { void RecurrentState::Update() { if (layer_indices_.empty()) return; + // Shared mode updates state contents in place, so swapping would only change + // the captured input/output addresses and defeat graph reuse. + if (past_present_share_buffer_) return; const int num_layers = static_cast(layer_indices_.size()); @@ -169,7 +196,12 @@ void RecurrentState::RewindTo(size_t index) { return; } - const int num_layers = static_cast(layer_indices_.size()); + if (past_present_share_buffer_) { + // Shared recurrent states keep stable input/output pointers for graph replay. + // Reset the state contents in place without rebinding. + ZeroStates(presents_); + return; + } // Zero existing buffers in-place instead of reallocating, to preserve // device pointers and avoid invalidating captured graphs. @@ -177,6 +209,7 @@ void RecurrentState::RewindTo(size_t index) { ZeroStates(presents_); // Re-bind state pointers (swap may have changed which OrtValue is past vs present) + const int num_layers = static_cast(layer_indices_.size()); 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(); diff --git a/src/models/recurrent_state.h b/src/models/recurrent_state.h index dc7d716a4c..fbc904f930 100644 --- a/src/models/recurrent_state.h +++ b/src/models/recurrent_state.h @@ -47,6 +47,8 @@ struct RecurrentState { std::vector conv_shape_; std::vector recurrent_shape_; + + bool past_present_share_buffer_{}; }; // Factory: returns nullptr if no recurrent layers are found in the session. diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index cc0a44bcc7..732934ac32 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -1625,6 +1625,11 @@ def _make_layernorm_op(self, layer_id, layernorm, skip, simple, location): root_input = inputs[0] skip_input = inputs[1] if skip else None + # Cast insertion can redirect SkipLayerNorm's fourth output to a casted + # value. Pass that redirected name to the primitive so QDQ export does + # not register two producers for the same original output_3 value. + primitive_output_3 = outputs[3] if skip and not self.layernorm_attrs["last_layernorm"] else None + if op_type == "SimplifiedLayerNormalization": self._make_simplified_layer_norm( name, @@ -1641,7 +1646,7 @@ def _make_layernorm_op(self, layer_id, layernorm, skip, simple, location): skip_input, weight, outputs[0], - output_3, + primitive_output_3, new_io_dtype, shape=["batch_size", "sequence_length", self.hidden_size], ) @@ -1653,7 +1658,7 @@ def _make_layernorm_op(self, layer_id, layernorm, skip, simple, location): weight, bias, outputs[0], - output_3, + primitive_output_3, new_io_dtype, shape=["batch_size", "sequence_length", self.hidden_size], ) From 1feee99d0d434ea6cae037966a22e8e3a88f56d8 Mon Sep 17 00:00:00 2001 From: yenshiw Date: Wed, 20 May 2026 11:32:41 -0700 Subject: [PATCH 02/11] Remove unnecessary layernorm output redirect --- src/python/py/models/builders/base.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index 732934ac32..cc0a44bcc7 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -1625,11 +1625,6 @@ def _make_layernorm_op(self, layer_id, layernorm, skip, simple, location): root_input = inputs[0] skip_input = inputs[1] if skip else None - # Cast insertion can redirect SkipLayerNorm's fourth output to a casted - # value. Pass that redirected name to the primitive so QDQ export does - # not register two producers for the same original output_3 value. - primitive_output_3 = outputs[3] if skip and not self.layernorm_attrs["last_layernorm"] else None - if op_type == "SimplifiedLayerNormalization": self._make_simplified_layer_norm( name, @@ -1646,7 +1641,7 @@ def _make_layernorm_op(self, layer_id, layernorm, skip, simple, location): skip_input, weight, outputs[0], - primitive_output_3, + output_3, new_io_dtype, shape=["batch_size", "sequence_length", self.hidden_size], ) @@ -1658,7 +1653,7 @@ def _make_layernorm_op(self, layer_id, layernorm, skip, simple, location): weight, bias, outputs[0], - primitive_output_3, + output_3, new_io_dtype, shape=["batch_size", "sequence_length", self.hidden_size], ) From 3c3deacc75c07aff9d1d7c6fd9a03e7f971a6f67 Mon Sep 17 00:00:00 2001 From: yenshiw Date: Wed, 20 May 2026 11:48:25 -0700 Subject: [PATCH 03/11] Remove TRT-RTX example EP alias --- examples/python/common.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/python/common.py b/examples/python/common.py index 1b88c87ade..a9002a3522 100644 --- a/examples/python/common.py +++ b/examples/python/common.py @@ -557,7 +557,6 @@ def get_ep_args(parser: argparse.ArgumentParser) -> None: "cpu", # CPU EP "cuda", # GenAI canonical name for CUDA EP "CUDAExecutionProvider", # CUDA EP - "NvTensorRtRtx", # GenAI canonical name for Nvidia IHV EP "NvTensorRTRTXExecutionProvider", # Nvidia IHV EP "OpenVINOExecutionProvider", # Intel IHV EP "QNNExecutionProvider", # Qualcomm IHV EP From 31d1c911ee462f50b00d9508d0266068a2225079 Mon Sep 17 00:00:00 2001 From: yenshiw Date: Wed, 20 May 2026 12:48:21 -0700 Subject: [PATCH 04/11] Map TRT-RTX GenAI EP name for example registration --- examples/python/common.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/python/common.py b/examples/python/common.py index a9002a3522..018be3d030 100644 --- a/examples/python/common.py +++ b/examples/python/common.py @@ -45,7 +45,10 @@ def register_ep(ep: str, ep_path: str, use_winml: bool) -> None: except Exception as e: print(f"Failed to register WinML execution providers: {e}") elif ep_path: - og.register_execution_provider_library(ep, ep_path) + if ep == "NvTensorRtRtx": + og.register_execution_provider_library("NvTensorRTRTXExecutionProvider", ep_path) + else: + og.register_execution_provider_library(ep, ep_path) print(f"Registered {ep} from {ep_path} successfully!") @@ -557,6 +560,7 @@ def get_ep_args(parser: argparse.ArgumentParser) -> None: "cpu", # CPU EP "cuda", # GenAI canonical name for CUDA EP "CUDAExecutionProvider", # CUDA EP + "NvTensorRtRtx", # GenAI canonical name for Nvidia IHV EP "NvTensorRTRTXExecutionProvider", # Nvidia IHV EP "OpenVINOExecutionProvider", # Intel IHV EP "QNNExecutionProvider", # Qualcomm IHV EP From f485a9d54115228c133a0d868f0b856507415593 Mon Sep 17 00:00:00 2001 From: yenshiw Date: Wed, 20 May 2026 16:03:59 -0700 Subject: [PATCH 05/11] Preserve configured providers with external EP libraries --- examples/python/common.py | 18 +++++++++++------- examples/python/model-chat.py | 2 +- examples/python/model-generate.py | 4 +++- examples/python/model-mm.py | 2 +- examples/python/model-qa.py | 2 +- 5 files changed, 17 insertions(+), 11 deletions(-) diff --git a/examples/python/common.py b/examples/python/common.py index 018be3d030..994fc1496d 100644 --- a/examples/python/common.py +++ b/examples/python/common.py @@ -45,15 +45,18 @@ def register_ep(ep: str, ep_path: str, use_winml: bool) -> None: except Exception as e: print(f"Failed to register WinML execution providers: {e}") elif ep_path: - if ep == "NvTensorRtRtx": - og.register_execution_provider_library("NvTensorRTRTXExecutionProvider", ep_path) - else: - og.register_execution_provider_library(ep, ep_path) + og.register_execution_provider_library(ep, ep_path) 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_options: dict[str, str] = {}, + search_options: dict[str, int] = {}, + ep_path: str = "", +) -> og.Config: """ Get og.Config object and set EP-specific and search-specific options inside it @@ -62,6 +65,8 @@ def get_config(path: str, ep: str, ep_options: dict[str, str] = {}, search_optio ep (str): Name of execution provider to set 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 + ep_path (str): Path to an external execution provider library. If set, the + registered library is used and providers from the GenAI config are preserved. Returns: og.Config: ORT GenAI config object with all options set """ @@ -69,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}") @@ -560,7 +565,6 @@ def get_ep_args(parser: argparse.ArgumentParser) -> None: "cpu", # CPU EP "cuda", # GenAI canonical name for CUDA EP "CUDAExecutionProvider", # CUDA EP - "NvTensorRtRtx", # GenAI canonical name for Nvidia IHV EP "NvTensorRTRTXExecutionProvider", # Nvidia IHV EP "OpenVINOExecutionProvider", # Intel IHV EP "QNNExecutionProvider", # Qualcomm IHV EP diff --git a/examples/python/model-chat.py b/examples/python/model-chat.py index 0edad43198..ff0f03dac8 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, ep_path=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..e3172815ff 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, ep_options={}, search_options=search_config, ep_path=args.ep_path + ) model = og.Model(config) if args.verbose: diff --git a/examples/python/model-mm.py b/examples/python/model-mm.py index 5a938cd670..516adcf726 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, ep_path=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..941a9302a0 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, ep_path=args.ep_path) model = og.Model(config) if args.verbose: print("Model loaded") From d70abd4fc9e23c8691cbee7095b2aa9aefbbc783 Mon Sep 17 00:00:00 2001 From: yenshiw Date: Thu, 21 May 2026 14:41:22 -0700 Subject: [PATCH 06/11] Require shared buffers for recurrent state --- src/models/recurrent_state.cpp | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/src/models/recurrent_state.cpp b/src/models/recurrent_state.cpp index 7cf1c9a63e..b9824796ea 100644 --- a/src/models/recurrent_state.cpp +++ b/src/models/recurrent_state.cpp @@ -110,19 +110,9 @@ RecurrentState::RecurrentState(State& state) ZeroStates(presents_); } else { - pasts_.resize(num_layers * 2); - - 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_); + throw std::runtime_error( + "RecurrentState requires past_present_share_buffer=true. " + "Set past_present_share_buffer to true in genai_config.json."); } } From 8c9b438418764592ab2aa5ab68bd87fa52c7750a Mon Sep 17 00:00:00 2001 From: yenshiw Date: Fri, 22 May 2026 02:07:26 -0700 Subject: [PATCH 07/11] Simplify recurrent shared buffer handling --- src/models/recurrent_state.cpp | 102 +++++++-------------------------- src/models/recurrent_state.h | 11 ---- 2 files changed, 22 insertions(+), 91 deletions(-) diff --git a/src/models/recurrent_state.cpp b/src/models/recurrent_state.cpp index b9824796ea..976ce34156 100644 --- a/src/models/recurrent_state.cpp +++ b/src/models/recurrent_state.cpp @@ -88,122 +88,64 @@ RecurrentState::RecurrentState(State& state) const int num_layers = static_cast(layer_indices_.size()); - past_present_share_buffer_ = state_.params_->IsPastPresentShareBufferEnabled(model_.config_->model.type); - if (g_log.enabled && past_present_share_buffer_) { + const bool past_present_share_buffer = state_.params_->IsPastPresentShareBufferEnabled(model_.config_->model.type); + if (g_log.enabled && past_present_share_buffer) { Log("info", "RecurrentState: using shared past/present buffers"); } + 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(); - if (past_present_share_buffer_) { - // 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) { - presents_.push_back(OrtValue::CreateTensor(allocator, conv_shape_, conv_type_)); - presents_.push_back(OrtValue::CreateTensor(allocator, recurrent_shape_, recurrent_type_)); - } - - ZeroStates(presents_); - } else { - throw std::runtime_error( - "RecurrentState requires past_present_share_buffer=true. " - "Set past_present_share_buffer to true in genai_config.json."); + // 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) { + presents_.push_back(OrtValue::CreateTensor(allocator, conv_shape_, conv_type_)); + presents_.push_back(OrtValue::CreateTensor(allocator, recurrent_shape_, recurrent_type_)); } + + 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) { - // In shared-buffer mode the same OrtValue is intentionally registered as - // input and output. Non-shared mode keeps the older ping-pong buffers. - auto* past = past_present_share_buffer_ ? presents_[i].get() : 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) { - auto& past = past_present_share_buffer_ ? presents_[i] : pasts_[i]; - past_byte_spans_.push_back(ByteWrapTensor(device, *past)); - present_byte_spans_.push_back(ByteWrapTensor(device, *presents_[i])); - } - } } void RecurrentState::Update() { - if (layer_indices_.empty()) return; - // Shared mode updates state contents in place, so swapping would only change - // the captured input/output addresses and defeat graph reuse. - if (past_present_share_buffer_) 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) { if (layer_indices_.empty()) return; if (index != 0) { - // Recurrent states cannot be partially rewound — they are compressed summaries + // Recurrent states cannot be partially rewound; they are compressed summaries // with no per-position history. Non-zero rewind is a no-op; the state remains unchanged. if (g_log.enabled) Log("warning", "RecurrentState::RewindTo(" + std::to_string(index) + ") is a no-op. Recurrent states cannot be partially rewound."); return; } - - if (past_present_share_buffer_) { - // Shared recurrent states keep stable input/output pointers for graph replay. - // Reset the state contents in place without rebinding. - ZeroStates(presents_); - return; - } - - // 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) - const int num_layers = static_cast(layer_indices_.size()); - 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 fbc904f930..e7c3c42734 100644 --- a/src/models/recurrent_state.h +++ b/src/models/recurrent_state.h @@ -27,28 +27,17 @@ 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_{}; std::vector conv_shape_; std::vector recurrent_shape_; - - bool past_present_share_buffer_{}; }; // Factory: returns nullptr if no recurrent layers are found in the session. From d934a98b67a1f6de1c3d60687d128af9b1871fb0 Mon Sep 17 00:00:00 2001 From: yenshiw Date: Fri, 22 May 2026 02:13:32 -0700 Subject: [PATCH 08/11] Remove recurrent shared-buffer info log --- src/models/recurrent_state.cpp | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/models/recurrent_state.cpp b/src/models/recurrent_state.cpp index 976ce34156..6be157dd72 100644 --- a/src/models/recurrent_state.cpp +++ b/src/models/recurrent_state.cpp @@ -89,9 +89,6 @@ RecurrentState::RecurrentState(State& state) const int num_layers = static_cast(layer_indices_.size()); const bool past_present_share_buffer = state_.params_->IsPastPresentShareBufferEnabled(model_.config_->model.type); - if (g_log.enabled && past_present_share_buffer) { - Log("info", "RecurrentState: using shared past/present buffers"); - } if (!past_present_share_buffer) { throw std::runtime_error( "RecurrentState requires past_present_share_buffer=true. " From 32a68ce5761f2d1ceb5b23d948cfb0cd1873fb27 Mon Sep 17 00:00:00 2001 From: yenshiw Date: Fri, 22 May 2026 02:15:17 -0700 Subject: [PATCH 09/11] Restore recurrent rewind comment wording --- src/models/recurrent_state.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/recurrent_state.cpp b/src/models/recurrent_state.cpp index 6be157dd72..8be92368a3 100644 --- a/src/models/recurrent_state.cpp +++ b/src/models/recurrent_state.cpp @@ -132,7 +132,7 @@ void RecurrentState::RewindTo(size_t index) { if (layer_indices_.empty()) return; if (index != 0) { - // Recurrent states cannot be partially rewound; they are compressed summaries + // Recurrent states cannot be partially rewound — they are compressed summaries // with no per-position history. Non-zero rewind is a no-op; the state remains unchanged. if (g_log.enabled) Log("warning", "RecurrentState::RewindTo(" + std::to_string(index) + From a1c01d1e876ffce58fd0752d5426622353167e24 Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Fri, 22 May 2026 19:21:52 +0000 Subject: [PATCH 10/11] Address PR comment --- examples/python/common.py | 6 +++--- examples/python/model-chat.py | 2 +- examples/python/model-generate.py | 2 +- examples/python/model-mm.py | 2 +- examples/python/model-qa.py | 2 +- examples/python/nemotron_speech.py | 2 +- 6 files changed, 8 insertions(+), 8 deletions(-) diff --git a/examples/python/common.py b/examples/python/common.py index 994fc1496d..1ebd1fee82 100644 --- a/examples/python/common.py +++ b/examples/python/common.py @@ -53,9 +53,9 @@ def register_ep(ep: str, ep_path: str, use_winml: bool) -> None: def get_config( path: str, ep: str, + ep_path: str | None, ep_options: dict[str, str] = {}, search_options: dict[str, int] = {}, - ep_path: str = "", ) -> og.Config: """ Get og.Config object and set EP-specific and search-specific options inside it @@ -63,10 +63,10 @@ def get_config( 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 - ep_path (str): Path to an external execution provider library. If set, the - registered library is used and providers from the GenAI config are preserved. Returns: og.Config: ORT GenAI config object with all options set """ diff --git a/examples/python/model-chat.py b/examples/python/model-chat.py index ff0f03dac8..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, ep_path=args.ep_path) + 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 e3172815ff..9d22564916 100644 --- a/examples/python/model-generate.py +++ b/examples/python/model-generate.py @@ -33,7 +33,7 @@ def main(args): 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, ep_path=args.ep_path + args.model_path, args.execution_provider, args.ep_path, ep_options={}, search_options=search_config ) model = og.Model(config) diff --git a/examples/python/model-mm.py b/examples/python/model-mm.py index 516adcf726..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, ep_path=args.ep_path) + 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 941a9302a0..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, ep_path=args.ep_path) + 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 ec03de7989..f098f883a1 100644 --- a/examples/python/nemotron_speech.py +++ b/examples/python/nemotron_speech.py @@ -53,7 +53,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) model = og.Model(config) processor = og.StreamingProcessor(model) From 50b4fdf2a97fdd4097c8458b4a216ddca99f1e8f Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Fri, 22 May 2026 19:23:52 +0000 Subject: [PATCH 11/11] Fix test --- test/test_models/qwen35-hybrid-preprocessing/genai_config.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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,