Enable graph capture for WebGPU models and DML continuous decoding tests

test/c_api_tests.cpp

@@ -1384,6 +1384,61 @@ 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.
+#if TEST_QWEN_2_5
+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
+
 #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,14 @@ 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, True),
     }
 
     ci_data_path = os.path.join(get_ci_data_path(), "pytorch")
@@ -76,12 +76,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 +94,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 +126,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 +143,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 +172,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")