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
 

test/c_api_tests.cpp

@@ -1384,6 +1384,62 @@ 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 with enableGraphCapture=1 set by
+// _test_utils.py), otherwise falls back to the dynamic mask path.
+// Skipped when qwen-2.5 model is not available.
+// TODO: This test fails on CUDA EP. Need to investigate the root cause.
+#if TEST_QWEN_2_5 && !USE_CUDA
+TEST(CAPITests, RewindQwen25CAPI) {
+  std::string 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 && !USE_CUDA
+
 #ifndef STREAMING_ASR_PATH
 #define STREAMING_ASR_PATH MODEL_PATH "nemotron-speech-streaming"
 #endif

test/guidance_tests.cpp

@@ -69,7 +69,7 @@ TEST_F(GuidanceTests, UseRegex) {
   EXPECT_TRUE(std::regex_match(output, std::regex("answer: .*")));
 }
 
-#if 0  // Temporarily disable JSON schema and LARK grammar tests
+#if 0   // Temporarily disable JSON schema and LARK grammar tests
 TEST_F(GuidanceTests, UseLarkGrammarSingleTurn) {
   auto input_string = get_qwen_2_5_prompt("What is the weather in Seattle?");
   auto input_sequences = OgaSequences::Create();
@@ -111,7 +111,6 @@ TEST_F(GuidanceTests, UseJsonSchemaSingleTurn) {
   EXPECT_EQ(output, expected_output);
 }
 
-#if !USE_DML  // DML doesn't support continuous decoding
 TEST_F(GuidanceTests, UseLarkGrammarMultiTurn) {
   auto params = OgaGeneratorParams::Create(*model_);
   params->SetSearchOption("max_length", 1024);
@@ -170,8 +169,7 @@ TEST_F(GuidanceTests, UseJsonSchemaMultiTurn) {
     EXPECT_EQ(output, expected_output);
   }
 }
-#endif        // USE_DML
-#endif        // #if 0
+#endif  // #if 0
 
 #endif  // TEST_QWEN_2_5
 

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)
+    # Note: enable_graph_capture currently only affects WebGPU EP (adds enable_webgpu_graph=true).
     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, True),
     }
 
     ci_data_path = os.path.join(get_ci_data_path(), "pytorch")
@@ -76,12 +77,13 @@ 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)
+    # Note: enable_graph_capture currently only affects WebGPU EP.
     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 +96,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=False):
     command = [
         sys.executable,
         "-m",
@@ -126,6 +128,8 @@ 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"]
+    if enable_graph_capture and device == "webgpu":
+        extra_options += ["enable_webgpu_graph=true"]
     if len(extra_options) > 1:
         command += extra_options
 
@@ -141,19 +145,20 @@ 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():
         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,
+                               enable_graph_capture=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():
         try:
             from huggingface_hub import model_info
 
@@ -169,7 +174,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,
+                           enable_graph_capture=graph_capture)
             output_paths.append(output_path)
 
     log.info(f"Successfully downloaded {len(output_paths)} models")

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,13 @@
 
 namespace test_utils {
 
-// Helper function to get the appropriate model path based on available models
+// Helper function to get the appropriate model path based on available models.
+// 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 +33,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