Enable graph capture for WebGPU models and DML continuous decoding tests

.github/workflows/win-webgpu-x64-build.yml

@@ -51,7 +51,7 @@ jobs:
         shell: pwsh
         run: |
           # Use Foundry package which includes WebGPU DLLs (dxil.dll, dxcompiler.dll)
-          $FOUNDRY_VERSION = "1.25.0-dev-20260210-0905-b214734cba"
+          $FOUNDRY_VERSION = "1.26.0-dev-20260410-0804-ce91376bdf"
           Write-Host "Downloading OnnxRuntime.Foundry version: $FOUNDRY_VERSION"
           nuget install Microsoft.ML.OnnxRuntime.Foundry -version $FOUNDRY_VERSION -Source https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/nuget/v3/index.json -x -NonInteractive -ExcludeVersion -DependencyVersion Ignore
 

src/models/position_inputs.cpp

@@ -418,12 +418,8 @@ void DefaultPositionInputs::RewindMask(size_t index) {
 // Currently triggered by:
 //   - DML (always uses graph capture, see IsGraphCaptureEnabled in config.cpp)
 //   - WebGPU with enableGraphCapture=1 in provider options
+//   - CUDA with enable_cuda_graph=1 in provider options
 //   - NvTensorRtRtx with past-present shared buffers
-// Not yet using this path:
-//   - CUDA: graph capture is currently disabled in GenAI due to bugs
-//     (IsGraphCaptureEnabled throws for CUDA). Once re-enabled, RewindMask's
-//     static path will work for CUDA as well since it uses device-agnostic
-//     CpuSpan/CopyCpuToDevice.
 bool DefaultPositionInputs::ShouldUseStaticMaskHandling() const {
   return state_.params_->use_graph_capture ||
          (state_.params_->IsPastPresentShareBufferEnabled(model_.config_->model.type) &&

test/c_api_tests.cpp

@@ -1425,6 +1425,74 @@ TEST(CAPITests, RewindGraphCaptureNvTensorRtRtxCAPI) {
   EXPECT_TRUE(0 == std::memcmp(first_output.data(), generator->GetSequenceData(0), seq_len * sizeof(int32_t)));
 }
 
+// Test RewindTo with the qwen-2.5 model. Exercises the static mask rewind path if
+// the EP supports it (DML by default, WebGPU/CUDA when graph capture is enabled
+// in model generation via _test_utils.py), otherwise falls back to the dynamic mask path.
+// Skipped when qwen-2.5 model is not available.
+//
+// CUDA is explicitly disabled: RewindTo(seq_len - 3) — a deep partial rewind near
+// the end of a completed sequence — produces incorrect output on CUDA. This is a
+// pre-existing runtime bug (not model-specific): it reproduces with both
+// qwen-2.5-0.5b-graph and tiny-qwen35-cuda models, and with both static-mask
+// (graph-capture) and dynamic-mask (baseline) code paths. Full rewind (RewindTo(0))
+// and shallow partial rewind (e.g. RewindTo(input_ids.size()-1)) work correctly.
+// TODO: Remove !USE_CUDA once the CUDA partial rewind bug is fixed.
+#if TEST_QWEN_2_5 && !USE_CUDA
+TEST(CAPITests, RewindQwen25CAPI) {
+  // Prefer graph-capture variant (exercises static mask rewind on CUDA/WebGPU/DML),
+  // fall back to baseline model when it is not available.
+  std::string model_path = QWEN_2_5_GRAPH_PATH;
+  if (!std::filesystem::exists(model_path)) {
+    model_path = QWEN_2_5_PATH;
+  }
+  if (!std::filesystem::exists(model_path)) {
+    GTEST_SKIP() << "qwen-2.5 model not available at " << model_path;
+  }
+
+  int max_length = 50;
+  std::vector<int32_t> input_ids{1, 2, 3, 4, 5};
+
+  auto model = OgaModel::Create(model_path.c_str());
+  auto params = OgaGeneratorParams::Create(*model);
+  params->SetSearchOption("max_length", max_length);
+  params->SetSearchOptionBool("do_sample", false);
+
+  auto generator = OgaGenerator::Create(*model, *params);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  // Save first-run output
+  auto seq_len = generator->GetSequenceCount(0);
+  std::vector<int32_t> first_output(seq_len);
+  std::memcpy(first_output.data(), generator->GetSequenceData(0), seq_len * sizeof(int32_t));
+
+  // RewindTo(0) - full rewind
+  generator->RewindTo(0);
+  generator->AppendTokens(input_ids.data(), input_ids.size());
+  while (!generator->IsDone()) {
+    generator->GenerateNextToken();
+  }
+
+  auto seq_len2 = generator->GetSequenceCount(0);
+  ASSERT_EQ(seq_len2, seq_len);
+  EXPECT_TRUE(0 == std::memcmp(first_output.data(), generator->GetSequenceData(0), seq_len * sizeof(int32_t)));
+
+  // Partial rewind
+  if (seq_len > 7) {
+    generator->RewindTo(seq_len - 3);
+    while (!generator->IsDone()) {
+      generator->GenerateNextToken();
+    }
+
+    seq_len2 = generator->GetSequenceCount(0);
+    ASSERT_EQ(seq_len2, seq_len);
+    EXPECT_TRUE(0 == std::memcmp(first_output.data(), generator->GetSequenceData(0), seq_len * sizeof(int32_t)));
+  }
+}
+#endif  // TEST_QWEN_2_5
+
 #ifndef STREAMING_ASR_PATH
 #define STREAMING_ASR_PATH MODEL_PATH "nemotron-speech-streaming"
 #endif

test/python/_test_utils.py

@@ -60,14 +60,15 @@ def get_ci_data_path():
 def get_model_paths():
     # TODO: Uncomment the following models as needed in the CI pipeline.
 
-    # Format: model alias: (HF repo name, create only 1 layer)
+    # Format: model alias: (HF repo name, create only 1 layer, enable graph capture)
     hf_paths = {
         # "olmo": "amd/AMD-OLMo-1B-SFT-DPO",
         # "phi-3.5": "microsoft/Phi-3.5-mini-instruct",
         # "llama-3.2": "meta-llama/Llama-3.2-1B-instruct",
         # "granite-3.0": "ibm-granite/granite-3.0-2b-instruct",
-        "phi-4-mini": ("microsoft/Phi-4-mini-instruct", True),
-        "qwen-2.5-0.5b": ("Qwen/Qwen2.5-0.5B-Instruct", False),
+        "phi-4-mini": ("microsoft/Phi-4-mini-instruct", True, False),
+        "qwen-2.5-0.5b": ("Qwen/Qwen2.5-0.5B-Instruct", False, False),
+        "qwen-2.5-0.5b-graph": ("Qwen/Qwen2.5-0.5B-Instruct", False, True),
     }
 
     ci_data_path = os.path.join(get_ci_data_path(), "pytorch")
@@ -76,12 +77,12 @@ def get_model_paths():
 
     # Note: If a model has over 4B parameters, please add a quantized version
     # to `ci_paths` instead of `hf_paths` to reduce file size and testing time.
-    # Format: model alias: (OS path, create only 1 layer)
+    # Format: model alias: (OS path, create only 1 layer, enable graph capture)
     ci_paths = {
         # "llama-2": os.path.join(ci_data_path, "Llama-2-7B-Chat-GPTQ"),
         # "llama-3": os.path.join(ci_data_path, "Meta-Llama-3-8B-AWQ"),
         # "mistral-v0.2": os.path.join(ci_data_path, "Mistral-7B-Instruct-v0.2-GPTQ"),
-        "phi-2": (os.path.join(ci_data_path, "phi2"), True),
+        "phi-2": (os.path.join(ci_data_path, "phi2"), True, False),
         # "gemma-2b": os.path.join(ci_data_path, "gemma-1.1-2b-it"),
         # "gemma-7b": os.path.join(ci_data_path, "gemma-7b-it-awq"),
         # "phi-3-mini": os.path.join(ci_data_path, "phi3-mini-128k-instruct"),
@@ -94,7 +95,7 @@ def get_model_paths():
     return ci_paths, hf_paths
 
 
-def download_model(model_name, input_path, output_path, precision, device, one_layer):
+def download_model(model_name, input_path, output_path, precision, device, one_layer, enable_graph_capture):
     command = [
         sys.executable,
         "-m",
@@ -126,12 +127,30 @@ def download_model(model_name, input_path, output_path, precision, device, one_l
         extra_options += ["int4_accuracy_level=4"]
     if one_layer:
         extra_options += ["num_hidden_layers=1"]
+    # Graph capture is a generic model option and maps to EP-specific builder flags.
+    if enable_graph_capture and device == "webgpu":
+        extra_options += ["enable_webgpu_graph=true"]
+    if enable_graph_capture and device == "cuda":
+        extra_options += ["enable_cuda_graph=true"]
     if len(extra_options) > 1:
         command += extra_options
 
     run_subprocess(command).check_returncode()
 
 
+# Devices that support graph capture. Models with enable_graph_capture=True
+# are only built for these devices.
+#
+# CUDA is intentionally excluded: the Windows CUDA CI consistently fails to
+# download this new model from Hugging Face. Will re-add "cuda" once the
+# CI download issue is resolved.
+#
+# Note: nvtensorrtrtx is included here for model generation but does not have
+# dedicated CI coverage yet — tests using NvTensorRtRtx models are guarded with
+# GTEST_SKIP when the model artifacts are not present.
+_GRAPH_CAPTURE_DEVICES = {"webgpu", "dml", "nvtensorrtrtx"}
+
+
 def download_models(download_path, precision, device, log):
     log.debug(f"Downloading models to {download_path} with precision {precision} and device {device}")
 
@@ -141,19 +160,24 @@ def download_models(download_path, precision, device, log):
     log.debug(f"Downloading {len(ci_paths)} PyTorch models and {len(hf_paths)} Hugging Face models")
 
     # python -m onnxruntime_genai.models.builder -i <input_path> -o <output_path> -p <precision> -e <device>
-    for model_name, (input_path, one_layer) in ci_paths.items():
+    for model_name, (input_path, one_layer, graph_capture) in ci_paths.items():
+        if graph_capture and device.lower() not in _GRAPH_CAPTURE_DEVICES:
+            continue
         try:
             output_path = os.path.join(download_path, model_name, precision, device)
             log.debug(f"Downloading {model_name} from {input_path} to {output_path}")
             if not os.path.exists(output_path):
-                download_model(None, input_path, output_path, precision, device, one_layer)
+                download_model(None, input_path, output_path, precision, device, one_layer,
+                               graph_capture)
                 output_paths.append(output_path)
         except Exception as e:
             log.warning(f"Error: {e}. Skipping CI model.")
             continue
 
     # python -m onnxruntime_genai.models.builder -m <model_name> -o <output_path> -p <precision> -e <device>
-    for model_name, (hf_name, one_layer) in hf_paths.items():
+    for model_name, (hf_name, one_layer, graph_capture) in hf_paths.items():
+        if graph_capture and device.lower() not in _GRAPH_CAPTURE_DEVICES:
+            continue
         try:
             from huggingface_hub import model_info
 
@@ -169,7 +193,8 @@ def download_models(download_path, precision, device, log):
         log.debug(f"Downloading {model_name} from {hf_name} to {output_path}")
 
         if not os.path.exists(output_path):
-            download_model(hf_name, "", output_path, precision, device, one_layer)
+            download_model(hf_name, "", output_path, precision, device, one_layer,
+                           graph_capture)
             output_paths.append(output_path)
 
     log.info(f"Successfully downloaded {len(output_paths)} models")

test/python/test_onnxruntime_genai_e2e.py

@@ -54,7 +54,8 @@ def run_whisper():
     for precision, execution_provider in [("fp16", "cuda"), ("fp32", "cuda"), ("fp32", "cpu")]:
         # Generate model via model builder
         built_model = os.path.join(cwd, "..", "test_models", f"whisper-tiny-{precision}-{execution_provider}")
-        download_model(model_name="openai/whisper-tiny", input_path="", output_path=built_model, precision=precision, device=execution_provider, one_layer=False)
+        download_model(model_name="openai/whisper-tiny", input_path="", output_path=built_model, precision=precision,
+                       device=execution_provider, one_layer=False, enable_graph_capture=False)
 
         # Get prebuilt model from CI
         ci_model = os.path.join(ci_data_path, "onnx", f"whisper-tiny-{precision}-{execution_provider}")

test/test_utils.h

@@ -5,6 +5,7 @@
 
 #include <filesystem>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 // Our working directory is generators/build so one up puts us in the root directory:
@@ -14,11 +15,15 @@
 
 namespace test_utils {
 
-// Helper function to get the appropriate model path based on available models
+// Tests run across different environments where only one EP's model artifacts
+// may be present. This helper probes cuda/dml/webgpu/cpu directories in order
+// and returns the first path containing genai_config.json.
+// Caches results per model_type so different models resolve independently.
 inline const std::string& GetModelPath(const std::string& model_type) {
-  static std::string model_path;
-  if (!model_path.empty()) {
-    return model_path;
+  static std::unordered_map<std::string, std::string> model_paths;
+  auto it = model_paths.find(model_type);
+  if (it != model_paths.end()) {
+    return it->second;
   }
 
   std::vector<std::string> candidate_paths = {
@@ -30,14 +35,12 @@ inline const std::string& GetModelPath(const std::string& model_type) {
   for (const auto& path : candidate_paths) {
     std::filesystem::path model_path_fs(path);
     if (std::filesystem::exists(model_path_fs / "genai_config.json")) {
-      model_path = path;
-      return model_path;
+      return model_paths.emplace(model_type, path).first->second;
     }
   }
 
   // Fallback to CPU path
-  model_path = std::string(MODEL_PATH) + model_type + "/int4/cpu";
-  return model_path;
+  return model_paths.emplace(model_type, std::string(MODEL_PATH) + model_type + "/int4/cpu").first->second;
 }
 
 // Helper to detect if we're using WebGPU or DML EP based on the model path
@@ -61,3 +64,7 @@ inline bool IsEngineTestsEnabled() {
 #ifndef QWEN_2_5_PATH
 #define QWEN_2_5_PATH test_utils::GetModelPath("qwen-2.5-0.5b").c_str()
 #endif
+
+#ifndef QWEN_2_5_GRAPH_PATH
+#define QWEN_2_5_GRAPH_PATH test_utils::GetModelPath("qwen-2.5-0.5b-graph").c_str()
+#endif