Enable Qwen3.5 TRT-RTX EP path with CUDA graph

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}")

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")

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:

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")

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")

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:

src/models/recurrent_state.cpp

@@ -88,73 +88,44 @@ RecurrentState::RecurrentState(State& state)
 
   const int num_layers = static_cast<int>(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<int>(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<int>(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<int>(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<std::unique_ptr<OrtValue>>& states) {

src/models/recurrent_state.h

@@ -27,21 +27,12 @@ struct RecurrentState {
   std::vector<int> layer_indices_;
 
   // Interleaved as [conv_0, recurrent_0, conv_1, recurrent_1, ...]
-  std::vector<std::unique_ptr<OrtValue>> pasts_;
   std::vector<std::unique_ptr<OrtValue>> presents_;
 
-  // Cached byte spans for graph-capture copy path (avoids recomputing
-  // tensor metadata on every decode step for fixed-shape tensors).
-  std::vector<DeviceSpan<uint8_t>> past_byte_spans_;
-  std::vector<DeviceSpan<uint8_t>> present_byte_spans_;
-
   // Kept alive for state_ const char* pointers
   std::vector<std::string> input_name_strings_;
   std::vector<std::string> output_name_strings_;
 
-  size_t input_index_{~0U};
-  size_t output_index_{~0U};
-
   ONNXTensorElementDataType conv_type_{};
   ONNXTensorElementDataType recurrent_type_{};
 

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,