Enable CUDA graph capture for CUDA EP to improve decode throughput

.gitignore

@@ -31,6 +31,7 @@ examples/csharp/ModelChat/models
 
 !test/test_models/hf-internal-testing/
 !test/test_models/hf-internal-testing/tiny-random-gpt2*/*.onnx
+!test/test_models/hf-internal-testing/tiny-qwen35-cuda/*.onnx
 !test/test_models/grammars/
 !test/test_models/qwen-vision-preprocessing/
 !test/test_models/qwen-vision-preprocessing/*.onnx

src/config.cpp

@@ -1396,12 +1396,12 @@ bool IsGraphCaptureEnabled(const Config::SessionOptions& session_options) {
                                                });
     if (provider_options != session_options.provider_options.end()) {
       if (provider_options->name == "cuda") {
-        // Graph Capture is currently broken for CUDA
         for (const auto& value : provider_options->options) {
           if (value.first == "enable_cuda_graph" && value.second == "1") {
-            throw std::runtime_error("Graph Capture is currently unsupported for CUDA");
+            return true;
           }
         }
+        return false;
       } else if (provider_options->name == "DML") {
         return true;
       } else if (provider_options->name == "WebGPU") {

src/cuda/session_options.cpp

@@ -84,6 +84,11 @@ DeviceInterface* AppendExecutionProvider(OrtSessionOptions& session_options,
                                          bool /*disable_graph_capture*/) {
   auto device = GetDeviceInterface(DeviceType::CUDA);
   AddCudaStreamConfig(session_options, device);
+
+  // CUDA EP honors enable_cuda_graph from the config as-is. Non-decoder
+  // sessions (vision, embedding, speech) should set enable_cuda_graph=0
+  // in their own session_options in genai_config.json.
+
   // Try pre-registered plugin path first
   if (!AppendExecutionProviderV2(session_options, provider_options,
                                  DeviceType::CUDA, "CUDAExecutionProvider")) {

src/models/recurrent_state.cpp

@@ -119,16 +119,41 @@ void RecurrentState::Add() {
     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());
-  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();
+
+  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();
+    }
   }
 }
 
