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

examples/python/common.py

@@ -50,7 +50,13 @@ 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_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
 
@@ -59,14 +65,16 @@ 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
     """
     # Create config with EP
     # - 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, ep_path=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, ep_options={}, search_options=search_config, ep_path=args.ep_path
+    )
 
     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, ep_path=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, ep_path=args.ep_path)
     model = og.Model(config)
     if args.verbose:
         print("Model loaded")

src/models/model_type.h

@@ -65,4 +65,4 @@ struct ModelType {
   }
 };
 
-}  // namespace Generators
+}  // namespace Generators

src/models/recurrent_state.cpp

@@ -88,22 +88,32 @@ RecurrentState::RecurrentState(State& state)
 
   const int num_layers = static_cast<int>(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 {
+    throw std::runtime_error(
+        "RecurrentState requires past_present_share_buffer=true. "
+        "Set past_present_share_buffer to true in genai_config.json.");
   }
-
-  // Zero-initialize past and present states
-  ZeroStates(pasts_);
-  ZeroStates(presents_);
 }
 
 void RecurrentState::Add() {
@@ -114,7 +124,10 @@ void RecurrentState::Add() {
 
   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());
+    // 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,14 +141,18 @@ 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]));
     }
   }
 }
 
 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<int>(layer_indices_.size());
 
@@ -169,14 +186,20 @@ void RecurrentState::RewindTo(size_t index) {
     return;
   }
 
-  const int num_layers = static_cast<int>(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.
   ZeroStates(pasts_);
   ZeroStates(presents_);
 
   // Re-bind state pointers (swap may have changed which OrtValue is past vs present)
+  const int num_layers = static_cast<int>(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();

src/models/recurrent_state.h

@@ -47,6 +47,8 @@ struct RecurrentState {
 
   std::vector<int64_t> conv_shape_;
   std::vector<int64_t> recurrent_shape_;
+
+  bool past_present_share_buffer_{};
 };
 
 // Factory: returns nullptr if no recurrent layers are found in the session.