@@ -147,7 +172,7 @@ void RecurrentState::RewindTo(size_t index) {
   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 CUDA graphs.
+  // device pointers and avoid invalidating captured graphs.
   ZeroStates(pasts_);
   ZeroStates(presents_);
 

src/models/recurrent_state.h

@@ -30,6 +30,11 @@ struct RecurrentState {
   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_;

test/create_gqa_model.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""Create a tiny model with GQA for testing CUDA graph capture.
+
+Generates a qwen2-type LLM (decoder-only) that supports
+past_present_share_buffer and CUDA graph capture.
+
+The decoder model is run through ORT's transformer optimizer to ensure proper
+CUDA EP partitioning for graph capture compatibility.
+
+Usage:
+  python create_gqa_model.py [--output_dir OUTPUT_DIR]
+"""
+
+import argparse
+import json
+import os
+
+import numpy as np
+import onnx
+from onnx import TensorProto, helper, numpy_helper
+from onnxruntime.transformers.optimizer import optimize_model
+
+# Tiny model config
+VOCAB_SIZE = 1000
+HIDDEN_SIZE = 64
+NUM_HEADS = 4
+NUM_KV_HEADS = 2
+HEAD_SIZE = HIDDEN_SIZE // NUM_HEADS  # 16
+NUM_LAYERS = 2
+INTERMEDIATE_SIZE = 128
+MAX_SEQ_LEN = 128
+
+
+def rand_init(name, shape, dtype=np.float16, scale=0.02):
+    np.random.seed(hash(name) % (2**31))
+    return numpy_helper.from_array((np.random.randn(*shape) * scale).astype(dtype), name=name)
+
+
+def zeros_init(name, shape, dtype=np.float16):
+    return numpy_helper.from_array(np.zeros(shape, dtype=dtype), name=name)
+
+
+def create_decoder(output_dir):
+    """Create and optimize a GQA decoder model with input_ids."""
+    onnx_dtype = TensorProto.FLOAT16
+    ms = "com.microsoft"
+    nodes, inits = [], []
+
+    # Embedding table
+    inits.append(rand_init("embed_tokens.weight", [VOCAB_SIZE, HIDDEN_SIZE]))
+
+    graph_inputs = [
+        helper.make_tensor_value_info("input_ids", TensorProto.INT32, ["batch_size", "seq_len"]),
+        helper.make_tensor_value_info("attention_mask", TensorProto.INT64, ["batch_size", "total_seq_len"]),
+        helper.make_tensor_value_info("position_ids", TensorProto.INT64, ["batch_size", "seq_len"]),
+    ]
+    graph_outputs = []
+
+    for i in range(NUM_LAYERS):
+        graph_inputs.append(helper.make_tensor_value_info(
+            f"past_key_values.{i}.key", onnx_dtype, ["batch_size", NUM_KV_HEADS, "past_seq_len", HEAD_SIZE]))
+        graph_inputs.append(helper.make_tensor_value_info(
+            f"past_key_values.{i}.value", onnx_dtype, ["batch_size", NUM_KV_HEADS, "past_seq_len", HEAD_SIZE]))
+
+    # Derive seqlens_k and total_sequence_length from attention_mask
+    # These are small CPU-side computations that ORT handles outside graph capture
+    nodes.append(helper.make_node("Cast", ["attention_mask"], ["attn_mask_i32"], to=TensorProto.INT32))
+    inits.append(numpy_helper.from_array(np.array([1], dtype=np.int64), "reduce_axes"))
+    nodes.append(helper.make_node("ReduceSum", ["attn_mask_i32", "reduce_axes"], ["seqlens_k_p1"],
+                                  keepdims=0, noop_with_empty_axes=0))
+    inits.append(numpy_helper.from_array(np.array([1], dtype=np.int32), "const_one"))
+    nodes.append(helper.make_node("Sub", ["seqlens_k_p1", "const_one"], ["seqlens_k"]))
+    nodes.append(helper.make_node("Shape", ["attention_mask"], ["mask_shape"]))
+    inits.append(numpy_helper.from_array(np.array(1, dtype=np.int64), "gather_idx"))
+    nodes.append(helper.make_node("Gather", ["mask_shape", "gather_idx"], ["total_sl_i64"], axis=0))
+    nodes.append(helper.make_node("Cast", ["total_sl_i64"], ["total_sl"], to=TensorProto.INT32))
+
+    # Embedding lookup: input_ids -> inputs_embeds
+    nodes.append(helper.make_node("Gather", ["embed_tokens.weight", "input_ids"], ["inputs_embeds"], axis=0))
+
+    prev = "inputs_embeds"
+
+    for li in range(NUM_LAYERS):
+        p = f"l{li}"
+
+        # RMSNorm (using SimplifiedLayerNormalization which is in com.microsoft)
+        # Use standard LayerNormalization instead for compatibility
+        inits.append(rand_init(f"{p}.ln1.w", [HIDDEN_SIZE]))
+        inits.append(zeros_init(f"{p}.ln1.b", [HIDDEN_SIZE]))
+        nodes.append(helper.make_node("LayerNormalization",
+                                      [prev, f"{p}.ln1.w", f"{p}.ln1.b"], [f"{p}.ln1"],
+                                      epsilon=1e-5, axis=-1))
+
+        # Q, K, V projections (separate, not packed)
+        inits.append(rand_init(f"{p}.q.w", [HIDDEN_SIZE, NUM_HEADS * HEAD_SIZE]))
+        inits.append(rand_init(f"{p}.k.w", [HIDDEN_SIZE, NUM_KV_HEADS * HEAD_SIZE]))
+        inits.append(rand_init(f"{p}.v.w", [HIDDEN_SIZE, NUM_KV_HEADS * HEAD_SIZE]))
+        nodes.append(helper.make_node("MatMul", [f"{p}.ln1", f"{p}.q.w"], [f"{p}.q"]))
+        nodes.append(helper.make_node("MatMul", [f"{p}.ln1", f"{p}.k.w"], [f"{p}.k"]))
+        nodes.append(helper.make_node("MatMul", [f"{p}.ln1", f"{p}.v.w"], [f"{p}.v"]))
+
+        # GQA with separate Q, K, V
+        nodes.append(helper.make_node("GroupQueryAttention",
+            [f"{p}.q", f"{p}.k", f"{p}.v",
+             f"past_key_values.{li}.key", f"past_key_values.{li}.value",
+             "seqlens_k", "total_sl",
+             "", "",  # cos/sin cache
+             ],
+            [f"{p}.attn", f"present.{li}.key", f"present.{li}.value"],
+            domain=ms, name=f"GQA_{li}",
+            num_heads=NUM_HEADS, kv_num_heads=NUM_KV_HEADS))
+
+        # Output projection + residual
+        inits.append(rand_init(f"{p}.o.w", [NUM_HEADS * HEAD_SIZE, HIDDEN_SIZE]))
+        nodes.append(helper.make_node("MatMul", [f"{p}.attn", f"{p}.o.w"], [f"{p}.o"]))
+        nodes.append(helper.make_node("Add", [prev, f"{p}.o"], [f"{p}.r1"]))
+
+        # Post-attn norm
+        inits.append(rand_init(f"{p}.ln2.w", [HIDDEN_SIZE]))
+        inits.append(zeros_init(f"{p}.ln2.b", [HIDDEN_SIZE]))
+        nodes.append(helper.make_node("LayerNormalization",
+                                      [f"{p}.r1", f"{p}.ln2.w", f"{p}.ln2.b"], [f"{p}.ln2"],
+                                      epsilon=1e-5, axis=-1))
+
+        # MLP
+        inits.append(rand_init(f"{p}.up.w", [HIDDEN_SIZE, INTERMEDIATE_SIZE]))
+        inits.append(rand_init(f"{p}.dn.w", [INTERMEDIATE_SIZE, HIDDEN_SIZE]))
+        nodes.append(helper.make_node("MatMul", [f"{p}.ln2", f"{p}.up.w"], [f"{p}.up"]))
+        nodes.append(helper.make_node("Relu", [f"{p}.up"], [f"{p}.act"]))
+        nodes.append(helper.make_node("MatMul", [f"{p}.act", f"{p}.dn.w"], [f"{p}.dn"]))
+
+        prev = f"{p}.r2"
+        nodes.append(helper.make_node("Add", [f"{p}.r1", f"{p}.dn"], [prev]))
+
+        for kv in ("key", "value"):
+            graph_outputs.append(helper.make_tensor_value_info(
+                f"present.{li}.{kv}", onnx_dtype,
+                ["batch_size", NUM_KV_HEADS, "total_seq_len", HEAD_SIZE]))
+
+    # Final norm + LM head
+    inits.append(rand_init("norm.w", [HIDDEN_SIZE]))
+    inits.append(zeros_init("norm.b", [HIDDEN_SIZE]))
+    nodes.append(helper.make_node("LayerNormalization",
+                                  [prev, "norm.w", "norm.b"], ["final_ln"],
+                                  epsilon=1e-5, axis=-1))
+    inits.append(rand_init("lm_head.w", [HIDDEN_SIZE, VOCAB_SIZE]))
+    nodes.append(helper.make_node("MatMul", ["final_ln", "lm_head.w"], ["logits"]))
+    graph_outputs.insert(0, helper.make_tensor_value_info(
+        "logits", onnx_dtype, ["batch_size", "seq_len", VOCAB_SIZE]))
+
+    graph = helper.make_graph(nodes, "decoder", graph_inputs, graph_outputs, initializer=inits)
+    model = helper.make_model(graph, opset_imports=[
+        helper.make_opsetid("", 21), helper.make_opsetid(ms, 1)])
+    model.ir_version = 10
+    onnx.checker.check_model(model)
+
+    # Save unoptimized first
+    raw_path = os.path.join(output_dir, "text_raw.onnx")
+    onnx.save(model, raw_path)
+
+    # Optimize with ORT transformer optimizer for proper CUDA EP placement
+    opt = optimize_model(raw_path, model_type="gpt2", opt_level=0,
+                         num_heads=NUM_HEADS, hidden_size=HIDDEN_SIZE)
+    final_path = os.path.join(output_dir, "decoder.onnx")
+    opt.save_model_to_file(final_path)
+    os.remove(raw_path)
+    print(f"  Saved decoder -> {final_path}")
+
+
+def create_config(output_dir):
+    config = {
+        "model": {
+            "type": "qwen2",
+            "bos_token_id": 1, "eos_token_id": 2, "pad_token_id": 0,
+            "vocab_size": VOCAB_SIZE, "context_length": MAX_SEQ_LEN,
+            "decoder": {
+                "session_options": {
+                    "log_id": "onnxruntime-genai",
+                    "provider_options": [{"cuda": {"enable_cuda_graph": "1"}}]
+                },
+                "filename": "decoder.onnx",
+                "num_attention_heads": NUM_HEADS, "num_key_value_heads": NUM_KV_HEADS,
+                "head_size": HEAD_SIZE, "hidden_size": HIDDEN_SIZE,
+                "num_hidden_layers": NUM_LAYERS,
+                "inputs": {
+                    "input_ids": "input_ids", "attention_mask": "attention_mask",
+                    "position_ids": "position_ids",
+                    "past_key_names": "past_key_values.%d.key",
+                    "past_value_names": "past_key_values.%d.value"
+                },
+                "outputs": {
+                    "logits": "logits",
+                    "present_key_names": "present.%d.key",
+                    "present_value_names": "present.%d.value"
+                }
+            }
+        },
+        "search": {"past_present_share_buffer": True}
+    }
+    path = os.path.join(output_dir, "genai_config.json")
+    with open(path, "w") as f:
+        json.dump(config, f, indent=4)
+    print(f"  Saved config -> {path}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_dir",
+                        default=os.path.join(os.path.dirname(__file__),
+                                             "test_models", "hf-internal-testing", "tiny-qwen35-cuda"))
+    args = parser.parse_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+    print(f"Creating tiny GQA test model in {args.output_dir}")
+    create_decoder(args.output_dir)
+    create_config(args.output_dir)
+    print("Done!")
+
+
+if __name__ == "__main__":
+    main()

test/model_tests.cpp

@@ -241,6 +241,75 @@ TEST(ModelTests, BeamSearchGptCuda) {
   for (auto model_path : c_tiny_gpt2_model_paths)
     Test_BeamSearch_Gpt_Cuda(model_path.first, model_path.second);
 }
+
+// Tiny GQA model for CUDA graph capture tests (GQA + past_present_share_buffer).
+// Generated by test/create_gqa_model.py
+static const char* c_cuda_graph_model_path = MODEL_PATH "hf-internal-testing/tiny-qwen35-cuda";
+
+// Test that CUDA graph capture produces correct outputs for greedy search.
+// Validates generation completes without CUDA errors under graph capture.
+TEST(ModelTests, GreedySearchCudaGraphCapture) {
+  if (!std::filesystem::exists(c_cuda_graph_model_path)) {
+    GTEST_SKIP() << "CUDA graph capture model not available at: " << c_cuda_graph_model_path;
+  }
+
+  std::vector<int32_t> input_ids{10, 20, 30, 40};
+  int max_length = 10;
+
+  auto model = OgaModel::Create(c_cuda_graph_model_path);
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetSearchOption("max_length", max_length);
+  params->SetSearchOption("batch_size", 1);
+
+  auto generator = OgaGenerator::Create(*model, *params);
+  generator->AppendTokens(input_ids);
+
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  auto seq = generator->GetSequence(0);
+  EXPECT_EQ(seq.size(), static_cast<size_t>(max_length));
+}
+
+// Test continuous decoding (RewindTo) with CUDA graph capture enabled.
+TEST(ModelTests, ContinuousDecodingCudaGraphCapture) {
+  if (!std::filesystem::exists(c_cuda_graph_model_path)) {
+    GTEST_SKIP() << "CUDA graph capture model not available at: " << c_cuda_graph_model_path;
+  }
+
+  std::vector<int32_t> input_ids{10, 20, 30, 40};
+  int max_length = 10;
+
+  auto model = OgaModel::Create(c_cuda_graph_model_path);
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetSearchOption("max_length", max_length);
+  params->SetSearchOption("batch_size", 1);
+
+  auto generator = OgaGenerator::Create(*model, *params);
+  generator->AppendTokens(input_ids);
+
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  auto seq1 = generator->GetSequence(0);
+  std::vector<int32_t> first_output(seq1.begin(), seq1.end());
+
+  // Rewind and regenerate — should produce same output
+  size_t rewind_pos = input_ids.size() - 1;
+  generator->RewindTo(rewind_pos);
+  std::vector<int32_t> resume_ids(input_ids.begin() + rewind_pos, input_ids.end());
+  generator->AppendTokens(resume_ids);
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  auto seq2 = generator->GetSequence(0);
+  ASSERT_EQ(seq2.size(), first_output.size());
+  EXPECT_TRUE(0 == std::memcmp(first_output.data(), seq2.data(), seq2.size() * sizeof(int32_t)))
+      << "Continuous decoding output differs after RewindTo with graph capture";
+}
 #endif
 
 // NvTensorRT test cases using Phi3 models