diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md
index 9cced1adb7..f55f9352d3 100644
--- a/.github/copilot-instructions.md
+++ b/.github/copilot-instructions.md
@@ -58,7 +58,7 @@ WinML integration downloads `Microsoft.WindowsAppSDK.ML` via NuGet and copies he
```bash
# Python tests with test models
-python -m pytest -sv test_onnxruntime_genai_api.py -k "test_name" --test_models ..\test_models
+python -m pytest -sv test_onnxruntime_genai_api.py -k "test_name" --test_models ..\models
# C++ unit tests via CMake/CTest
ctest --build-config Release --output-on-failure
@@ -160,4 +160,4 @@ Tests are organized by language binding:
- **Python tests**: `test/python/`, includes end-to-end model testing
- **Platform tests**: Android/iOS tests run via emulator/simulator
-Always test with actual model files from `test/test_models/` directory rather than mock data.
+Always test with actual model files from `test/models/` directory rather than mock data.
diff --git a/.github/instructions/python-model-builder.instructions.md b/.github/instructions/python-model-builder.instructions.md
index 74477c65bf..8f9ade9e95 100644
--- a/.github/instructions/python-model-builder.instructions.md
+++ b/.github/instructions/python-model-builder.instructions.md
@@ -11,4 +11,9 @@ When generating or reviewing code in the Python Model Builder (`src/python/py/mo
Read both documents to understand the intended usage, supported models, design principles, and architectural constraints before suggesting or reviewing any code changes in this area.
-When a node is inserted into the model, prefer using `self.make_op_name` as the wrapper method for `self.make_node` + `self.make_value` calls.
+## Code Style Guidelines
+
+1. When a node is inserted into the model, prefer using `self.make_op_name` as the wrapper method for `self.make_node` + `self.make_value` calls.
+2. Ignore any CodeQL warnings about how an __init__ method calls an overridden method. These warnings are false positives and can be safely ignored. The warning message is: "this call to ABC in an initialization method is overwritten by XYZ".
+3. Find ways to reduce code duplication by reusing existing functionality and implementing common patterns.
+4. Discover ways to leverage the use of shared code in the base classes to avoid code duplication and improve maintainability.
diff --git a/.github/workflows/linux-cpu-x64-build.yml b/.github/workflows/linux-cpu-x64-build.yml
index 2094821a24..96a42a535d 100644
--- a/.github/workflows/linux-cpu-x64-build.yml
+++ b/.github/workflows/linux-cpu-x64-build.yml
@@ -131,12 +131,12 @@ jobs:
set -e -x
ctest --test-dir build/cpu/src/java --build-config Release --verbose --timeout 10800
- # This will also download all the test models to the test/test_models directory
+ # This will also download all the test models to the test/models directory
# These models are used by the python tests as well as C#, C++ and others.
- name: Run the Python tests
run: |
export ORTGENAI_LOG_ORT_LIB=1
- python3 test/python/test_onnxruntime_genai.py --cwd test/python --test_models test/test_models
+ python3 test/python/test_onnxruntime_genai.py --cwd test/python --test_models test/models
- name: Build the C# API and Run the C# Tests
run: |
@@ -155,8 +155,8 @@ jobs:
- name: Test the C# LLM Example with Tool Calling
run: |
export ORTGENAI_LOG_ORT_LIB=1
- python3 test/python/special_tokens.py -p test/test_models/qwen-2.5-0.5b/int4/cpu/tokenizer.json -s "" -e ""
- ./examples/csharp/ModelChat/bin/Release/net8.0/ModelChat -m test/test_models/qwen-2.5-0.5b/int4/cpu/ -e cpu --response_format lark_grammar --tools_file test/test_models/tool-definitions/weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
+ python3 test/python/special_tokens.py -p test/models/qwen-2.5-0.5b/int4/cpu/tokenizer.json -s "" -e ""
+ ./examples/csharp/ModelChat/bin/Release/net8.0/ModelChat -m test/models/qwen-2.5-0.5b/int4/cpu/ -e cpu --response_format lark_grammar --tools_file test/models/tool-definitions/weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
- name: Run tests
run: |
diff --git a/.github/workflows/linux-cpu-x64-nightly-build.yml b/.github/workflows/linux-cpu-x64-nightly-build.yml
index a7c506699c..e76dea328f 100644
--- a/.github/workflows/linux-cpu-x64-nightly-build.yml
+++ b/.github/workflows/linux-cpu-x64-nightly-build.yml
@@ -136,7 +136,7 @@ jobs:
- name: Run the Python tests
run: |
export ORTGENAI_LOG_ORT_LIB=1
- python3 test/python/test_onnxruntime_genai.py --cwd test/python --test_models test/test_models --e2e
+ python3 test/python/test_onnxruntime_genai.py --cwd test/python --test_models test/models --e2e
- name: Build the C# API and Run the C# Tests
run: |
@@ -155,8 +155,8 @@ jobs:
- name: Test the C# LLM Example with Tool Calling
run: |
export ORTGENAI_LOG_ORT_LIB=1
- python3 test/python/special_tokens.py -p test/test_models/qwen-2.5-0.5b/int4/cpu/tokenizer.json -s "" -e ""
- ./examples/csharp/ModelChat/bin/Release/net8.0/ModelChat -m test/test_models/qwen-2.5-0.5b/int4/cpu/ -e cpu --response_format lark_grammar --tools_file test/test_models/tool-definitions/weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
+ python3 test/python/special_tokens.py -p test/models/qwen-2.5-0.5b/int4/cpu/tokenizer.json -s "" -e ""
+ ./examples/csharp/ModelChat/bin/Release/net8.0/ModelChat -m test/models/qwen-2.5-0.5b/int4/cpu/ -e cpu --response_format lark_grammar --tools_file test/models/tool-definitions/weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
- name: Run Q&A Example
run: |
diff --git a/.github/workflows/linux-gpu-x64-build.yml b/.github/workflows/linux-gpu-x64-build.yml
index 4afece681d..4a0dec58f0 100644
--- a/.github/workflows/linux-gpu-x64-build.yml
+++ b/.github/workflows/linux-gpu-x64-build.yml
@@ -180,7 +180,7 @@ jobs:
${{ env.PYTHON_EXECUTABLE }} -m pip install -r test/python/cuda/torch/requirements.txt --user && \
${{ env.PYTHON_EXECUTABLE }} -m pip install -r test/python/cuda/ort/requirements.txt --user && \
${{ env.PYTHON_EXECUTABLE }} -m pip install /ort_genai_src/build/cuda/wheel/onnxruntime_genai*manylinux*.whl --no-deps --user && \
- ${{ env.PYTHON_EXECUTABLE }} test/python/test_onnxruntime_genai.py --cwd test/python --test_models test/test_models --e2e"
+ ${{ env.PYTHON_EXECUTABLE }} test/python/test_onnxruntime_genai.py --cwd test/python --test_models test/models --e2e"
# TODO: Enable this by adding dotnet to the docker image
# - name: Build the C# API and Run the C# Tests
diff --git a/.github/workflows/win-cpu-arm64-build.yml b/.github/workflows/win-cpu-arm64-build.yml
index 56d28dc1e9..c4b4396623 100644
--- a/.github/workflows/win-cpu-arm64-build.yml
+++ b/.github/workflows/win-cpu-arm64-build.yml
@@ -114,7 +114,7 @@ jobs:
- name: Run the Python Tests
run: |
- python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\test_models"
+ python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\models"
- name: Build the C# API and Run the C# Tests
run: |
@@ -130,8 +130,8 @@ jobs:
- name: Test the C# LLM Example with Tool Calling
run: |
- python3 test\python\special_tokens.py -p test\test_models\qwen-2.5-0.5b\int4\cpu\tokenizer.json -s "" -e ""
- .\examples\csharp\ModelChat\bin\Release\net8.0\ModelChat.exe -m test\test_models\qwen-2.5-0.5b\int4\cpu\ -e cpu --response_format lark_grammar --tools_file test\test_models\tool-definitions\weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
+ python3 test\python\special_tokens.py -p test\models\qwen-2.5-0.5b\int4\cpu\tokenizer.json -s "" -e ""
+ .\examples\csharp\ModelChat\bin\Release\net8.0\ModelChat.exe -m test\models\qwen-2.5-0.5b\int4\cpu\ -e cpu --response_format lark_grammar --tools_file test\models\tool-definitions\weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
- name: Verify Build Artifacts
if: always()
diff --git a/.github/workflows/win-cpu-x64-build.yml b/.github/workflows/win-cpu-x64-build.yml
index 453df0422a..6c3d339d5c 100644
--- a/.github/workflows/win-cpu-x64-build.yml
+++ b/.github/workflows/win-cpu-x64-build.yml
@@ -125,7 +125,7 @@ jobs:
- name: Run the Python Tests
run: |
- python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\test_models"
+ python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\models"
- name: Build the C# API and Run the C# Tests
run: |
@@ -141,8 +141,8 @@ jobs:
- name: Test the C# LLM Example with Tool Calling
run: |
- python3 test\python\special_tokens.py -p test\test_models\qwen-2.5-0.5b\int4\cpu\tokenizer.json -s "" -e ""
- .\examples\csharp\ModelChat\bin\Release\net8.0\ModelChat.exe -m test\test_models\qwen-2.5-0.5b\int4\cpu\ -e cpu --response_format lark_grammar --tools_file test\test_models\tool-definitions\weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
+ python3 test\python\special_tokens.py -p test\models\qwen-2.5-0.5b\int4\cpu\tokenizer.json -s "" -e ""
+ .\examples\csharp\ModelChat\bin\Release\net8.0\ModelChat.exe -m test\models\qwen-2.5-0.5b\int4\cpu\ -e cpu --response_format lark_grammar --tools_file test\models\tool-definitions\weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
- name: Verify Build Artifacts
if: always()
diff --git a/.github/workflows/win-cuda-x64-build.yml b/.github/workflows/win-cuda-x64-build.yml
index d799950d1c..43b6e6c90e 100644
--- a/.github/workflows/win-cuda-x64-build.yml
+++ b/.github/workflows/win-cuda-x64-build.yml
@@ -158,7 +158,7 @@ jobs:
- name: Run the Python Tests
run: |
- python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\test_models" --e2e
+ python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\models" --e2e
- name: Verify Build Artifacts
if: always()
@@ -182,8 +182,8 @@ jobs:
- name: Test the C# LLM Example with Tool Calling
run: |
- python test\python\special_tokens.py -p test\test_models\qwen-2.5-0.5b\int4\cpu\tokenizer.json -s "" -e ""
- .\examples\csharp\ModelChat\bin\Release\net8.0\ModelChat.exe -m test\test_models\qwen-2.5-0.5b\int4\cpu\ -e cpu --response_format lark_grammar --tools_file test\test_models\tool-definitions\weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
+ python test\python\special_tokens.py -p test\models\qwen-2.5-0.5b\int4\cpu\tokenizer.json -s "" -e ""
+ .\examples\csharp\ModelChat\bin\Release\net8.0\ModelChat.exe -m test\models\qwen-2.5-0.5b\int4\cpu\ -e cpu --response_format lark_grammar --tools_file test\models\tool-definitions\weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
- name: Prepend CUDA to PATH and Run tests
run: |-
diff --git a/.github/workflows/win-directml-x64-build.yml b/.github/workflows/win-directml-x64-build.yml
index 0263b9ccae..b77584ee96 100644
--- a/.github/workflows/win-directml-x64-build.yml
+++ b/.github/workflows/win-directml-x64-build.yml
@@ -129,7 +129,7 @@ jobs:
- name: Run the Python Tests
run: |
- python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\test_models" --e2e
+ python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\models" --e2e
- name: Verify Build Artifacts
if: always()
diff --git a/.github/workflows/win-webgpu-x64-build.yml b/.github/workflows/win-webgpu-x64-build.yml
index a72e4dfd7d..5ca09c6bd7 100644
--- a/.github/workflows/win-webgpu-x64-build.yml
+++ b/.github/workflows/win-webgpu-x64-build.yml
@@ -103,7 +103,7 @@ jobs:
- name: Run the Python Tests
run: |
- python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\test_models" --e2e
+ python test/python/test_onnxruntime_genai.py --cwd "test\python" --test_models "test\models" --e2e
- name: Verify Build Artifacts
if: always()
diff --git a/.gitignore b/.gitignore
index 07fc122d62..d8a1b8691a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,7 +10,7 @@ src/csharp/AssemblyInfo.cs
/ort
/build
/build_*
-/test/test_models/*
+/test/models/*
/cache_models
/onnxruntime-linux-x64-*
*.csv
@@ -29,24 +29,17 @@ examples/python/genai_models
examples/python/hf_cache
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
-!test/test_models/qwen3-vl-vision-preprocessing/
-!test/test_models/qwen3-vl-vision-preprocessing/*.onnx
-!test/test_models/qwen35-hybrid-preprocessing/
-!test/test_models/qwen35-hybrid-preprocessing/*.onnx
-!test/test_models/mistral3-vision-preprocessing/
-!test/test_models/gemma4-vision-preprocessing/
-!test/test_models/gemma4-vision-preprocessing/*.onnx
-!test/test_models/multimodal-decoder-no-input-ids/
-!test/test_models/multimodal-decoder-no-input-ids/*
-!test/test_models/multimodal-decoder-with-input-ids/
-!test/test_models/multimodal-decoder-with-input-ids/*
-!test/test_models/create_decoder_input_ids_test_models.py
+!test/models/gemma4/*
+!test/models/hf-internal-testing/*
+!test/models/mistral3/*
+!test/models/multimodal-decoder-no-input-ids/*
+!test/models/multimodal-decoder-with-input-ids/*
+!test/models/phi3-v/*
+!test/models/pipeline-model/*.json
+!test/models/qwen2-5-vl/*
+!test/models/qwen3-5/*
+!test/models/qwen3-vl/*
+!test/models/whisper/*
.ipynb_checkpoints/
/src/java/.gradle
diff --git a/.pipelines/mac-cpu-arm64-build.yml b/.pipelines/mac-cpu-arm64-build.yml
index 13777f68d1..4c6c59e544 100644
--- a/.pipelines/mac-cpu-arm64-build.yml
+++ b/.pipelines/mac-cpu-arm64-build.yml
@@ -208,7 +208,7 @@ jobs:
source genai-macos-venv/bin/activate
export ORTGENAI_LOG_ORT_LIB=1
python3 -m pip install requests
- python3 test/python/test_onnxruntime_genai.py --cwd test/python --test_models test/test_models
+ python3 test/python/test_onnxruntime_genai.py --cwd test/python --test_models test/models
displayName: 'Run the Python tests'
workingDirectory: '$(Build.SourcesDirectory)'
@@ -233,8 +233,8 @@ jobs:
- bash: |
set -e -x
export ORTGENAI_LOG_ORT_LIB=1
- python3 test/python/special_tokens.py -p test/test_models/qwen-2.5-0.5b/int4/cpu/tokenizer.json -s "" -e ""
- ./examples/csharp/ModelChat/bin/Release/net8.0/ModelChat -m test/test_models/qwen-2.5-0.5b/int4/cpu/ -e cpu --response_format lark_grammar --tools_file test/test_models/tool-definitions/weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
+ python3 test/python/special_tokens.py -p test/models/qwen-2.5-0.5b/int4/cpu/tokenizer.json -s "" -e ""
+ ./examples/csharp/ModelChat/bin/Release/net8.0/ModelChat -m test/models/qwen-2.5-0.5b/int4/cpu/ -e cpu --response_format lark_grammar --tools_file test/models/tool-definitions/weather.json --tool_call_start "" --tool_call_end "" --user_prompt "What is the weather in Redmond, WA?" --tool_output --non_interactive --verbose
displayName: 'Test the C# LLM example with tool calling'
workingDirectory: '$(Build.SourcesDirectory)'
diff --git a/examples/c/src/common.cpp b/examples/c/src/common.cpp
index 975d64bb8d..451f9e2896 100644
--- a/examples/c/src/common.cpp
+++ b/examples/c/src/common.cpp
@@ -195,7 +195,7 @@ bool ParseArgs(
app.add_option("-p,--top_p", generator_params_args.top_p, "Top p probability to sample with")->group(generator_params);
app.add_option("--response_format", guidance_args.response_format, "Provide response format for the model")->group(guidance);
- app.add_option("--tools_file", guidance_args.tools_file, "Path to file containing list of OpenAI-compatible tool definitions. Ex: test/test_models/tool-definitions/weather.json")->group(guidance);
+ app.add_option("--tools_file", guidance_args.tools_file, "Path to file containing list of OpenAI-compatible tool definitions. Ex: test/models/tool-definitions/weather.json")->group(guidance);
app.add_flag("--text_output", guidance_args.text_output, "Produce a text response in the output")->group(guidance);
app.add_flag("--tool_output", guidance_args.tool_output, "Produce a tool call in the output")->group(guidance);
app.add_option("--tool_call_start", guidance_args.tool_call_start, "String representation of tool call start (ex: <|tool_call|>). Needs to be marked as special in tokenizer.json for guidance to work.")->group(guidance);
diff --git a/examples/csharp/Common/Common.cs b/examples/csharp/Common/Common.cs
index 6ebac6baaf..56b75f829a 100644
--- a/examples/csharp/Common/Common.cs
+++ b/examples/csharp/Common/Common.cs
@@ -878,7 +878,7 @@ public static void GetGuidanceArgs(RootCommand parser)
{
Arity = ArgumentArity.ExactlyOne,
DefaultValueFactory = (_) => "",
- Description = "Path to file containing list of OpenAI-compatible tool definitions. Ex: test/test_models/tool-definitions/weather.json"
+ Description = "Path to file containing list of OpenAI-compatible tool definitions. Ex: test/models/tool-definitions/weather.json"
};
tools_file.Validators.Add(result =>
{
diff --git a/examples/python/common.py b/examples/python/common.py
index 1ebd1fee82..80b8d05398 100644
--- a/examples/python/common.py
+++ b/examples/python/common.py
@@ -660,7 +660,7 @@ def get_guidance_args(parser: argparse.ArgumentParser) -> None:
"--tools_file",
type=str,
default="",
- help="Path to file containing list of OpenAI-compatible tool definitions. Ex: test/test_models/tool-definitions/weather.json",
+ help="Path to file containing list of OpenAI-compatible tool definitions. Ex: test/models/tool-definitions/weather.json",
)
guidance.add_argument(
"-text", "--text_output", action="store_true", default=False, help="Produce a text response in the output"
diff --git a/src/java/CMakeLists.txt b/src/java/CMakeLists.txt
index 4724b65def..26a9798d8c 100644
--- a/src/java/CMakeLists.txt
+++ b/src/java/CMakeLists.txt
@@ -196,7 +196,7 @@ if (ANDROID)
# Copy the test model to the assets folder in the test app
add_custom_command(TARGET onnxruntime-genai-jni POST_BUILD
COMMAND ${CMAKE_COMMAND} -E copy_directory_if_different
- ${REPO_ROOT}/test/test_models/hf-internal-testing/tiny-random-gpt2-fp32
+ ${REPO_ROOT}/test/models/hf-internal-testing/tiny-random-gpt2-fp32
${ANDROID_TEST_PACKAGE_APP_ASSETS_DIR}/model)
# Copy the Android AAR package we built to the libs folder of our test app
diff --git a/src/java/src/test/android/README.md b/src/java/src/test/android/README.md
index e28aff853a..a856a379d0 100644
--- a/src/java/src/test/android/README.md
+++ b/src/java/src/test/android/README.md
@@ -6,7 +6,7 @@ This directory contains a simple android application for testing the ONNX Runtim
This android application is mainly aimed for testing:
-- Model used: test/test_models/hf-internal-testing/tiny-random-gpt2-fp32
+- Model used: test/models/hf-internal-testing/tiny-random-gpt2-fp32
- Main test file: An android instrumentation test under `app\src\androidtest\java\ai.onnxruntime.genai.example.javavalidator\SimpleTest.kt`
- The main dependency of this application is `onnxruntime-genai` aar package under `app\libs`.
- The onnxruntime dependency is provided by the latest released onnxruntime-android package.
diff --git a/src/java/src/test/java/ai/onnxruntime/genai/GenerationTest.java b/src/java/src/test/java/ai/onnxruntime/genai/GenerationTest.java
index bc25395a1f..11e996b504 100644
--- a/src/java/src/test/java/ai/onnxruntime/genai/GenerationTest.java
+++ b/src/java/src/test/java/ai/onnxruntime/genai/GenerationTest.java
@@ -24,7 +24,7 @@ public class GenerationTest {
// phi-2 can be used in full end-to-end testing but needs to be manually downloaded.
// it's also used this way in the C# unit tests.
private static final String phi2ModelPath() {
- return TestUtils.getTestResourcePath("phi-2/int4/cpu");
+ return TestUtils.getTestModelPath("phi-2/int4/cpu");
}
@SuppressWarnings("unused") // Used in EnabledIf
diff --git a/src/java/src/test/java/ai/onnxruntime/genai/MultiModalProcessorTest.java b/src/java/src/test/java/ai/onnxruntime/genai/MultiModalProcessorTest.java
index 7a29677b4e..6bba1fd7ca 100644
--- a/src/java/src/test/java/ai/onnxruntime/genai/MultiModalProcessorTest.java
+++ b/src/java/src/test/java/ai/onnxruntime/genai/MultiModalProcessorTest.java
@@ -23,9 +23,7 @@ public void testBatchEncodeDecode() throws GenAIException {
new String(
"<|user|>\n<|image_1|>\n Can you convert the table to markdown format?\n<|end|>\n<|assistant|>\n");
try (Images image =
- new Images(
- TestUtils.getFilePathFromDisk(
- TestUtils.getTestResourcePath("images/sheet.png")));
+ new Images(TestUtils.getFilePathFromDisk(TestUtils.getTestImagePath("sheet.png")));
NamedTensors processed = multiModalProcessor.processImages(inputs, image); ) {
assertNotNull(processed);
}
diff --git a/src/java/src/test/java/ai/onnxruntime/genai/TestUtils.java b/src/java/src/test/java/ai/onnxruntime/genai/TestUtils.java
index a0c2972daa..d6394ba72c 100644
--- a/src/java/src/test/java/ai/onnxruntime/genai/TestUtils.java
+++ b/src/java/src/test/java/ai/onnxruntime/genai/TestUtils.java
@@ -12,27 +12,31 @@ public class TestUtils {
private static final Logger logger = Logger.getLogger(TestUtils.class.getName());
public static final String testAdapterTestModelPath() {
- return getFilePathFromDisk(getTestResourcePath("adapters"));
+ return getFilePathFromDisk(getTestModelPath("adapters"));
}
public static final String testAdapterTestAdaptersPath() {
- return getFilePathFromDisk(getTestResourcePath("adapters/adapters.onnx_adapter"));
+ return getFilePathFromDisk(getTestModelPath("adapters/adapters.onnx_adapter"));
}
public static final String tinyGpt2ModelPath() {
- return getFilePathFromDisk(getTestResourcePath("hf-internal-testing/tiny-random-gpt2-fp32"));
+ return getFilePathFromDisk(getTestModelPath("hf-internal-testing/tiny-random-gpt2-fp32"));
}
public static final String phi2ModelPath() {
- return getFilePathFromDisk(getTestResourcePath("phi-2/int4/cpu"));
+ return getFilePathFromDisk(getTestModelPath("phi-2/int4/cpu"));
}
public static final String testVisionModelPath() {
- return getFilePathFromDisk(getTestResourcePath("vision-preprocessing"));
+ return getFilePathFromDisk(getTestModelPath("phi3-v"));
}
- public static final String getTestResourcePath(String relativeResourcePath) {
- return getFilePathFromDisk(getRepoRoot() + "test/test_models/" + relativeResourcePath);
+ public static final String getTestModelPath(String relativeResourcePath) {
+ return getFilePathFromDisk(getRepoRoot() + "test/models/" + relativeResourcePath);
+ }
+
+ public static final String getTestImagePath(String relativeResourcePath) {
+ return getFilePathFromDisk(getRepoRoot() + "test/images/" + relativeResourcePath);
}
public static final String getRepoRoot() {
diff --git a/src/ort_genai_c.h b/src/ort_genai_c.h
index b507ca5b14..62f71c690d 100644
--- a/src/ort_genai_c.h
+++ b/src/ort_genai_c.h
@@ -433,8 +433,8 @@ OGA_EXPORT OgaResult* OGA_API_CALL OgaGeneratorParamsSetSearchBool(OgaGeneratorP
/**
* \brief Sets the guidance type and data for the Generator params
* \param[in] params The generator params to set the guidance on
- * \param[in] type The type of the guidance. Currently, we support json_schema, regex and lark_grammar
- * \param[in] data The input string, which is the guidance data. Examples are present in test/test_models/grammars folder
+ * \param[in] type The type of the guidance. Currently, we support json_schema, regex and lark_grammar.
+ * \param[in] data The input string, which is the guidance data.
* \param[in] enable_ff_tokens Whether to enable ff_tokens generation. This feature allows guidance to force-forward tokens that satisfy input grammar without calling model, hence speeding up generation process. Only valid when guidance type is set and batch_size is 1 and beam_size is 1.
* \return OgaResult containing the error message if the setting of the guidance failed
*/
diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 2aab41b83d..04c1d30dff 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -416,7 +416,7 @@ def get_args():
int4_block_size = 16/32/64/128/256: Specify the block size for int4 quantization (MatMulNBits).
Default value is 32.
qmoe_block_size = 16/32/64/128/256: Specify the block size for QMoE expert weights quantization.
- Default is 128 for trt-rtx, 32 for others. Supported EPs: cpu, webgpu, trt-rtx.
+ Default is 128 for CUDA and TRT-RTX, 32 for others. Supported EPs: CPU, CUDA, WebGPU, TRT-RTX.
int4_is_symmetric = Quantize the weights symmetrically. Default is true.
If true, quantization is done to int4. If false, quantization is done to uint4.
int4_op_types_to_quantize = MatMul/Gather: Specify op types to target for int4 quantization.
diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py
index 061862f153..de416c18d1 100644
--- a/src/python/py/models/builders/base.py
+++ b/src/python/py/models/builders/base.py
@@ -52,6 +52,7 @@ def parse_hf_token(hf_token):
class Model:
def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
+ # Model attributes from config
self.context_length = config.seq_length if hasattr(config, "seq_length") else config.max_position_embeddings
self.original_context_length = (
config.original_max_position_embeddings
@@ -131,6 +132,13 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
},
"trt-rtx": {"enable_cuda_graph": "1"},
}
+ self.graph_capture = (
+ extra_options.get("enable_cuda_graph", False) or
+ extra_options.get("enable_webgpu_graph", False) or
+ self.ep in {"dml", "trt-rtx"}
+ )
+ # Initialize EP-specific expansions
+ self.make_ep_expansions_init()
# Map input names to their types and shapes
self.input_names = {
@@ -226,10 +234,10 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
"use_lora": is_lora, # Use LoRA/QLoRA format
}
- # RotaryEmbedding-specific variables
+ # RoPE-specific variables
position_scale = config.rope_position_scale if hasattr(config, "rope_position_scale") else 1
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
- rotemb_dim = int(self.head_size * partial_rotary_factor) if partial_rotary_factor != 1.0 else 0
+ rope_dim = int(self.head_size * partial_rotary_factor) if partial_rotary_factor != 1.0 else 0
rope_theta = (
config.rope_theta
if hasattr(config, "rope_theta")
@@ -248,7 +256,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
"theta": rope_theta, # Base value if calculating cos/sin caches from scratch
"partial_rotary_factor": partial_rotary_factor, # Factor for partial rotary embeddings
"interleaved": 0, # Interleave the rotary embeddings (e.g. [0, 0, 0, 1, 1, 1] to [0, 1, 0, 1, 0, 1], RotaryEmbedding kernel expects a default value of 0)
- "rotary_embedding_dim": rotemb_dim, # For partial rotary embeddings (RotaryEmbedding kernel expects a default value of 0)
+ "rotary_embedding_dim": rope_dim, # For partial rotary embeddings (RotaryEmbedding kernel expects a default value of 0)
"rescale_factors": 1, # Rescale factors when calculating `inv_freq` in rotary embeddings
"t_dtype": torch.int64, # Torch dtype when calculating `t` in rotary embeddings
"position_scale": position_scale, # Scale value when calculating `t` in rotary embeddings
@@ -260,13 +268,6 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
# Attention-specific variables (MHA, GQA, GQA + Rot.Emb., etc.)
attn_softcap = config.attn_logit_softcapping if hasattr(config, "attn_logit_softcapping") and config.attn_logit_softcapping is not None else 0.0 # default is 0.0 in GroupQueryAttention kernel
-
- # Block-sparse attention-specific variables
- sparse_block_size = config.blocksparse_block_size if hasattr(config, "blocksparse_block_size") else 0
- kernel_block_size = config.blocksparse_triton_kernel_block_size if hasattr(config, "blocksparse_triton_kernel_block_size") else 0
- local_blocks = config.blocksparse_num_local_blocks if hasattr(config, "blocksparse_num_local_blocks") else 0
- vert_block_stride = config.blocksparse_vert_stride if hasattr(config, "blocksparse_vert_stride") else 0
- homo_head = config.blocksparse_homo_head_pattern if hasattr(config, "blocksparse_homo_head_pattern") else False
self.attention_attrs = {
# Attributes for MHA, GQA, etc:
"q_path": "", # Q path to attention
@@ -277,13 +278,6 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
"softcap": attn_softcap, # Softcap value to prevent values from exploding in attention
"use_rope_in_attn": False, # Use rotary embeddings within attention (instead of a separate RotaryEmbedding op)
"use_packed_matmul": False, # Use packed MatMul (instead of 3 separate MatMuls for Q/K/V)
- "block_sparse": { # Block-sparse attention-specific variables
- "sparse_block_size": sparse_block_size, # Sparse block size for SparseAttention op
- "kernel_block_size": kernel_block_size, # Kernel block size for sparse attention
- "local_blocks": local_blocks, # Number of local blocks for sparse attention
- "vert_stride": vert_block_stride, # Vertical stride to use for sparse attention
- "homo_head": homo_head, # Use homo head pattern for sparse attention
- },
"rope": True, # Use rotary embeddings in attention subgraph
"q_norm": False, # LayerNorm after MatMul in Q path
"k_norm": False, # LayerNorm after MatMul in K path
@@ -296,7 +290,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
"unidirectional": False, # Whether every token can only attend to previous tokens
"use_matmul_in_attn": False, # Use MatMuls with attention (instead of separate MatMul ops)
}
- self.make_attention_init()
+ self.make_attention_init(config)
# MLP-specific variables
self.mlp_attrs = {
@@ -331,73 +325,54 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
"mask": None, # LM head mask for tokens in the vocabulary
"softcap": lm_head_softcap, # Softcap value to prevent values from exploding in LM head
}
- if hasattr(config, "dummy_token_indices"):
- # Create LM head mask for tokens in the vocabulary
- dummy_tokens_mask = torch.zeros(self.vocab_size).bool()
- dummy_tokens_mask[config.dummy_token_indices] = True
- self.lm_head_attrs["mask"] = dummy_tokens_mask
+ self.make_lm_head_init(config)
# Quantization-specific variables (INT4, INT8, etc.)
- int4_algo_config = self.make_int4_algo_config(extra_options.get("int4_algo_config", "default"))
- self.int4_block_size = extra_options.get("int4_block_size", 32)
-
- # CPU, WebGPU, and TRT-RTX support block-wise quantization for QMoE.
- # TRT-RTX defaults to 128; others default to 32 for consistency with MatMulNBits.
- supported_blockwise_eps = ["cpu", "webgpu", "trt-rtx"]
- default_qmoe_block_size = 128 if self.ep == "trt-rtx" else 32
- self.qmoe_block_size = int(extra_options.get("qmoe_block_size", default_qmoe_block_size))
-
- # Validate that unsupported EPs don't explicitly request block-wise quantization
- if self.ep not in supported_blockwise_eps and "qmoe_block_size" in extra_options and moe_op_type == "QMoE":
- raise ValueError(
- f"The 'qmoe_block_size' option is not supported for {self.ep} execution provider with QMoE. "
- f"Block-wise quantization is only supported for: {', '.join(supported_blockwise_eps)}."
- )
-
+ algo_config = self.make_algo_config(extra_options.get("int4_algo_config", "default"))
+ self.matmul_block_size = int(extra_options.get("int4_block_size", 32))
+ self.qmoe_block_size = int(extra_options.get("qmoe_block_size", 128 if self.ep in {"cuda", "trt-rtx"} else 32))
self.quant_attrs = {
- "int4": {
- "accuracy_level": int(extra_options.get("int4_accuracy_level", 4 if self.ep in ["cpu", "webgpu"] else 0)),
- "qmoe_block_size": int(self.qmoe_block_size),
- "qdq_block_size": int(self.int4_block_size),
- "is_symmetric": extra_options.get("int4_is_symmetric", True),
- "op_types_to_quantize": extra_options.get("int4_op_types_to_quantize", ("MatMul",)),
- "nodes_to_exclude": extra_options.get("int4_nodes_to_exclude", []),
- "algo_config": int4_algo_config,
- },
+ "accuracy_level": int(extra_options.get("int4_accuracy_level", 4 if self.ep in ["cpu", "webgpu"] else 0)),
+ "qmoe_block_size": int(self.qmoe_block_size),
+ "qdq_block_size": int(self.matmul_block_size),
+ "is_symmetric": extra_options.get("int4_is_symmetric", True),
+ "op_types_to_quantize": extra_options.get("int4_op_types_to_quantize", ("MatMul",)),
+ "nodes_to_exclude": extra_options.get("int4_nodes_to_exclude", []),
+ "algo_config": algo_config,
"use_qdq": extra_options.get("use_qdq", False),
}
+ self.make_quant_init(config)
- # Propagate block_size to MoE/QMoE op when supported.
- # QMoE on supported EPs uses block-wise quantization via the 'block_size' attribute.
- # Ensure the attribute is set on the MoE op so runtime kernels can honor it.
- if self.moe_attrs.get("op_type") == "QMoE" and self.ep in supported_blockwise_eps:
- self.moe_attrs["block_size"] = int(self.qmoe_block_size)
- if self.quant_type is not None:
- # Create quantized attributes from quantization config
- self.quant_attrs["config"] = config.quantization_config
- self.quant_attrs["use_g_idx"] = (
- config.quantization_config["desc_act"] if "desc_act" in config.quantization_config else False
- )
+ # Initialize tied embeddings
+ self.make_tied_embeddings_init(config)
- # Determine if lm_head is unquantized. int4/8 can have options to int4_nodes_to_exclude. FP models are always unquantized.
- self.unquantized_lm_head = "/lm_head/MatMul" in self.quant_attrs["int4"]["nodes_to_exclude"] or self.onnx_dtype in {ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16}
- self.shared_embeddings = extra_options.get(
- "shared_embeddings",
- config.tie_word_embeddings
- if hasattr(config, "tie_word_embeddings") and config.tie_word_embeddings is not None
- else False,
- )
- self.int8_lm_head = extra_options.get("int4_algo_config", "default") in {"k_quant_mixed", "k_quant_last", "k_quant_linear", "rtn_last"}
+ def make_ep_expansions_init(self):
+ """
+ Replace the current class's methods with the appropriate expansion class's methods.
+
+ For an EP with specific subgraph requirements, this can be used to extend the current class
+ with additional functionality provided by the expansion class.
+ """
+ if self.ep == "trt-rtx":
+ from .expansions import TRT_RTX
- # shared_embeddings conflicts with exclude_embeds and exclude_lm_head
- if self.shared_embeddings and (self.exclude_embeds or self.exclude_lm_head):
- self.shared_embeddings = False
- elif self.shared_embeddings and not self.unquantized_lm_head:
- # matmul_nbits_quantizer.py has a different naming for default quantization, so lm_head.MatMul.weight_Q{}G{} does not match.
- self.shared_embeddings = self.int8_lm_head or extra_options.get("int4_algo_config", "default") in {"rtn", "k_quant"}
+ self.make_layernorm_subgraph = TRT_RTX.make_layernorm_subgraph.__get__(self, self.__class__)
+ self.make_skip_simplified_layer_norm = TRT_RTX.make_skip_simplified_layer_norm.__get__(self, self.__class__)
+ self.make_skip_layer_norm = TRT_RTX.make_skip_layer_norm.__get__(self, self.__class__)
+ self.make_simplified_layer_norm = TRT_RTX.make_simplified_layer_norm.__get__(self, self.__class__)
+ self.make_padded_cache = TRT_RTX.make_padded_cache.__get__(self, self.__class__)
+ self.make_split_if_nodes = TRT_RTX.make_split_if_nodes.__get__(self, self.__class__)
- def to_str_dtype(self, dtype: ir.DataType) -> str:
- return dtype.name
+ elif self.ep == "webgpu":
+ from .expansions import WebGPU
+
+ if self.extra_options.get("enable_webgpu_graph", False):
+ self.make_attention_mask_reformatting_for_gqa = (
+ WebGPU.make_attention_mask_graph_capture_reformatting_for_gqa.__get__(self, self.__class__)
+ )
+
+ else:
+ return
def make_inputs_init(self):
self.exclude_embeds = self.extra_options.get("exclude_embeds", False)
@@ -426,39 +401,7 @@ def make_outputs_init(self):
del self.output_names["logits"]
def make_rope_init(self, config):
- if "short_factor" in config.rope_scaling:
- # For models with multiple rotary embedding caches (e.g. Phi-3 mini 128K)
- self.rope_attrs["mscale_policy"] = config.rope_scaling["type"]
- short_factor = torch.tensor(config.rope_scaling["short_factor"], dtype=torch.float32)
- long_factor = torch.tensor(config.rope_scaling["long_factor"], dtype=torch.float32)
-
- short_mscale = config.rope_scaling["short_mscale"] if "short_mscale" in config.rope_scaling else 0
- long_mscale = config.rope_scaling["long_mscale"] if "long_mscale" in config.rope_scaling else 0
- short_mscale = short_mscale if short_mscale > 0 else self.make_mscale(self.context_length / self.original_context_length)
- long_mscale = long_mscale if long_mscale > 0 else self.make_mscale(self.context_length / self.original_context_length)
-
- self.rope_attrs["multi_cache"] = {
- "short_factor": short_factor, # Short factor when calculating `inv_freq` in rotary embeddings
- "long_factor": long_factor, # Long factor when calculating `inv_freq` in rotary embeddings
- "short_mscale": short_mscale, # Magnitude scaling for short factor when scaling `emb.cos()/emb.sin()` in rotary embeddings
- "long_mscale": long_mscale, # Magnitude scaling for long factor when scaling `emb.cos()/emb.sin()` in rotary embeddings
- }
-
- elif "low_freq_factor" in config.rope_scaling:
- # For models that rescale `inv_freq` using `low_freq_factor` and `high_freq_factor` (e.g. LLaMA-3.1)
- factor = config.rope_scaling["factor"] if "factor" in config.rope_scaling else 0
- low_freq_factor = config.rope_scaling["low_freq_factor"] if "low_freq_factor" in config.rope_scaling else 0
- high_freq_factor = (
- config.rope_scaling["high_freq_factor"] if "high_freq_factor" in config.rope_scaling else 0
- )
-
- self.rope_attrs["rescale_inv_freq"] = {
- "factor": factor, # Scale factor when calculating `new_freq` in rotary embeddings
- "low_freq_factor": low_freq_factor, # Low freq factor when calculating `low_freq_wavelen` in rotary embeddings
- "high_freq_factor": high_freq_factor, # High freq factor when calculating `high_freq_wavelen` in rotary embeddings
- }
-
- elif "beta_fast" in config.rope_scaling:
+ if "beta_fast" in config.rope_scaling:
# For models that use YARN (e.g. OpenAI OS-minier, Ministral3)
factor = config.rope_scaling["factor"] if "factor" in config.rope_scaling else 0
beta_slow = config.rope_scaling["beta_slow"] if "beta_slow" in config.rope_scaling else 0
@@ -531,7 +474,7 @@ def is_packed_attn_supported(self) -> bool:
def is_fused_rope_supported(self):
return self.ep not in ["dml"]
- def make_attention_init(self):
+ def make_attention_init(self, config):
self.q_size = self.num_attn_heads * self.head_size
self.kv_size = self.num_kv_heads * self.head_size
@@ -564,6 +507,30 @@ def make_attention_init(self):
self.past_present_share_buffer = self.attention_attrs["op_type"] == "GroupQueryAttention"
+ def make_lm_head_init(self, config):
+ pass
+
+ def make_quant_init(self, config):
+ if self.quant_type is not None:
+ # Create quantized attributes from quantization config
+ self.quant_attrs["config"] = config.quantization_config
+ self.quant_attrs["use_g_idx"] = (
+ config.quantization_config["desc_act"] if "desc_act" in config.quantization_config else False
+ )
+
+ def make_tied_embeddings_init(self, config):
+ # Determine if lm_head is unquantized. int4/8 can have options to int4_nodes_to_exclude. FP models are always unquantized.
+ self.unquantized_lm_head = "/lm_head/MatMul" in self.quant_attrs["nodes_to_exclude"] or self.onnx_dtype in {ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16}
+ self.shared_embeddings = self.extra_options.get("shared_embeddings", config.tie_word_embeddings if hasattr(config, "tie_word_embeddings") and config.tie_word_embeddings is not None else False)
+ self.int8_lm_head = self.extra_options.get("int4_algo_config", "default") in {"k_quant_mixed", "k_quant_last", "k_quant_linear", "rtn_last"}
+
+ # shared_embeddings conflicts with exclude_embeds and exclude_lm_head
+ if self.exclude_embeds or self.exclude_lm_head:
+ self.shared_embeddings = False
+ elif self.shared_embeddings and not self.unquantized_lm_head:
+ # matmul_nbits_quantizer.py has a different naming for default quantization, so lm_head.MatMul.weight_Q{}G{} does not match.
+ self.shared_embeddings = self.int8_lm_head or self.extra_options.get("int4_algo_config", "default") in {"rtn", "k_quant"}
+
def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir):
# Create config with attributes from config.json and generation_config.json (if latter file exists)
config = AutoConfig.from_pretrained(
@@ -689,7 +656,9 @@ def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir):
json.dump(genai_config, f, indent=4)
def update_genai_config(self, genai_config):
- """Override in subclasses to modify genai_config before it is written to disk."""
+ """
+ Override in subclasses to modify genai_config before it is written to disk.
+ """
pass
def make_key_value_cache_names(self, layer_id):
@@ -721,14 +690,14 @@ def save_processing(self, model_name_or_path, extra_kwargs, out_dir):
print(f"Saving processing files in {out_dir} for GenAI")
tokenizer.save_pretrained(out_dir)
- def make_int4_algo_config(self, quant_method: str):
+ def make_algo_config(self, quant_method: str):
customized_weight_config = {}
- int4_algo_config = None
+ algo_config = None
if quant_method in {"rtn", "rtn_last"}:
if quant_method == "rtn_last":
customized_weight_config["/lm_head/MatMul"] = {"bits": 8}
- int4_algo_config = RTNWeightOnlyQuantConfig(customized_weight_config=customized_weight_config)
+ algo_config = RTNWeightOnlyQuantConfig(customized_weight_config=customized_weight_config)
elif quant_method in {"k_quant", "k_quant_mixed", "k_quant_last", "k_quant_linear"}:
if quant_method != "k_quant":
@@ -762,20 +731,20 @@ def make_int4_algo_config(self, quant_method: str):
customized_weight_config[f"/model/layers.{i}/mlp/{proj}/MatMul"] = {"bits": 8}
customized_weight_config["/lm_head/MatMul"] = {"bits": 8}
- int4_algo_config = KQuantWeightOnlyQuantConfig(customized_weight_config=customized_weight_config)
+ algo_config = KQuantWeightOnlyQuantConfig(customized_weight_config=customized_weight_config)
- return int4_algo_config
+ return algo_config
def to_int4(self) -> ir.Model:
quant = MatMulNBitsQuantizer(
model=ir.to_proto(self.model),
- block_size=self.quant_attrs["int4"]["qdq_block_size"],
- is_symmetric=self.quant_attrs["int4"]["is_symmetric"],
- accuracy_level=self.quant_attrs["int4"]["accuracy_level"],
- nodes_to_exclude=self.quant_attrs["int4"]["nodes_to_exclude"],
+ block_size=self.quant_attrs["qdq_block_size"],
+ is_symmetric=self.quant_attrs["is_symmetric"],
+ accuracy_level=self.quant_attrs["accuracy_level"],
+ nodes_to_exclude=self.quant_attrs["nodes_to_exclude"],
quant_format=QuantFormat.QDQ if self.quant_attrs["use_qdq"] else QuantFormat.QOperator,
- op_types_to_quantize=self.quant_attrs["int4"]["op_types_to_quantize"],
- algo_config=self.quant_attrs["int4"]["algo_config"],
+ op_types_to_quantize=self.quant_attrs["op_types_to_quantize"],
+ algo_config=self.quant_attrs["algo_config"],
)
quant.process()
return ir.from_proto(quant.model.model)
@@ -828,6 +797,9 @@ def callback(tensor: ir.TensorProtocol, metadata: dict):
if not os.listdir(self.cache_dir):
os.rmdir(self.cache_dir)
+ def to_str_dtype(self, dtype: ir.DataType) -> str:
+ return dtype.name
+
def make_initializer(self, tensor: torch.Tensor | np.ndarray | ir.TensorProtocol, /, name: str, to: ir.DataType | None = None):
if to is not None:
# Cast the tensor lazily if `to` is provided
@@ -1207,7 +1179,7 @@ def make_matmul_int4(self, matmul, basename, root_input, **kwargs):
outputs=[output],
name=name,
domain="com.microsoft",
- accuracy_level=self.quant_attrs["int4"]["accuracy_level"],
+ accuracy_level=self.quant_attrs["accuracy_level"],
bits=matmul.bits,
block_size=matmul.group_size,
K=matmul.in_features,
@@ -1439,6 +1411,7 @@ def make_packed_add(self, q_add, k_add, v_add, name, root_input, **kwargs):
def make_embedding(self, embedding):
basename = "/model/embed_tokens"
+
# Use GatherBlockQuantized if and only if tied embeddings are enabled and export model is quantized. quantized d_type in set_onnx_dtype is INT4/UINT4
if self.shared_embeddings and self.onnx_dtype in {ir.DataType.INT4, ir.DataType.UINT4}:
gather_name = f"{basename}/GatherBlockQuantized"
@@ -1448,18 +1421,17 @@ def make_embedding(self, embedding):
bits = 8 if self.int8_lm_head else 4
flat_dim = self.hidden_size * bits // 8
weight_reshape_inputs = [
- f"lm_head.MatMul.weight_Q{bits}G{self.int4_block_size}",
+ f"lm_head.MatMul.weight_Q{bits}G{self.matmul_block_size}",
f"/model/constants/INT64/[{self.vocab_size}, {flat_dim}]",
]
weight_reshape_output = f"{weight_reshape_name}/output_0"
- # quantized weight dtype is uint8, see here
+ # Quantized weight dtype is uint8. See here for more info:
# https://github.com/microsoft/onnxruntime/blob/0c9356cb986fd4cd2c5d510909d31186010ba226/onnxruntime/python/tools/quantization/neural_compressor/weight_only.py#L73
- self.make_reshape(
- weight_reshape_name, weight_reshape_inputs, dtype=ir.DataType.UINT8, shape=[self.vocab_size, flat_dim]
- )
+ self.make_reshape(weight_reshape_name, weight_reshape_inputs, dtype=ir.DataType.UINT8, shape=[self.vocab_size, flat_dim])
input_names = [weight_reshape_output, self.input_names["input_ids"], "lm_head.MatMul.weight_scale"];
- if not self.quant_attrs["int4"]["is_symmetric"]:
+ if not self.quant_attrs["is_symmetric"]:
input_names.append("lm_head.MatMul.weight_zp")
+
self.make_node(
"GatherBlockQuantized",
inputs=input_names,
@@ -1467,10 +1439,11 @@ def make_embedding(self, embedding):
name=gather_name,
domain="com.microsoft",
bits=bits,
- block_size=int(self.int4_block_size),
+ block_size=int(self.matmul_block_size),
gather_axis=0,
quantize_axis=1,
)
+
# Use Transpose + Gather for tied embeddings for float embedding layers
elif self.shared_embeddings and self.unquantized_lm_head:
transpose_name = f"{basename}/Transpose"
@@ -1486,6 +1459,7 @@ def make_embedding(self, embedding):
gather_name = f"{basename}/Gather"
gather_output = f"{gather_name}/output_0"
self.make_node("Gather", inputs=[transpose_output, self.input_names["input_ids"]], outputs=[gather_output], name=gather_name)
+
else:
weight = "model.embed_tokens.weight"
self.make_initializer(embedding, weight, to=self.io_dtype)
@@ -1526,13 +1500,6 @@ def make_embedding(self, embedding):
self.layernorm_attrs["skip_input"] = layernorm_attrs_value
def make_layernorm(self, layer_id, layernorm, skip, simple, location):
- if self.ep == "trt-rtx" and (skip or simple):
- # Fall back to primitive ops
- self._make_layernorm_op(layer_id, layernorm, skip, simple, location)
- else:
- self.make_layernorm_op(layer_id, layernorm, skip, simple, location)
-
- def make_layernorm_op(self, layer_id, layernorm, skip, simple, location):
root_input = self.layernorm_attrs["root_input"]
skip_input = self.layernorm_attrs["skip_input"]
@@ -1571,8 +1538,16 @@ def make_layernorm_op(self, layer_id, layernorm, skip, simple, location):
if cast:
inputs, outputs = self.make_layernorm_casts(name, inputs, outputs, old_io_dtype, new_io_dtype)
- # Make op and its shape
- self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
+ # Make op subgraph and its shapes
+ self.make_layernorm_subgraph(
+ name,
+ op_type=op_type,
+ inputs=inputs,
+ outputs=outputs,
+ skip=skip,
+ new_io_dtype=new_io_dtype,
+ **kwargs,
+ )
if not use_hidden_states_as_output:
# Add shape only if not graph output
self.make_value(outputs[0], new_io_dtype, shape=["batch_size", "sequence_length", self.hidden_size])
@@ -1587,93 +1562,6 @@ def make_layernorm_op(self, layer_id, layernorm, skip, simple, location):
# Assign output 3 of current SkipLayerNorm as root input to next SkipLayerNorm
self.layernorm_attrs["root_input"] = output_3
- def _make_layernorm_op(self, layer_id, layernorm, skip, simple, location):
- root_input = self.layernorm_attrs["root_input"]
- skip_input = self.layernorm_attrs["skip_input"]
-
- # Get precision types to use
- old_io_dtype = self.io_dtype
- new_io_dtype = ir.DataType.FLOAT if self.layernorm_attrs["cast"]["use_fp32"] else self.io_dtype
- cast = old_io_dtype != new_io_dtype
-
- # Create weight and bias tensors
- weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
- self.make_initializer(layernorm.weight + self.layernorm_attrs["add_offset"], weight, to=new_io_dtype)
- bias = f"model.layers.{layer_id}.{location}_layernorm.bias"
- if not simple:
- self.make_initializer(layernorm.bias, bias, to=new_io_dtype)
-
- # Create input names for op
- inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
- if not simple:
- inputs.append(bias)
-
- name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
- op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
- kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
- if not skip:
- kwargs.update({"axis": -1, "stash_type": 1})
-
- # Create output names for op
- output_0 = f"/model/layers.{layer_id}/{location}_layernorm/output_0"
- output_3 = f"/model/layers.{layer_id}/{location}_layernorm/output_3"
- use_hidden_states_as_output = self.layernorm_attrs["last_layernorm"] and (self.include_hidden_states or self.exclude_lm_head)
- if use_hidden_states_as_output:
- output_0 = self.output_names["hidden_states"]
- outputs = [output_0, "", "", output_3] if skip and not self.layernorm_attrs["last_layernorm"] else [output_0]
-
- # Create Cast nodes for inputs and outputs if old_dtype != new_dtype
- if cast:
- inputs, outputs = self.make_layernorm_casts(name, inputs, outputs, old_io_dtype, new_io_dtype)
- root_input = inputs[0]
- skip_input = inputs[1] if skip else None
-
- if op_type == "SimplifiedLayerNormalization":
- self._make_simplified_layer_norm(
- name,
- root_input,
- weight,
- outputs[0],
- new_io_dtype,
- shape=["batch_size", "sequence_length", self.hidden_size],
- )
- elif op_type == "SkipSimplifiedLayerNormalization":
- self._make_skip_simplified_layer_norm(
- name,
- root_input,
- skip_input,
- weight,
- outputs[0],
- output_3,
- new_io_dtype,
- shape=["batch_size", "sequence_length", self.hidden_size],
- )
- elif op_type == "SkipLayerNormalization":
- self._make_skip_layer_norm(
- name,
- root_input,
- skip_input,
- weight,
- bias,
- outputs[0],
- output_3,
- new_io_dtype,
- shape=["batch_size", "sequence_length", self.hidden_size],
- )
- else:
- raise ValueError(f"Invalid op_type: {op_type}")
-
- if skip and not self.layernorm_attrs["last_layernorm"]:
- self.make_value(outputs[3], new_io_dtype, shape=["batch_size", "sequence_length", self.hidden_size])
-
- # Update LayerNorm attributes
- self.layernorm_attrs["output_0"] = output_0
- if skip and not self.layernorm_attrs["last_layernorm"]:
- self.layernorm_attrs["output_3"] = output_3
-
- # Assign output 3 of current SkipLayerNorm as root input to next SkipLayerNorm
- self.layernorm_attrs["root_input"] = output_3
-
def make_layernorm_casts(self, name, inputs, outputs, old_dtype, new_dtype):
# Name = name of original LayerNorm op as if the cast nodes did not exist
# Inputs = inputs into the original LayerNorm op as if the cast nodes did not exist
@@ -1732,6 +1620,21 @@ def make_layernorm_casts(self, name, inputs, outputs, old_dtype, new_dtype):
return inputs, outputs
+ def make_layernorm_subgraph(self, name, **kwargs):
+ # This method can be used to create multiple LayerNorm operations
+ op_type = kwargs.pop("op_type")
+ inputs = kwargs.pop("inputs")
+ outputs = kwargs.pop("outputs")
+ skip = kwargs.pop("skip")
+ new_io_dtype = kwargs.pop("new_io_dtype")
+
+ # Create LayerNorm op
+ self.make_layernorm_op(name, op_type, inputs, outputs, skip, new_io_dtype, **kwargs)
+
+ def make_layernorm_op(self, name, op_type, inputs, outputs, skip, new_io_dtype, **kwargs):
+ # Create the LayerNorm, SimplifiedLayerNorm, SkipLayerNorm, or SkipSimplifiedLayerNorm op
+ self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
+
def make_mscale_su(self, mscale):
if mscale <= 1.0:
return 1.0
@@ -1873,157 +1776,6 @@ def make_rotary_embedding_caches(self, **kwargs):
return cos_cache_name, sin_cache_name
- def make_padded_cache(self, small_cache, large_cache, pad_value=0.0):
- """Pad small cache to match large cache shape for uniform If node branches.
-
- This is used for TRT-RTX EP which requires uniform dimensions in both branches of If nodes.
-
- Args:
- small_cache: The smaller cache tensor to pad
- large_cache: The larger cache tensor (defines target shape)
- pad_value: Value to use for padding (1.0 for cos_cache, 0.0 for sin_cache)
- """
- target_shape = large_cache.shape
- if small_cache.shape == target_shape:
- return small_cache
-
- # Create padded tensor filled with pad_value
- padded_cache = torch.full(target_shape, pad_value, dtype=small_cache.dtype)
- # Copy original data to the beginning
- padded_cache[: small_cache.shape[0], :] = small_cache
- return padded_cache
-
- def _make_split_if_nodes_for_trt_rtx(
- self,
- basename,
- greater_name,
- cos_cache_name,
- sin_cache_name,
- cos_cache_large,
- sin_cache_large,
- cos_cache_small,
- sin_cache_small,
- cos_cache_large_name,
- sin_cache_large_name,
- cos_cache_small_name,
- sin_cache_small_name,
- small_cache_shape,
- ):
- """Create split If nodes for TRT-RTX to workaround trt-rtx multi-output bug.
-
- This is a TEMPORARY workaround for TRT-RTX bug where If nodes with
- multiple outputs
-
- Creates two separate If nodes instead of one:
- - {basename}/cos/If: Outputs cos_cache only
- - {basename}/sin/If: Outputs sin_cache only
-
- Both If nodes use the same condition and independently select their respective caches.
- """
- cos_if_name = f"{basename}/cos/If"
-
- cos_large_for_split = ir.node(
- "Constant",
- [],
- outputs=[
- ir.Value(
- name=f"{cos_cache_large_name}_split",
- type=ir.TensorType(self.io_dtype),
- shape=ir.Shape(cos_cache_large.shape),
- )
- ],
- name="/large/cos_cache/Constant_split_cos",
- attributes=dict(value=ir.tensor(cos_cache_large)),
- )
-
- cos_small_for_split = ir.node(
- "Constant",
- [],
- outputs=[
- ir.Value(
- name=f"{cos_cache_small_name}_split",
- type=ir.TensorType(self.io_dtype),
- shape=ir.Shape(small_cache_shape),
- )
- ],
- name="/small/cos_cache/Constant_split_cos",
- attributes=dict(value=ir.tensor(cos_cache_small)),
- )
-
- self.make_node(
- "If",
- inputs=[f"{greater_name}/output_0"],
- outputs=[cos_cache_name],
- name=cos_if_name,
- then_branch=ir.Graph(
- inputs=[],
- outputs=[cos_large_for_split.outputs[0]],
- nodes=[cos_large_for_split],
- name="large_cos_cache_graph",
- ),
- else_branch=ir.Graph(
- inputs=[],
- outputs=[cos_small_for_split.outputs[0]],
- nodes=[cos_small_for_split],
- name="small_cos_cache_graph",
- ),
- )
-
- # Create separate If node for sin_cache only
- sin_if_name = f"{basename}/sin/If"
-
- # Create unique constant nodes for sin to avoid tensor sharing
- sin_large_for_split = ir.node(
- "Constant",
- [],
- outputs=[
- ir.Value(
- name=f"{sin_cache_large_name}_split",
- type=ir.TensorType(self.io_dtype),
- shape=ir.Shape(sin_cache_large.shape),
- )
- ],
- name="/large/sin_cache/Constant_split_sin",
- attributes=dict(value=ir.tensor(sin_cache_large)),
- )
-
- sin_small_for_split = ir.node(
- "Constant",
- [],
- outputs=[
- ir.Value(
- name=f"{sin_cache_small_name}_split",
- type=ir.TensorType(self.io_dtype),
- shape=ir.Shape(small_cache_shape),
- )
- ],
- name="/small/sin_cache/Constant_split_sin",
- attributes=dict(value=ir.tensor(sin_cache_small)),
- )
-
- self.make_node(
- "If",
- inputs=[f"{greater_name}/output_0"],
- outputs=[sin_cache_name],
- name=sin_if_name,
- then_branch=ir.Graph(
- inputs=[],
- outputs=[sin_large_for_split.outputs[0]],
- nodes=[sin_large_for_split],
- name="large_sin_cache_graph",
- ),
- else_branch=ir.Graph(
- inputs=[],
- outputs=[sin_small_for_split.outputs[0]],
- nodes=[sin_small_for_split],
- name="small_sin_cache_graph",
- ),
- )
-
- # Create output values
- self.make_value(cos_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
- self.make_value(sin_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
-
def make_rotary_embedding(self, name, root_input, **kwargs):
cos_cache_name, sin_cache_name = self.make_rotary_embedding_caches()
num_heads = self.num_kv_heads if "k_rotary" in name else self.num_attn_heads
@@ -2098,7 +1850,7 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
sin_cache_small = self.make_padded_cache(sin_cache_small, sin_cache_large, pad_value=0.0)
# Create Greater condition node for If nodes
- basename = "/model/rotemb_caches_subgraph"
+ basename = "/model/rope_caches_subgraph"
gather_name = ""
if self.attention_attrs["op_type"] == "GroupQueryAttention":
gather_name = "/model/attn_mask_reformat/attn_mask_subgraph/Gather"
@@ -2110,7 +1862,7 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
self.make_greater(greater_name, greater_inputs, shape=[])
# Create split If nodes and return early
- self._make_split_if_nodes_for_trt_rtx(
+ self.make_split_if_nodes(
basename=basename,
greater_name=greater_name,
cos_cache_name=cos_cache_name,
@@ -2128,14 +1880,14 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
self.ep_attrs["trt-rtx"]["enable_cuda_graph"] = "0"
return
- # For other EPs (CUDA, CPU, WebGPU), create regular If node with multiple outputs
+ # For other EPs (CPU, CUDA, WebGPU), create regular If node with multiple outputs
# Make the following subgraph to decide which cos/sin caches to use in the rotary embeddings
#
# attention_mask --> Shape --> Gather --> Greater --> If --> (cos_cache, sin_cache)
# (idx=1)
#
- basename = "/model/rotemb_caches_subgraph"
+ basename = "/model/rope_caches_subgraph"
gather_name = ""
if self.attention_attrs["op_type"] == "GroupQueryAttention":
gather_name = "/model/attn_mask_reformat/attn_mask_subgraph/Gather"
@@ -2208,7 +1960,7 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
cos_cache_large_node,
sin_cache_large_node,
],
- name="large_rotemb_caches_graph",
+ name="large_rope_caches_graph",
),
else_branch=ir.Graph(
inputs=[],
@@ -2220,127 +1972,12 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
cos_cache_small_node,
sin_cache_small_node,
],
- name="small_rotemb_caches_graph",
+ name="small_rope_caches_graph",
),
)
self.make_value(cos_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
self.make_value(sin_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
- # This expansion of contrib-op can be updated / deprecated in future.
- def _make_skip_simplified_layer_norm(
- self, basename, root_input, skip_input, weight_name, output_0, output_3, io_dtype, shape
- ):
- # root_input skip_input
- # | |
- # +------------------+
- # |
- # Add-------------> output (1)
- # |
- # SimplifiedLayerNorm----> output (0)
- make_add_name = f"{basename}/Add"
- output_3 = f"{basename}/Add/output_0" if output_3 is None else output_3
- self.make_node("Add", inputs=[root_input, skip_input], outputs=[output_3], name=make_add_name)
- self.make_value(output_3, io_dtype, shape=["batch_size", "sequence_length", self.hidden_size])
-
- make_simplified_layer_norm_name = f"{basename}/skip_simplified_layer_norm"
- self._make_simplified_layer_norm(
- make_simplified_layer_norm_name, output_3, weight_name, output_0, io_dtype, shape=shape
- )
-
- # This expansion contrib-op can be updated / deprecated in the future.
- def _make_skip_layer_norm(
- self, basename, root_input, skip_input, weight_name, bias_name, output_0, output_3, io_dtype, shape
- ):
- # root_input skip_input
- # | |
- # +------------------+
- # |
- # Add-------------> output (1)
- # |
- # LayerNormalization-----> output (0)
- output_3 = f"{basename}/Add/output_0" if output_3 is None else output_3
- make_add_name = f"{basename}/Add"
- self.make_node("Add", inputs=[root_input, skip_input], outputs=[output_3], name=make_add_name)
- self.make_value(output_3, io_dtype, shape=["batch_size", "sequence_length", self.hidden_size])
-
- make_layer_norm_name = f"{basename}/LayerNormalization"
- inputs = [output_3, weight_name, bias_name]
-
- kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
- kwargs.update({"axis": -1, "stash_type": 1})
-
- self.make_node("LayerNormalization", inputs=inputs, outputs=[output_0], name=make_layer_norm_name, **kwargs)
- self.make_value(output_0, io_dtype, shape=shape)
-
- # This expansion contrib-op can be updated / deprecated in the future.
- def _make_simplified_layer_norm(self, basename, root_input, weight_name, output_0, io_dtype, shape):
- # Cast (float32) - most calc happens in higher precision
- # |
- # +-------+-------+
- # | |
- # Pow |
- # | |
- # ReduceMean |
- # | |
- # Add |
- # | |
- # Sqrt |
- # | |
- # Div |
- # | |
- # +-------+-------+
- # |
- # Mul
- # |
- # Cast_1 (io_dtype - float16)
- # |
- # Mul_1
-
- make_cast_name = f"{basename}/Cast"
- self.make_cast(make_cast_name, root_input, ir.DataType.FLOAT, shape=shape)
-
- make_pow_name = f"{basename}/Pow"
- make_pow_inputs = [f"{make_cast_name}/output_0", "/model/constants/FLOAT/2"]
-
- self.make_node(
- "Pow", inputs=make_pow_inputs, outputs=[f"{make_pow_name}/output_0"], name=make_pow_name, domain=""
- )
- self.make_value(f"{make_pow_name}/output_0", ir.DataType.FLOAT, shape=shape)
-
- make_reducemean_name = f"{basename}/ReduceMean"
- make_reducemean_inputs = [f"{make_pow_name}/output_0", "/model/constants/INT64/[-1]"]
- self.make_reduce_mean(
- make_reducemean_name, make_reducemean_inputs, ir.DataType.FLOAT, keepdims=True, shape=shape
- )
-
- make_add_name = f"{basename}/Add"
- make_add_inputs = [
- f"{make_reducemean_name}/output_0",
- f"/model/constants/FLOAT/{self.layernorm_attrs['epsilon']}",
- ]
- self.make_add(make_add_name, make_add_inputs, ir.DataType.FLOAT, shape=shape)
-
- make_sqrt_name = f"{basename}/Sqrt"
- make_sqrt_inputs = [f"{make_add_name}/output_0"]
- self.make_sqrt(make_sqrt_name, make_sqrt_inputs, ir.DataType.FLOAT, shape=shape)
-
- make_div_name = f"{basename}/Div"
- make_div_inputs = ["/model/constants/FLOAT/1", f"{make_sqrt_name}/output_0"]
- self.make_div(make_div_name, make_div_inputs, ir.DataType.FLOAT, shape=shape)
-
- make_mul_name = f"{basename}/Mul"
- make_mul_inputs = [f"{make_div_name}/output_0", f"{make_cast_name}/output_0"]
- self.make_mul(make_mul_name, make_mul_inputs, ir.DataType.FLOAT, shape=shape)
-
- make_cast_1_name = f"{basename}/Cast_1"
- self.make_cast(make_cast_1_name, f"{make_mul_name}/output_0", dtype=io_dtype, shape=shape)
-
- make_mul_1_name = f"{basename}/Mul_1"
- make_mul_1_inputs = [f"{make_cast_1_name}/output_0", weight_name]
-
- self.make_node("Mul", inputs=make_mul_1_inputs, outputs=[output_0], name=make_mul_1_name)
- self.make_value(output_0, dtype=io_dtype, shape=shape)
-
def make_qk_norm(self, layer_id, attention):
# Make subgraph to compute SimplifiedLayerNorm after Q and K MatMuls in attention:
#
@@ -3607,12 +3244,12 @@ def make_qmoe_weights(self, weights):
qweight, scales = None, None
# Use block-wise quantization for supported EPs when qmoe_block_size > 0.
- # TRT-RTX defaults to 128; others default to 32.
- supported_blockwise_eps = ["cpu", "webgpu", "trt-rtx"]
+ # CUDA and TRT-RTX default to 128; others default to 32.
+ supported_blockwise_eps = ["cpu", "cuda", "webgpu", "trt-rtx"]
use_blockwise_quant = self.ep in supported_blockwise_eps and self.qmoe_block_size > 0
if use_blockwise_quant:
- block_size = self.quant_attrs["int4"]["qmoe_block_size"]
+ block_size = self.quant_attrs["qmoe_block_size"]
try:
qweight, scales = self._symmetric_blockwise_quantize(weights, block_size)
self.moe_attrs["block_size"] = block_size
@@ -3904,6 +3541,8 @@ def make_activation(self, layer_id, root_input):
return output_name
def make_lm_head(self, lm_head):
+ basename = "/lm_head"
+
# Check if there are ops to insert after MatMul
bias_exists = lm_head.bias is not None
scale_exists = self.lm_head_attrs["scale"] != 1
@@ -3915,7 +3554,7 @@ def make_lm_head(self, lm_head):
# Add new checks to the end of the list and after the below if condition checks.
exists_checks = [bias_exists, scale_exists, mask_exists, softcap_exists, cast_exists]
- matmul_basename = "/lm_head/MatMul"
+ matmul_basename = f"{basename}/MatMul"
root_input = self.layernorm_attrs["output_0"]
# Sequence dimension for shape annotations ("sequence_length" normally, 1 when pruned)
@@ -3928,23 +3567,12 @@ def make_lm_head(self, lm_head):
seq_dim = 1
# Gather: [B, S, H] + scalar(-1) -> [B, H]
- gather_name = "/lm_head/prune/Gather"
- self.make_gather(
- gather_name,
- inputs=[root_input, "/model/constants/INT64/-1"],
- dtype=self.io_dtype,
- shape=["batch_size", self.hidden_size],
- axis=1,
- )
+ gather_name = f"{basename}/prune/Gather"
+ self.make_gather(gather_name, inputs=[root_input, "/model/constants/INT64/-1"], dtype=self.io_dtype, shape=["batch_size", self.hidden_size], axis=1)
# Unsqueeze: [B, H] -> [B, 1, H]
- unsqueeze_name = "/lm_head/prune/Unsqueeze"
- self.make_unsqueeze(
- unsqueeze_name,
- inputs=[f"{gather_name}/output_0", "/model/constants/INT64/[1]"],
- dtype=self.io_dtype,
- shape=["batch_size", 1, self.hidden_size],
- )
+ unsqueeze_name = f"{basename}/prune/Unsqueeze"
+ self.make_unsqueeze(unsqueeze_name, inputs=[f"{gather_name}/output_0", "/model/constants/INT64/[1]"], dtype=self.io_dtype, shape=["batch_size", 1, self.hidden_size])
root_input = f"{unsqueeze_name}/output_0"
@@ -3955,18 +3583,13 @@ def make_lm_head(self, lm_head):
lm_name = matmul_name
if bias_exists:
- add_name = "/lm_head/Add"
- self.make_add_bias(
- lm_head.bias, add_name, root_input=f"{lm_name}/output_0", logits=not any(exists_checks[1:]), seq_dim=seq_dim
- )
+ add_name = f"{basename}/Add"
+ self.make_add_bias(lm_head.bias, add_name, root_input=f"{lm_name}/output_0", logits=not any(exists_checks[1:]), seq_dim=seq_dim)
lm_name = add_name
if scale_exists:
- mul_name = "/lm_head/Mul"
- mul_inputs = [
- f"{lm_name}/output_0",
- f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{self.lm_head_attrs['scale']}",
- ]
+ mul_name = f"{basename}/Mul"
+ mul_inputs = [f"{lm_name}/output_0", f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{self.lm_head_attrs['scale']}"]
mul_output = "logits" if not any(exists_checks[2:]) else f"{mul_name}/output_0"
self.make_node("Mul", inputs=mul_inputs, outputs=[mul_output], name=mul_name)
self.make_value(mul_output, self.io_dtype, shape=["batch_size", seq_dim, self.vocab_size])
@@ -3977,12 +3600,8 @@ def make_lm_head(self, lm_head):
logits_mask_name = "logits_mask"
self.make_initializer(self.lm_head_attrs["mask"], logits_mask_name)
- where_name = "/lm_head/Where"
- where_inputs = [
- logits_mask_name,
- f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{torch.finfo(to_torch_dtype(self.io_dtype)).min}",
- f"{lm_name}/output_0",
- ]
+ where_name = f"{basename}/Where"
+ where_inputs = [logits_mask_name, f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{torch.finfo(to_torch_dtype(self.io_dtype)).min}", f"{lm_name}/output_0"]
where_output = "logits" if not any(exists_checks[3:]) else f"{where_name}/output_0"
self.make_node("Where", inputs=where_inputs, outputs=[where_output], name=where_name)
self.make_value(where_output, self.io_dtype, shape=["batch_size", seq_dim, self.vocab_size])
@@ -3990,28 +3609,15 @@ def make_lm_head(self, lm_head):
if softcap_exists:
# Add final logit softcapping (Div --> Tanh --> Mul)
- div_name = "/lm_head/softcap/Div"
- div_inputs = [
- f"{lm_name}/output_0",
- f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{self.lm_head_attrs['softcap']}",
- ]
- self.make_div(
- div_name, div_inputs, dtype=self.io_dtype, shape=["batch_size", seq_dim, self.vocab_size]
- )
+ div_name = f"{basename}/softcap/Div"
+ div_inputs = [f"{lm_name}/output_0", f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{self.lm_head_attrs['softcap']}"]
+ self.make_div(div_name, div_inputs, dtype=self.io_dtype, shape=["batch_size", seq_dim, self.vocab_size])
- tanh_name = "/lm_head/softcap/Tanh"
- self.make_tanh(
- tanh_name,
- f"{div_name}/output_0",
- dtype=self.io_dtype,
- shape=["batch_size", seq_dim, self.vocab_size],
- )
+ tanh_name = f"{basename}/softcap/Tanh"
+ self.make_tanh(tanh_name, f"{div_name}/output_0", dtype=self.io_dtype, shape=["batch_size", seq_dim, self.vocab_size])
- mul_name = "/lm_head/softcap/Mul"
- mul_inputs = [
- f"{tanh_name}/output_0",
- f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{self.lm_head_attrs['softcap']}",
- ]
+ mul_name = f"{basename}/softcap/Mul"
+ mul_inputs = [f"{tanh_name}/output_0", f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{self.lm_head_attrs['softcap']}"]
mul_output = "logits" if not any(exists_checks[4:]) else f"{mul_name}/output_0"
self.make_node("Mul", inputs=mul_inputs, outputs=[mul_output], name=mul_name)
self.make_value(mul_output, self.io_dtype, shape=["batch_size", seq_dim, self.vocab_size])
@@ -4019,37 +3625,17 @@ def make_lm_head(self, lm_head):
if cast_exists:
# Add final cast from io_dtype to logits_dtype
- cast_name = "/lm_head/Cast"
+ cast_name = f"{basename}/Cast"
cast_output = "logits"
- self.make_node(
- "Cast",
- inputs=[f"{lm_name}/output_0"],
- outputs=[cast_output],
- name=cast_name,
- to=self.output_types["logits"],
- )
- self.make_value(
- cast_output, self.output_types["logits"], shape=["batch_size", seq_dim, self.vocab_size]
- )
+ self.make_node("Cast", inputs=[f"{lm_name}/output_0"], outputs=[cast_output], name=cast_name, to=self.output_types["logits"])
+ self.make_value(cast_output, self.output_types["logits"], shape=["batch_size", seq_dim, self.vocab_size])
def make_layer(self, layer_id, layer):
# Each LLM decoder layer is typically defined as:
# input_layernorm --> attention --> output_layernorm --> MLP
- self.make_layernorm(
- layer_id,
- layer.input_layernorm,
- skip=not self.layernorm_attrs["first_layernorm"],
- simple=self.layernorm_attrs["simple"],
- location="input",
- )
+ self.make_layernorm(layer_id, layer.input_layernorm, skip=not self.layernorm_attrs["first_layernorm"], simple=self.layernorm_attrs["simple"], location="input")
self.make_attention(layer_id, layer.self_attn, root_input=self.layernorm_attrs["output_0"])
- self.make_layernorm(
- layer_id,
- layer.post_attention_layernorm,
- skip=True,
- simple=self.layernorm_attrs["simple"],
- location="post_attention",
- )
+ self.make_layernorm(layer_id, layer.post_attention_layernorm, skip=True, simple=self.layernorm_attrs["simple"], location="post_attention")
self.make_mlp(layer_id, layer.mlp, root_input=self.layernorm_attrs["output_0"])
self.layernorm_attrs["first_layernorm"] = False
@@ -4083,6 +3669,7 @@ def load_weights(self, input_path):
from quantized_model import QuantModel
except ImportError:
from onnxruntime_genai.models.quantized_model import QuantModel
+
q_size = self.num_attn_heads * self.head_size
kv_size = self.num_kv_heads * self.head_size
model = QuantModel.from_pretrained(
@@ -4242,19 +3829,6 @@ def make_preprocessing_nodes(self):
self.make_attention_mask_reformatting()
def make_attention_mask_reformatting(self):
- if (
- self.extra_options.get("enable_cuda_graph", False)
- or self.extra_options.get("enable_webgpu_graph", False)
- or self.ep == "dml"
- ):
- # ORT does not allow nodes to be placed on mulitple execution providers
- # with graph capture enabled. We've only verified it works with GQA and with
- # past_present_share_buffer enabled(so the total_seq_len in GQA is hardcoded
- # to a fixed value by logic).
- # For other models, we need to check if it works and update the logic here.
- # This assertion is temporary.
- assert self.past_present_share_buffer
-
if self.attention_attrs["op_type"] == "GroupQueryAttention":
self.make_attention_mask_reformatting_for_gqa()
elif self.attention_attrs["op_type"] == "MultiHeadAttention":
@@ -4267,9 +3841,6 @@ def make_attention_mask_reformatting(self):
# 4D causal attention mask
self.make_attention_mask_reformatting_for_mha()
- if self.attention_attrs["block_sparse"]["sparse_block_size"] != 0:
- self.make_attention_mask_reformatting_for_sparse_attn()
-
def make_attention_mask_reformatting_for_mha(self):
# Make nodes for the attention mask subgraphs that reformat the
# 2D attention mask (B, S) to 4D causal attention mask (B, N, S, T)
@@ -4625,49 +4196,7 @@ def make_common_mask_reformat_subgraph(
return expand_name
- def make_attention_mask_graph_capture_reformatting_for_gqa(self, attn_mask_basename):
- # Make nodes for the attention mask subgraph that calculates
- # attributes about the 2D attention mask to use in GroupQueryAttention
- #
- # Key difference vs make_attention_mask_standard_reformatting_for_gqa:
- # - Standard mode: total_seq_len is calculated from Shape op (always runs on CPU)
- # - Graph capture mode: No Shape ops inserted to ensure all ops run on GPU (no CPU ops)
- #
- # attention_mask
- # |
- # Cast to int32
- # |
- # ReduceSum (keepdims=0)
- # / \
- # / \
- # Sub ReduceMax
- # | |
- # seqlens_k total_seq_len
- # (1D) (int)
-
- # Calculate ReduceSum from attention_mask
- cast_1_name = f"{attn_mask_basename}/Cast"
- self.make_cast(
- cast_1_name, self.input_names["attention_mask"], dtype=ir.DataType.INT32, shape=["batch_size", "total_sequence_length"]
- )
- reduce_sum_name = f"{attn_mask_basename}/ReduceSum"
- reduce_sum_inputs = [f"{cast_1_name}/output_0", "/model/constants/INT64/[1]"]
- self.make_reduce_sum(reduce_sum_name, reduce_sum_inputs, dtype=ir.DataType.INT32, shape=["batch_size"])
-
- # Left branch: Calculate seqlens_k = ReduceSum - 1
- sub_name = f"{attn_mask_basename}/Sub"
- sub_inputs = [f"{reduce_sum_name}/output_0", "/model/constants/INT32/[1]"]
- self.make_sub(sub_name, sub_inputs, dtype=ir.DataType.INT32, shape=["batch_size"])
-
- # Right branch: ReduceMax to get maximum int value for total_seq_len
- reduce_max_name = f"{attn_mask_basename}/ReduceMax"
- reduce_max_inputs = [f"{reduce_sum_name}/output_0"]
- self.make_reduce_max(reduce_max_name, reduce_max_inputs, dtype=ir.DataType.INT32, shape=[])
-
- self.mask_attrs["seqlens_k"] = sub_name
- self.mask_attrs["total_seq_len"] = reduce_max_name
-
- def make_attention_mask_standard_reformatting_for_gqa(self, attn_mask_basename):
+ def make_attention_mask_reformatting_for_gqa(self):
# Make nodes for the attention mask subgraph that calculates
# attributes about the 2D attention mask to use in GroupQueryAttention
#
@@ -4682,6 +4211,8 @@ def make_attention_mask_standard_reformatting_for_gqa(self, attn_mask_basename):
# | |
# seqlens_k total_seq_len
# (1D) (int)
+ basename = "/model/attn_mask_reformat"
+ attn_mask_basename = f"{basename}/attn_mask_subgraph"
# Left path
reduce_sum_name = f"{attn_mask_basename}/ReduceSum"
@@ -4705,55 +4236,6 @@ def make_attention_mask_standard_reformatting_for_gqa(self, attn_mask_basename):
self.mask_attrs["seqlens_k"] = cast_1_name
self.mask_attrs["total_seq_len"] = cast_2_name
- def make_attention_mask_reformatting_for_gqa(self):
- # Make nodes for the attention mask subgraph that calculates
- # attributes about the 2D attention mask to use in GroupQueryAttention
- basename = "/model/attn_mask_reformat"
- attn_mask_basename = f"{basename}/attn_mask_subgraph"
-
- if self.extra_options.get("enable_webgpu_graph", False):
- self.make_attention_mask_graph_capture_reformatting_for_gqa(attn_mask_basename)
- else:
- self.make_attention_mask_standard_reformatting_for_gqa(attn_mask_basename)
-
- def make_attention_mask_reformatting_for_sparse_attn(self):
- # Make nodes for the attention mask subgraph that calculates
- # attributes about the 2D attention mask to use in SparseAttention
- #
- # attention_mask
- # / \
- # ReduceSum Shape
- # (keepdims=0) |
- # | |
- # Cast to int32 Gather
- # | |
- # key_total_seq_lens Cast to int32
- # (1D) |
- # total_seq_len
- # (int)
-
- basename = "/model/attn_mask_reformat"
- attn_mask_basename = f"{basename}/attn_mask_subgraph"
-
- # Left path
- reduce_sum_name = f"{attn_mask_basename}/ReduceSum"
- reduce_sum_inputs = [self.input_names["attention_mask"], "/model/constants/INT64/[1]"]
- self.make_reduce_sum(reduce_sum_name, reduce_sum_inputs, dtype=ir.DataType.INT64, shape=["batch_size"])
- cast_1_name = f"{attn_mask_basename}/ReduceSum/Cast"
- self.make_cast(cast_1_name, f"{reduce_sum_name}/output_0", dtype=ir.DataType.INT32, shape=["batch_size"])
-
- # Right path
- shape_name = f"{attn_mask_basename}/Shape"
- self.make_shape(shape_name, self.input_names["attention_mask"], shape=[2])
- gather_name = f"{attn_mask_basename}/Gather"
- gather_inputs = [f"{shape_name}/output_0", "/model/constants/INT64/1"]
- self.make_gather(gather_name, gather_inputs, dtype=ir.DataType.INT64, shape=[], axis=0)
- cast_2_name = f"{attn_mask_basename}/Gather/Cast"
- self.make_cast(cast_2_name, f"{gather_name}/output_0", dtype=ir.DataType.INT32, shape=None)
-
- self.mask_attrs["key_total_seq_lens"] = cast_1_name
- self.mask_attrs["total_seq_len"] = cast_2_name
-
def make_position_ids_reformatting(self):
# For most cases, position_ids are already properly formatted as 2D tensors
# with int64 values matching input_ids shape, so we can use them directly
diff --git a/src/python/py/models/builders/expansions/__init__.py b/src/python/py/models/builders/expansions/__init__.py
new file mode 100644
index 0000000000..8427bacd6f
--- /dev/null
+++ b/src/python/py/models/builders/expansions/__init__.py
@@ -0,0 +1,12 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+from .trt_rtx import TRT_RTX
+from .webgpu import WebGPU
+
+__all__ = [
+ "TRT_RTX",
+ "WebGPU",
+]
diff --git a/src/python/py/models/builders/expansions/trt_rtx.py b/src/python/py/models/builders/expansions/trt_rtx.py
new file mode 100644
index 0000000000..1aaedb748b
--- /dev/null
+++ b/src/python/py/models/builders/expansions/trt_rtx.py
@@ -0,0 +1,327 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+import onnx_ir as ir
+import torch
+
+
+class TRT_RTX:
+ """
+ TRT-RTX specific subgraph expansions
+ """
+ def make_layernorm_subgraph(self, name, **kwargs):
+ # This method can be used to create multiple LayerNorm operations
+ op_type = kwargs.pop("op_type")
+ inputs = kwargs.pop("inputs")
+ outputs = kwargs.pop("outputs")
+ skip = kwargs.pop("skip")
+ new_io_dtype = kwargs.pop("new_io_dtype")
+
+ if op_type == "LayerNormalization":
+ # Create LayerNorm op
+ self.make_layernorm_op(name, op_type, inputs, outputs, skip, new_io_dtype, **kwargs)
+
+ elif op_type == "SkipLayerNormalization":
+ # Create subgraph to calculate SkipLayerNorm
+ self.make_skip_layer_norm(
+ name,
+ root_input=inputs[0],
+ skip_input=inputs[1],
+ weight_name=inputs[2],
+ bias_name=inputs[3],
+ output_0=outputs[0],
+ output_3=outputs[3] if len(outputs) > 3 else None,
+ io_dtype=new_io_dtype,
+ shape=["batch_size", "sequence_length", self.hidden_size],
+ )
+
+ elif op_type == "SimplifiedLayerNormalization":
+ # Create subgraph to calculate RMSNorm
+ self.make_simplified_layer_norm(
+ name,
+ root_input=inputs[0],
+ weight_name=inputs[1],
+ output_0=outputs[0],
+ io_dtype=new_io_dtype,
+ shape=["batch_size", "sequence_length", self.hidden_size],
+ )
+
+ elif op_type == "SkipSimplifiedLayerNormalization":
+ # Create subgraph to calculate SkipRMSNorm
+ self.make_skip_simplified_layer_norm(
+ name,
+ root_input=inputs[0],
+ skip_input=inputs[1],
+ weight_name=inputs[2],
+ output_0=outputs[0],
+ output_3=outputs[3] if len(outputs) > 3 else None,
+ io_dtype=new_io_dtype,
+ shape=["batch_size", "sequence_length", self.hidden_size],
+ )
+
+ def make_skip_simplified_layer_norm(
+ self, basename, root_input, skip_input, weight_name, output_0, output_3, io_dtype, shape
+ ):
+ # root_input skip_input
+ # | |
+ # +------------------+
+ # |
+ # Add-------------> output (1)
+ # |
+ # SimplifiedLayerNorm----> output (0)
+ make_add_name = f"{basename}/Add"
+ output_3 = f"{make_add_name}/output_0" if output_3 is None else output_3
+ self.make_node("Add", inputs=[root_input, skip_input], outputs=[output_3], name=make_add_name)
+ self.make_value(output_3, io_dtype, shape=["batch_size", "sequence_length", self.hidden_size])
+
+ make_simplified_layer_norm_name = f"{basename}/skip_simplified_layer_norm"
+ self.make_simplified_layer_norm(
+ make_simplified_layer_norm_name, output_3, weight_name, output_0, io_dtype, shape=shape
+ )
+
+ def make_skip_layer_norm(
+ self, basename, root_input, skip_input, weight_name, bias_name, output_0, output_3, io_dtype, shape
+ ):
+ # root_input skip_input
+ # | |
+ # +------------------+
+ # |
+ # Add-------------> output (1)
+ # |
+ # LayerNormalization-----> output (0)
+ make_add_name = f"{basename}/Add"
+ output_3 = f"{make_add_name}/output_0" if output_3 is None else output_3
+ self.make_node("Add", inputs=[root_input, skip_input], outputs=[output_3], name=make_add_name)
+ self.make_value(output_3, io_dtype, shape=["batch_size", "sequence_length", self.hidden_size])
+
+ make_layer_norm_name = f"{basename}/LayerNormalization"
+ inputs = [output_3, weight_name, bias_name]
+
+ kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
+ kwargs.update({"axis": -1, "stash_type": 1})
+
+ self.make_node("LayerNormalization", inputs=inputs, outputs=[output_0], name=make_layer_norm_name, **kwargs)
+ self.make_value(output_0, io_dtype, shape=shape)
+
+ # This expansion contrib-op can be updated / deprecated in the future.
+ def make_simplified_layer_norm(self, basename, root_input, weight_name, output_0, io_dtype, shape):
+ # Cast (float32) - most calc happens in higher precision
+ # |
+ # +-------+-------+
+ # | |
+ # Pow |
+ # | |
+ # ReduceMean |
+ # | |
+ # Add |
+ # | |
+ # Sqrt |
+ # | |
+ # Div |
+ # | |
+ # +-------+-------+
+ # |
+ # Mul
+ # |
+ # Cast_1 (io_dtype - float16)
+ # |
+ # Mul_1
+
+ make_cast_name = f"{basename}/Cast"
+ self.make_cast(make_cast_name, root_input, ir.DataType.FLOAT, shape=shape)
+
+ make_pow_name = f"{basename}/Pow"
+ make_pow_inputs = [f"{make_cast_name}/output_0", "/model/constants/FLOAT/2"]
+
+ self.make_node(
+ "Pow", inputs=make_pow_inputs, outputs=[f"{make_pow_name}/output_0"], name=make_pow_name, domain=""
+ )
+ self.make_value(f"{make_pow_name}/output_0", ir.DataType.FLOAT, shape=shape)
+
+ make_reducemean_name = f"{basename}/ReduceMean"
+ make_reducemean_inputs = [f"{make_pow_name}/output_0", "/model/constants/INT64/[-1]"]
+ self.make_reduce_mean(
+ make_reducemean_name, make_reducemean_inputs, ir.DataType.FLOAT, keepdims=True, shape=shape
+ )
+
+ make_add_name = f"{basename}/Add"
+ make_add_inputs = [
+ f"{make_reducemean_name}/output_0",
+ f"/model/constants/FLOAT/{self.layernorm_attrs['epsilon']}",
+ ]
+ self.make_add(make_add_name, make_add_inputs, ir.DataType.FLOAT, shape=shape)
+
+ make_sqrt_name = f"{basename}/Sqrt"
+ make_sqrt_inputs = [f"{make_add_name}/output_0"]
+ self.make_sqrt(make_sqrt_name, make_sqrt_inputs, ir.DataType.FLOAT, shape=shape)
+
+ make_div_name = f"{basename}/Div"
+ make_div_inputs = ["/model/constants/FLOAT/1", f"{make_sqrt_name}/output_0"]
+ self.make_div(make_div_name, make_div_inputs, ir.DataType.FLOAT, shape=shape)
+
+ make_mul_name = f"{basename}/Mul"
+ make_mul_inputs = [f"{make_div_name}/output_0", f"{make_cast_name}/output_0"]
+ self.make_mul(make_mul_name, make_mul_inputs, ir.DataType.FLOAT, shape=shape)
+
+ make_cast_1_name = f"{basename}/Cast_1"
+ self.make_cast(make_cast_1_name, f"{make_mul_name}/output_0", dtype=io_dtype, shape=shape)
+
+ make_mul_1_name = f"{basename}/Mul_1"
+ make_mul_1_inputs = [f"{make_cast_1_name}/output_0", weight_name]
+
+ self.make_node("Mul", inputs=make_mul_1_inputs, outputs=[output_0], name=make_mul_1_name)
+ self.make_value(output_0, dtype=io_dtype, shape=shape)
+
+ def make_padded_cache(self, small_cache, large_cache, pad_value=0.0):
+ """Pad small cache to match large cache shape for uniform If node branches.
+
+ This is used for TRT-RTX EP which requires uniform dimensions in both branches of If nodes.
+
+ Args:
+ small_cache: The smaller cache tensor to pad
+ large_cache: The larger cache tensor (defines target shape)
+ pad_value: Value to use for padding (1.0 for cos_cache, 0.0 for sin_cache)
+ """
+ target_shape = large_cache.shape
+ if small_cache.shape == target_shape:
+ return small_cache
+
+ # Create padded tensor filled with pad_value
+ padded_cache = torch.full(target_shape, pad_value, dtype=small_cache.dtype)
+ # Copy original data to the beginning
+ padded_cache[: small_cache.shape[0], :] = small_cache
+ return padded_cache
+
+ def make_split_if_nodes(
+ self,
+ basename,
+ greater_name,
+ cos_cache_name,
+ sin_cache_name,
+ cos_cache_large,
+ sin_cache_large,
+ cos_cache_small,
+ sin_cache_small,
+ cos_cache_large_name,
+ sin_cache_large_name,
+ cos_cache_small_name,
+ sin_cache_small_name,
+ small_cache_shape,
+ ):
+ """Create split If nodes for TRT-RTX to workaround trt-rtx multi-output bug.
+
+ This is a TEMPORARY workaround for TRT-RTX bug where If nodes with
+ multiple outputs
+
+ Creates two separate If nodes instead of one:
+ - {basename}/cos/If: Outputs cos_cache only
+ - {basename}/sin/If: Outputs sin_cache only
+
+ Both If nodes use the same condition and independently select their respective caches.
+ """
+ cos_if_name = f"{basename}/cos/If"
+
+ cos_large_for_split = ir.node(
+ "Constant",
+ [],
+ outputs=[
+ ir.Value(
+ name=f"{cos_cache_large_name}_split",
+ type=ir.TensorType(self.io_dtype),
+ shape=ir.Shape(cos_cache_large.shape),
+ )
+ ],
+ name="/large/cos_cache/Constant_split_cos",
+ attributes=dict(value=ir.tensor(cos_cache_large)),
+ )
+
+ cos_small_for_split = ir.node(
+ "Constant",
+ [],
+ outputs=[
+ ir.Value(
+ name=f"{cos_cache_small_name}_split",
+ type=ir.TensorType(self.io_dtype),
+ shape=ir.Shape(small_cache_shape),
+ )
+ ],
+ name="/small/cos_cache/Constant_split_cos",
+ attributes=dict(value=ir.tensor(cos_cache_small)),
+ )
+
+ self.make_node(
+ "If",
+ inputs=[f"{greater_name}/output_0"],
+ outputs=[cos_cache_name],
+ name=cos_if_name,
+ then_branch=ir.Graph(
+ inputs=[],
+ outputs=[cos_large_for_split.outputs[0]],
+ nodes=[cos_large_for_split],
+ name="large_cos_cache_graph",
+ ),
+ else_branch=ir.Graph(
+ inputs=[],
+ outputs=[cos_small_for_split.outputs[0]],
+ nodes=[cos_small_for_split],
+ name="small_cos_cache_graph",
+ ),
+ )
+
+ # Create separate If node for sin_cache only
+ sin_if_name = f"{basename}/sin/If"
+
+ # Create unique constant nodes for sin to avoid tensor sharing
+ sin_large_for_split = ir.node(
+ "Constant",
+ [],
+ outputs=[
+ ir.Value(
+ name=f"{sin_cache_large_name}_split",
+ type=ir.TensorType(self.io_dtype),
+ shape=ir.Shape(sin_cache_large.shape),
+ )
+ ],
+ name="/large/sin_cache/Constant_split_sin",
+ attributes=dict(value=ir.tensor(sin_cache_large)),
+ )
+
+ sin_small_for_split = ir.node(
+ "Constant",
+ [],
+ outputs=[
+ ir.Value(
+ name=f"{sin_cache_small_name}_split",
+ type=ir.TensorType(self.io_dtype),
+ shape=ir.Shape(small_cache_shape),
+ )
+ ],
+ name="/small/sin_cache/Constant_split_sin",
+ attributes=dict(value=ir.tensor(sin_cache_small)),
+ )
+
+ self.make_node(
+ "If",
+ inputs=[f"{greater_name}/output_0"],
+ outputs=[sin_cache_name],
+ name=sin_if_name,
+ then_branch=ir.Graph(
+ inputs=[],
+ outputs=[sin_large_for_split.outputs[0]],
+ nodes=[sin_large_for_split],
+ name="large_sin_cache_graph",
+ ),
+ else_branch=ir.Graph(
+ inputs=[],
+ outputs=[sin_small_for_split.outputs[0]],
+ nodes=[sin_small_for_split],
+ name="small_sin_cache_graph",
+ ),
+ )
+
+ # Create output values
+ self.make_value(cos_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
+ self.make_value(sin_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
+
diff --git a/src/python/py/models/builders/expansions/webgpu.py b/src/python/py/models/builders/expansions/webgpu.py
new file mode 100644
index 0000000000..002dd9f59e
--- /dev/null
+++ b/src/python/py/models/builders/expansions/webgpu.py
@@ -0,0 +1,55 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+import onnx_ir as ir
+
+
+class WebGPU:
+ """
+ WebGPU specific subgraph expansions
+ """
+ def make_attention_mask_graph_capture_reformatting_for_gqa(self):
+ # Make nodes for the attention mask subgraph that calculates
+ # attributes about the 2D attention mask to use in GroupQueryAttention
+ #
+ # Key difference vs make_attention_mask_standard_reformatting_for_gqa:
+ # - Standard mode: total_seq_len is calculated from Shape op (always runs on CPU)
+ # - Graph capture mode: No Shape ops inserted to ensure all ops run on GPU (no CPU ops)
+ #
+ # attention_mask
+ # |
+ # Cast to int32
+ # |
+ # ReduceSum (keepdims=0)
+ # / \
+ # / \
+ # Sub ReduceMax
+ # | |
+ # seqlens_k total_seq_len
+ # (1D) (int)
+ basename = "/model/attn_mask_reformat"
+ attn_mask_basename = f"{basename}/attn_mask_subgraph"
+
+ # Calculate ReduceSum from attention_mask
+ cast_1_name = f"{attn_mask_basename}/Cast"
+ self.make_cast(
+ cast_1_name, self.input_names["attention_mask"], dtype=ir.DataType.INT32, shape=["batch_size", "total_sequence_length"]
+ )
+ reduce_sum_name = f"{attn_mask_basename}/ReduceSum"
+ reduce_sum_inputs = [f"{cast_1_name}/output_0", "/model/constants/INT64/[1]"]
+ self.make_reduce_sum(reduce_sum_name, reduce_sum_inputs, dtype=ir.DataType.INT32, shape=["batch_size"])
+
+ # Left branch: Calculate seqlens_k = ReduceSum - 1
+ sub_name = f"{attn_mask_basename}/Sub"
+ sub_inputs = [f"{reduce_sum_name}/output_0", "/model/constants/INT32/[1]"]
+ self.make_sub(sub_name, sub_inputs, dtype=ir.DataType.INT32, shape=["batch_size"])
+
+ # Right branch: ReduceMax to get maximum int value for total_seq_len
+ reduce_max_name = f"{attn_mask_basename}/ReduceMax"
+ reduce_max_inputs = [f"{reduce_sum_name}/output_0"]
+ self.make_reduce_max(reduce_max_name, reduce_max_inputs, dtype=ir.DataType.INT32, shape=[])
+
+ self.mask_attrs["seqlens_k"] = sub_name
+ self.mask_attrs["total_seq_len"] = reduce_max_name
diff --git a/src/python/py/models/builders/gemma.py b/src/python/py/models/builders/gemma.py
index d230d81b49..0845f7800b 100644
--- a/src/python/py/models/builders/gemma.py
+++ b/src/python/py/models/builders/gemma.py
@@ -130,10 +130,10 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
def is_local(self, layer_id):
return bool((layer_id + 1) % 6)
- def make_attention_init(self):
+ def make_attention_init(self, config):
self.attention_attrs["q_norm"] = True
self.attention_attrs["k_norm"] = True
- super().make_attention_init()
+ super().make_attention_init(config)
def make_rotary_embedding_multi_cache(self):
self.cos_cache_global_name, self.sin_cache_global_name = "cos_cache_global", "sin_cache_global"
diff --git a/src/python/py/models/builders/hunyuan.py b/src/python/py/models/builders/hunyuan.py
index c2f9a0c100..7bb3a79e06 100644
--- a/src/python/py/models/builders/hunyuan.py
+++ b/src/python/py/models/builders/hunyuan.py
@@ -48,10 +48,10 @@ def is_fused_rope_supported(self):
# Force explicit RotaryEmbedding nodes so QK norms can be placed after them.
return False
- def make_attention_init(self):
+ def make_attention_init(self, config):
self.attention_attrs["q_norm"] = True
self.attention_attrs["k_norm"] = True
- super().make_attention_init()
+ super().make_attention_init(config)
def make_attention_qk_rope_and_norm(self, layer_id, attention, **kwargs):
"""
diff --git a/src/python/py/models/builders/lfm2.py b/src/python/py/models/builders/lfm2.py
index 226850ae79..1399b09de1 100644
--- a/src/python/py/models/builders/lfm2.py
+++ b/src/python/py/models/builders/lfm2.py
@@ -37,10 +37,10 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
self.conv_L_cache = config.conv_L_cache
self.kv_layer_indices = [i for i, t in enumerate(self.layer_types) if t == "full_attention"]
- def make_attention_init(self):
+ def make_attention_init(self, config):
self.attention_attrs["q_norm"] = True
self.attention_attrs["k_norm"] = True
- super().make_attention_init()
+ super().make_attention_init(config)
def make_inputs_and_outputs(self):
# Replace the base class's all-layer KV lists with attention-layer-only lists,
diff --git a/src/python/py/models/builders/llama.py b/src/python/py/models/builders/llama.py
index f4055b0da1..5a757b0716 100644
--- a/src/python/py/models/builders/llama.py
+++ b/src/python/py/models/builders/llama.py
@@ -9,3 +9,18 @@
class LlamaModel(Model):
def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
+
+ def make_rope_init(self, config):
+ if "low_freq_factor" in config.rope_scaling:
+ # For models that rescale `inv_freq` using `low_freq_factor` and `high_freq_factor` (e.g. LLaMA-3.1)
+ factor = config.rope_scaling["factor"] if "factor" in config.rope_scaling else 0
+ low_freq_factor = config.rope_scaling["low_freq_factor"] if "low_freq_factor" in config.rope_scaling else 0
+ high_freq_factor = (
+ config.rope_scaling["high_freq_factor"] if "high_freq_factor" in config.rope_scaling else 0
+ )
+
+ self.rope_attrs["rescale_inv_freq"] = {
+ "factor": factor, # Scale factor when calculating `new_freq` in rotary embeddings
+ "low_freq_factor": low_freq_factor, # Low freq factor when calculating `low_freq_wavelen` in rotary embeddings
+ "high_freq_factor": high_freq_factor, # High freq factor when calculating `high_freq_wavelen` in rotary embeddings
+ }
diff --git a/src/python/py/models/builders/phi.py b/src/python/py/models/builders/phi.py
index b346ef8f8d..520f759736 100644
--- a/src/python/py/models/builders/phi.py
+++ b/src/python/py/models/builders/phi.py
@@ -70,6 +70,25 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
# position_ids won't be used since rotary embeddings are handled in GQA
self.position_ids_name = None
+ def make_rope_init(self, config):
+ if "short_factor" in config.rope_scaling:
+ # For models with multiple rotary embedding caches (e.g. Phi-3 mini 128K)
+ self.rope_attrs["mscale_policy"] = config.rope_scaling["type"]
+ short_factor = torch.tensor(config.rope_scaling["short_factor"], dtype=torch.float32)
+ long_factor = torch.tensor(config.rope_scaling["long_factor"], dtype=torch.float32)
+
+ short_mscale = config.rope_scaling["short_mscale"] if "short_mscale" in config.rope_scaling else 0
+ long_mscale = config.rope_scaling["long_mscale"] if "long_mscale" in config.rope_scaling else 0
+ short_mscale = short_mscale if short_mscale > 0 else self.make_mscale(self.context_length / self.original_context_length)
+ long_mscale = long_mscale if long_mscale > 0 else self.make_mscale(self.context_length / self.original_context_length)
+
+ self.rope_attrs["multi_cache"] = {
+ "short_factor": short_factor, # Short factor when calculating `inv_freq` in rotary embeddings
+ "long_factor": long_factor, # Long factor when calculating `inv_freq` in rotary embeddings
+ "short_mscale": short_mscale, # Magnitude scaling for short factor when scaling `emb.cos()/emb.sin()` in rotary embeddings
+ "long_mscale": long_mscale, # Magnitude scaling for long factor when scaling `emb.cos()/emb.sin()` in rotary embeddings
+ }
+
def make_position_ids_reformatting(self):
if self.ep not in self.eps_without_if_support:
position_ids_input_to_rotemb = super().make_position_ids_reformatting()
@@ -167,6 +186,74 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
self.clamp_limit = config.gegelu_limit
+ def make_attention_init(self, config):
+ # Block-sparse attention-specific variables
+ sparse_block_size = config.blocksparse_block_size if hasattr(config, "blocksparse_block_size") else 0
+ kernel_block_size = config.blocksparse_triton_kernel_block_size if hasattr(config, "blocksparse_triton_kernel_block_size") else 0
+ local_blocks = config.blocksparse_num_local_blocks if hasattr(config, "blocksparse_num_local_blocks") else 0
+ vert_block_stride = config.blocksparse_vert_stride if hasattr(config, "blocksparse_vert_stride") else 0
+ homo_head = config.blocksparse_homo_head_pattern if hasattr(config, "blocksparse_homo_head_pattern") else False
+
+ self.attention_attrs["block_sparse"] = {
+ "sparse_block_size": sparse_block_size, # Sparse block size for SparseAttention op
+ "kernel_block_size": kernel_block_size, # Kernel block size for sparse attention
+ "local_blocks": local_blocks, # Number of local blocks for sparse attention
+ "vert_stride": vert_block_stride, # Vertical stride to use for sparse attention
+ "homo_head": homo_head, # Use homo head pattern for sparse attention
+ }
+
+ super().make_attention_init(config)
+
+ def make_lm_head_init(self, config):
+ if hasattr(config, "dummy_token_indices"):
+ # Create LM head mask for tokens in the vocabulary
+ dummy_tokens_mask = torch.zeros(self.vocab_size).bool()
+ dummy_tokens_mask[config.dummy_token_indices] = True
+ self.lm_head_attrs["mask"] = dummy_tokens_mask
+
+ def make_attention_mask_reformatting(self):
+ super().make_attention_mask_reformatting()
+ if self.attention_attrs["block_sparse"]["sparse_block_size"] != 0:
+ self.make_attention_mask_reformatting_for_sparse_attn()
+
+ def make_attention_mask_reformatting_for_sparse_attn(self):
+ # Make nodes for the attention mask subgraph that calculates
+ # attributes about the 2D attention mask to use in SparseAttention
+ #
+ # attention_mask
+ # / \
+ # ReduceSum Shape
+ # (keepdims=0) |
+ # | |
+ # Cast to int32 Gather
+ # | |
+ # key_total_seq_lens Cast to int32
+ # (1D) |
+ # total_seq_len
+ # (int)
+
+ basename = "/model/attn_mask_reformat"
+ attn_mask_basename = f"{basename}/attn_mask_subgraph"
+
+ # Left path
+ reduce_sum_name = f"{attn_mask_basename}/ReduceSum"
+ reduce_sum_inputs = [self.input_names["attention_mask"], "/model/constants/INT64/[1]"]
+ self.make_reduce_sum(reduce_sum_name, reduce_sum_inputs, dtype=ir.DataType.INT64, shape=["batch_size"])
+ cast_1_name = f"{attn_mask_basename}/ReduceSum/Cast"
+ self.make_cast(cast_1_name, f"{reduce_sum_name}/output_0", dtype=ir.DataType.INT32, shape=["batch_size"])
+
+ # Right path
+ shape_name = f"{attn_mask_basename}/Shape"
+ self.make_shape(shape_name, self.input_names["attention_mask"], shape=[2])
+ gather_name = f"{attn_mask_basename}/Gather"
+ gather_inputs = [f"{shape_name}/output_0", "/model/constants/INT64/1"]
+ self.make_gather(gather_name, gather_inputs, dtype=ir.DataType.INT64, shape=[], axis=0)
+ cast_2_name = f"{attn_mask_basename}/Gather/Cast"
+ self.make_cast(cast_2_name, f"{gather_name}/output_0", dtype=ir.DataType.INT32, shape=None)
+
+ self.mask_attrs["key_total_seq_lens"] = cast_1_name
+ self.mask_attrs["total_seq_len"] = cast_2_name
+
def calculate_cdiv(self, a, b):
return -(a // -b)
diff --git a/src/python/py/models/builders/qwen.py b/src/python/py/models/builders/qwen.py
index b5974d63bb..9e31b2fcb5 100644
--- a/src/python/py/models/builders/qwen.py
+++ b/src/python/py/models/builders/qwen.py
@@ -29,10 +29,10 @@ class Qwen3Model(QwenModel):
def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
- def make_attention_init(self):
+ def make_attention_init(self, config):
self.attention_attrs["q_norm"] = True
self.attention_attrs["k_norm"] = True
- super().make_attention_init()
+ super().make_attention_init(config)
class Qwen25VLTextModel(Model):
@@ -945,6 +945,8 @@ def load_weights(self, input_path):
token=self.hf_token,
**extra_kwargs,
)
+
+
class Qwen35TextModel(Model):
"""Qwen3.5 hybrid model builder.
@@ -2117,7 +2119,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
# MoE layers use MoE/QMoE ops instead of individual MatMul nodes,
# so remove any /mlp/ MatMul overrides that don't apply.
- algo_config = self.quant_attrs["int4"].get("algo_config")
+ algo_config = self.quant_attrs.get("algo_config")
if algo_config is not None and hasattr(algo_config, "customized_weight_config"):
keys_to_remove = [k for k in algo_config.customized_weight_config if "/mlp/" in k]
for k in keys_to_remove:
@@ -2256,7 +2258,7 @@ def make_shared_expert(self, layer_id, shared_expert, shared_expert_gate, root_i
self.make_sigmoid(gate_sigmoid_name, f"{gate_matmul_name}/output_0", self.io_dtype,
shape=["batch_size", "sequence_length", 1])
- gated_mul_name = f"{basename}/GatedMul"
+ gated_mul_name = f"{basename}/Mul"
self.make_mul(gated_mul_name,
[f"{down_matmul}/output_0", f"{gate_sigmoid_name}/output_0"],
dtype=self.io_dtype,
diff --git a/src/python/py/models/quantized_model.py b/src/python/py/models/quantized_model.py
index 0ed952cdbf..3f5a5753c5 100644
--- a/src/python/py/models/quantized_model.py
+++ b/src/python/py/models/quantized_model.py
@@ -228,7 +228,7 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme
self.layers = {}
self.num_layers = num_layers
self._quant_attrs = quant_attrs
- self._load_quant_config(quant_attrs) # codeql[py/init-calls-subclass]
+ self._load_quant_config(quant_attrs)
lm_head_tensors = {}
for weight_file in os.listdir(input_path):
@@ -242,8 +242,8 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme
continue
# Per-layer quantization support
- local_bits = self.get_layer_bits(name) # codeql[py/init-calls-subclass]
- local_group_size = self.get_layer_group_size(name) # codeql[py/init-calls-subclass]
+ local_bits = self.get_layer_bits(name)
+ local_group_size = self.get_layer_group_size(name)
if name == "model.embed_tokens.weight" or name == "transformer.embedding.word_embeddings.weight":
self.embedding.weight = tensor
@@ -698,7 +698,7 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme
self.set_properties()
# Canonical name mapping for lm_head tensors (transformer.output_layer.* -> lm_head.*)
- _LM_HEAD_NAME_MAP = { # noqa: RUF012
+ _LM_HEAD_NAME_MAP = {
"transformer.output_layer.weight": "lm_head.weight",
"transformer.output_layer.bias": "lm_head.bias",
"transformer.output_layer.qweight": "lm_head.qweight",
diff --git a/src/python/setup.py.in b/src/python/setup.py.in
index fd25e15891..f06c9ecd3a 100644
--- a/src/python/setup.py.in
+++ b/src/python/setup.py.in
@@ -56,7 +56,12 @@ setup(
description='ONNX Runtime GenAI',
long_description=long_description,
long_description_content_type='text/markdown',
- packages=['onnxruntime_genai', 'onnxruntime_genai.models', 'onnxruntime_genai.models.builders'],
+ packages=[
+ 'onnxruntime_genai',
+ 'onnxruntime_genai.models',
+ 'onnxruntime_genai.models.builders',
+ 'onnxruntime_genai.models.builders.expansions',
+ ],
include_package_data=True,
package_data={'': ['*.pyd', '*.dll', '*.so*', '*.dylib'] + extras},
install_requires=_get_install_requires(),
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index ce64136033..f9091f28e3 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -54,7 +54,7 @@ if(USE_CUDA AND CMAKE_CUDA_COMPILER AND ENABLE_CUDA_KERNEL_TESTS)
add_dependencies(unit_tests onnxruntime-genai-cuda)
endif()
-set(TEST_MODEL_SRC_DIR "${CMAKE_CURRENT_SOURCE_DIR}/test_models/")
+set(TEST_MODEL_SRC_DIR "${CMAKE_CURRENT_SOURCE_DIR}/models/")
add_compile_definitions(MODEL_PATH="${TEST_MODEL_SRC_DIR}")
set_target_properties(unit_tests PROPERTIES FOLDER "Tests")
diff --git a/test/test_models/audios/1272-141231-0002.mp3 b/test/audios/1272-141231-0002.mp3
similarity index 100%
rename from test/test_models/audios/1272-141231-0002.mp3
rename to test/audios/1272-141231-0002.mp3
diff --git a/test/audios/common_voice_en_59751.mp3 b/test/audios/common_voice_en_59751.mp3
new file mode 100644
index 0000000000..da9a6f1179
Binary files /dev/null and b/test/audios/common_voice_en_59751.mp3 differ
diff --git a/test/test_models/audios/jfk.flac b/test/audios/jfk.flac
similarity index 100%
rename from test/test_models/audios/jfk.flac
rename to test/audios/jfk.flac
diff --git a/test/test_models/audios/tedlium_long_120s.flac b/test/audios/tedlium_long_120s.flac
similarity index 100%
rename from test/test_models/audios/tedlium_long_120s.flac
rename to test/audios/tedlium_long_120s.flac
diff --git a/test/csharp/Microsoft.ML.OnnxRuntimeGenAI.Tests.csproj b/test/csharp/Microsoft.ML.OnnxRuntimeGenAI.Tests.csproj
index 6cf90b0245..5d117478d0 100644
--- a/test/csharp/Microsoft.ML.OnnxRuntimeGenAI.Tests.csproj
+++ b/test/csharp/Microsoft.ML.OnnxRuntimeGenAI.Tests.csproj
@@ -80,10 +80,10 @@
false
-
+
PreserveNewest
false
- "test_models\"
+ "models\"
diff --git a/test/csharp/TestOnnxRuntimeGenAIAPI.cs b/test/csharp/TestOnnxRuntimeGenAIAPI.cs
index 932e3a425f..260393b420 100644
--- a/test/csharp/TestOnnxRuntimeGenAIAPI.cs
+++ b/test/csharp/TestOnnxRuntimeGenAIAPI.cs
@@ -46,9 +46,9 @@ private static string GetDirectoryInTreeThatContains(string currentDirectory, st
private static Lazy _lazyPhi2Path = new Lazy(() =>
{
string cpuModelPath = Path.Combine(GetDirectoryInTreeThatContains(Directory.GetCurrentDirectory(), "test"),
- "test_models", "phi-2", "int4", "cpu");
+ "models", "phi-2", "int4", "cpu");
string cudaModelPath = Path.Combine(GetDirectoryInTreeThatContains(Directory.GetCurrentDirectory(), "test"),
- "test_models", "phi-2", "int4", "cuda");
+ "models", "phi-2", "int4", "cuda");
// Prefer CUDA model if available.
if (System.IO.Directory.Exists(cudaModelPath))
{
@@ -65,7 +65,7 @@ private static string GetDirectoryInTreeThatContains(string currentDirectory, st
private static Lazy _lazyTinyRandomGpt2ModelPath = new Lazy(() =>
{
string modelPath = Path.Combine(GetDirectoryInTreeThatContains(Directory.GetCurrentDirectory(), "test"),
- "test_models", "hf-internal-testing", "tiny-random-gpt2-fp32");
+ "models", "hf-internal-testing", "tiny-random-gpt2-fp32");
if (System.IO.Directory.Exists(modelPath))
{
return modelPath;
@@ -79,7 +79,7 @@ private static string GetDirectoryInTreeThatContains(string currentDirectory, st
private static Lazy _lazyAdaptersPath = new Lazy(() =>
{
string modelPath = Path.Combine(GetDirectoryInTreeThatContains(Directory.GetCurrentDirectory(), "test"),
- "test_models", "adapters");
+ "models", "adapters");
if (System.IO.Directory.Exists(modelPath))
{
return modelPath;
@@ -238,7 +238,7 @@ public void TestLoadModelFromMemory()
public void TestAudioOpenBytes()
{
byte[] audioBytes = File.ReadAllBytes(Path.Combine(GetDirectoryInTreeThatContains(Directory.GetCurrentDirectory(), "test"),
- "test_models", "audios", "1272-141231-0002.mp3"));
+ "audios", "1272-141231-0002.mp3"));
var audios = Audios.Load(audioBytes);
Assert.NotNull(audios);
}
@@ -247,7 +247,7 @@ public void TestAudioOpenBytes()
public void TestImageOpenBytes()
{
byte[] imageBytes = File.ReadAllBytes(Path.Combine(GetDirectoryInTreeThatContains(Directory.GetCurrentDirectory(), "test"),
- "test_models", "images", "10809054.jpg"));
+ "images", "10809054.jpg"));
var images = Images.Load(imageBytes);
Assert.NotNull(images);
}
diff --git a/test/cuda_kernel/cuda_sampling_benchmark.cpp b/test/cuda/cuda_sampling_benchmark.cpp
similarity index 100%
rename from test/cuda_kernel/cuda_sampling_benchmark.cpp
rename to test/cuda/cuda_sampling_benchmark.cpp
diff --git a/test/cuda_kernel/cuda_sampling_tests.cpp b/test/cuda/cuda_sampling_tests.cpp
similarity index 100%
rename from test/cuda_kernel/cuda_sampling_tests.cpp
rename to test/cuda/cuda_sampling_tests.cpp
diff --git a/test/cuda_kernel/cuda_topk_benchmark.cpp b/test/cuda/cuda_topk_benchmark.cpp
similarity index 100%
rename from test/cuda_kernel/cuda_topk_benchmark.cpp
rename to test/cuda/cuda_topk_benchmark.cpp
diff --git a/test/cuda_kernel/cuda_topk_tests.cpp b/test/cuda/cuda_topk_tests.cpp
similarity index 100%
rename from test/cuda_kernel/cuda_topk_tests.cpp
rename to test/cuda/cuda_topk_tests.cpp
diff --git a/test/test_models/images/10809054.jpg b/test/images/10809054.jpg
similarity index 100%
rename from test/test_models/images/10809054.jpg
rename to test/images/10809054.jpg
diff --git a/test/test_models/images/australia.jpg b/test/images/australia.jpg
similarity index 100%
rename from test/test_models/images/australia.jpg
rename to test/images/australia.jpg
diff --git a/test/images/cars.jpg b/test/images/cars.jpg
new file mode 100644
index 0000000000..7a5d05274d
Binary files /dev/null and b/test/images/cars.jpg differ
diff --git a/test/test_models/images/landscape.jpg b/test/images/landscape.jpg
similarity index 100%
rename from test/test_models/images/landscape.jpg
rename to test/images/landscape.jpg
diff --git a/test/test_models/images/sheet.png b/test/images/sheet.png
similarity index 100%
rename from test/test_models/images/sheet.png
rename to test/images/sheet.png
diff --git a/test/test_models/gemma4-vision-preprocessing/audio_feature_extraction.json b/test/models/gemma4/audio_feature_extraction.json
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/audio_feature_extraction.json
rename to test/models/gemma4/audio_feature_extraction.json
diff --git a/test/test_models/gemma4-vision-preprocessing/dummy_embedding.onnx b/test/models/gemma4/dummy_embedding.onnx
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/dummy_embedding.onnx
rename to test/models/gemma4/dummy_embedding.onnx
diff --git a/test/test_models/gemma4-vision-preprocessing/dummy_speech.onnx b/test/models/gemma4/dummy_speech.onnx
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/dummy_speech.onnx
rename to test/models/gemma4/dummy_speech.onnx
diff --git a/test/test_models/gemma4-vision-preprocessing/dummy_text.onnx b/test/models/gemma4/dummy_text.onnx
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/dummy_text.onnx
rename to test/models/gemma4/dummy_text.onnx
diff --git a/test/test_models/gemma4-vision-preprocessing/dummy_vision.onnx b/test/models/gemma4/dummy_vision.onnx
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/dummy_vision.onnx
rename to test/models/gemma4/dummy_vision.onnx
diff --git a/test/test_models/gemma4-vision-preprocessing/genai_config.json b/test/models/gemma4/genai_config.json
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/genai_config.json
rename to test/models/gemma4/genai_config.json
diff --git a/test/test_models/gemma4-vision-preprocessing/processor_config.json b/test/models/gemma4/processor_config.json
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/processor_config.json
rename to test/models/gemma4/processor_config.json
diff --git a/test/test_models/gemma4-vision-preprocessing/special_tokens_map.json b/test/models/gemma4/special_tokens_map.json
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/special_tokens_map.json
rename to test/models/gemma4/special_tokens_map.json
diff --git a/test/test_models/gemma4-vision-preprocessing/tokenizer.json b/test/models/gemma4/tokenizer.json
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/tokenizer.json
rename to test/models/gemma4/tokenizer.json
diff --git a/test/test_models/gemma4-vision-preprocessing/tokenizer_config.json b/test/models/gemma4/tokenizer_config.json
similarity index 100%
rename from test/test_models/gemma4-vision-preprocessing/tokenizer_config.json
rename to test/models/gemma4/tokenizer_config.json
diff --git a/test/test_models/hf-internal-testing/tiny-qwen35-cuda/decoder.onnx b/test/models/hf-internal-testing/tiny-qwen35-cuda/decoder.onnx
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-qwen35-cuda/decoder.onnx
rename to test/models/hf-internal-testing/tiny-qwen35-cuda/decoder.onnx
diff --git a/test/test_models/hf-internal-testing/tiny-qwen35-cuda/genai_config.json b/test/models/hf-internal-testing/tiny-qwen35-cuda/genai_config.json
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-qwen35-cuda/genai_config.json
rename to test/models/hf-internal-testing/tiny-qwen35-cuda/genai_config.json
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp16-cuda/genai_config.json b/test/models/hf-internal-testing/tiny-random-gpt2-fp16-cuda/genai_config.json
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp16-cuda/genai_config.json
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp16-cuda/genai_config.json
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp16-cuda/past.onnx b/test/models/hf-internal-testing/tiny-random-gpt2-fp16-cuda/past.onnx
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp16-cuda/past.onnx
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp16-cuda/past.onnx
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp32-cuda/genai_config.json b/test/models/hf-internal-testing/tiny-random-gpt2-fp32-cuda/genai_config.json
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp32-cuda/genai_config.json
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp32-cuda/genai_config.json
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp32-cuda/past.onnx b/test/models/hf-internal-testing/tiny-random-gpt2-fp32-cuda/past.onnx
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp32-cuda/past.onnx
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp32-cuda/past.onnx
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/genai_config.json b/test/models/hf-internal-testing/tiny-random-gpt2-fp32/genai_config.json
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/genai_config.json
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp32/genai_config.json
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/past.onnx b/test/models/hf-internal-testing/tiny-random-gpt2-fp32/past.onnx
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/past.onnx
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp32/past.onnx
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/tokenizer.json b/test/models/hf-internal-testing/tiny-random-gpt2-fp32/tokenizer.json
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/tokenizer.json
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp32/tokenizer.json
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/tokenizer_config.json b/test/models/hf-internal-testing/tiny-random-gpt2-fp32/tokenizer_config.json
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/tokenizer_config.json
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp32/tokenizer_config.json
diff --git a/test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/vocab.json b/test/models/hf-internal-testing/tiny-random-gpt2-fp32/vocab.json
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-gpt2-fp32/vocab.json
rename to test/models/hf-internal-testing/tiny-random-gpt2-fp32/vocab.json
diff --git a/test/test_models/hf-internal-testing/tiny-random-lfm2-fp32/decoder.onnx b/test/models/hf-internal-testing/tiny-random-lfm2-fp32/decoder.onnx
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-lfm2-fp32/decoder.onnx
rename to test/models/hf-internal-testing/tiny-random-lfm2-fp32/decoder.onnx
diff --git a/test/test_models/hf-internal-testing/tiny-random-lfm2-fp32/genai_config.json b/test/models/hf-internal-testing/tiny-random-lfm2-fp32/genai_config.json
similarity index 100%
rename from test/test_models/hf-internal-testing/tiny-random-lfm2-fp32/genai_config.json
rename to test/models/hf-internal-testing/tiny-random-lfm2-fp32/genai_config.json
diff --git a/test/test_models/mistral3-vision-preprocessing/processor_config.json b/test/models/mistral3/processor_config.json
similarity index 100%
rename from test/test_models/mistral3-vision-preprocessing/processor_config.json
rename to test/models/mistral3/processor_config.json
diff --git a/test/test_models/multimodal-decoder-no-input-ids/dummy_embedding.onnx b/test/models/multimodal-decoder-no-input-ids/dummy_embedding.onnx
similarity index 100%
rename from test/test_models/multimodal-decoder-no-input-ids/dummy_embedding.onnx
rename to test/models/multimodal-decoder-no-input-ids/dummy_embedding.onnx
diff --git a/test/test_models/multimodal-decoder-no-input-ids/dummy_text.onnx b/test/models/multimodal-decoder-no-input-ids/dummy_text.onnx
similarity index 100%
rename from test/test_models/multimodal-decoder-no-input-ids/dummy_text.onnx
rename to test/models/multimodal-decoder-no-input-ids/dummy_text.onnx
diff --git a/test/test_models/multimodal-decoder-no-input-ids/dummy_vision.onnx b/test/models/multimodal-decoder-no-input-ids/dummy_vision.onnx
similarity index 100%
rename from test/test_models/multimodal-decoder-no-input-ids/dummy_vision.onnx
rename to test/models/multimodal-decoder-no-input-ids/dummy_vision.onnx
diff --git a/test/test_models/multimodal-decoder-no-input-ids/genai_config.json b/test/models/multimodal-decoder-no-input-ids/genai_config.json
similarity index 100%
rename from test/test_models/multimodal-decoder-no-input-ids/genai_config.json
rename to test/models/multimodal-decoder-no-input-ids/genai_config.json
diff --git a/test/test_models/multimodal-decoder-no-input-ids/tokenizer.json b/test/models/multimodal-decoder-no-input-ids/tokenizer.json
similarity index 100%
rename from test/test_models/multimodal-decoder-no-input-ids/tokenizer.json
rename to test/models/multimodal-decoder-no-input-ids/tokenizer.json
diff --git a/test/test_models/multimodal-decoder-no-input-ids/tokenizer_config.json b/test/models/multimodal-decoder-no-input-ids/tokenizer_config.json
similarity index 100%
rename from test/test_models/multimodal-decoder-no-input-ids/tokenizer_config.json
rename to test/models/multimodal-decoder-no-input-ids/tokenizer_config.json
diff --git a/test/test_models/multimodal-decoder-with-input-ids/dummy_embedding.onnx b/test/models/multimodal-decoder-with-input-ids/dummy_embedding.onnx
similarity index 100%
rename from test/test_models/multimodal-decoder-with-input-ids/dummy_embedding.onnx
rename to test/models/multimodal-decoder-with-input-ids/dummy_embedding.onnx
diff --git a/test/test_models/multimodal-decoder-with-input-ids/dummy_text.onnx b/test/models/multimodal-decoder-with-input-ids/dummy_text.onnx
similarity index 100%
rename from test/test_models/multimodal-decoder-with-input-ids/dummy_text.onnx
rename to test/models/multimodal-decoder-with-input-ids/dummy_text.onnx
diff --git a/test/test_models/multimodal-decoder-with-input-ids/dummy_vision.onnx b/test/models/multimodal-decoder-with-input-ids/dummy_vision.onnx
similarity index 100%
rename from test/test_models/multimodal-decoder-with-input-ids/dummy_vision.onnx
rename to test/models/multimodal-decoder-with-input-ids/dummy_vision.onnx
diff --git a/test/test_models/multimodal-decoder-with-input-ids/genai_config.json b/test/models/multimodal-decoder-with-input-ids/genai_config.json
similarity index 100%
rename from test/test_models/multimodal-decoder-with-input-ids/genai_config.json
rename to test/models/multimodal-decoder-with-input-ids/genai_config.json
diff --git a/test/test_models/multimodal-decoder-with-input-ids/tokenizer.json b/test/models/multimodal-decoder-with-input-ids/tokenizer.json
similarity index 100%
rename from test/test_models/multimodal-decoder-with-input-ids/tokenizer.json
rename to test/models/multimodal-decoder-with-input-ids/tokenizer.json
diff --git a/test/test_models/multimodal-decoder-with-input-ids/tokenizer_config.json b/test/models/multimodal-decoder-with-input-ids/tokenizer_config.json
similarity index 100%
rename from test/test_models/multimodal-decoder-with-input-ids/tokenizer_config.json
rename to test/models/multimodal-decoder-with-input-ids/tokenizer_config.json
diff --git a/test/test_models/vision-preprocessing/dummy_embedding.onnx b/test/models/phi3-v/dummy_embedding.onnx
similarity index 100%
rename from test/test_models/vision-preprocessing/dummy_embedding.onnx
rename to test/models/phi3-v/dummy_embedding.onnx
diff --git a/test/test_models/vision-preprocessing/dummy_text.onnx b/test/models/phi3-v/dummy_text.onnx
similarity index 100%
rename from test/test_models/vision-preprocessing/dummy_text.onnx
rename to test/models/phi3-v/dummy_text.onnx
diff --git a/test/test_models/vision-preprocessing/dummy_vision.onnx b/test/models/phi3-v/dummy_vision.onnx
similarity index 100%
rename from test/test_models/vision-preprocessing/dummy_vision.onnx
rename to test/models/phi3-v/dummy_vision.onnx
diff --git a/test/test_models/vision-preprocessing/genai_config.json b/test/models/phi3-v/genai_config.json
similarity index 100%
rename from test/test_models/vision-preprocessing/genai_config.json
rename to test/models/phi3-v/genai_config.json
diff --git a/test/test_models/vision-preprocessing/processor_config.json b/test/models/phi3-v/processor_config.json
similarity index 100%
rename from test/test_models/vision-preprocessing/processor_config.json
rename to test/models/phi3-v/processor_config.json
diff --git a/test/test_models/vision-preprocessing/special_tokens_map.json b/test/models/phi3-v/special_tokens_map.json
similarity index 100%
rename from test/test_models/vision-preprocessing/special_tokens_map.json
rename to test/models/phi3-v/special_tokens_map.json
diff --git a/test/test_models/vision-preprocessing/tokenizer.json b/test/models/phi3-v/tokenizer.json
similarity index 100%
rename from test/test_models/vision-preprocessing/tokenizer.json
rename to test/models/phi3-v/tokenizer.json
diff --git a/test/test_models/vision-preprocessing/tokenizer_config.json b/test/models/phi3-v/tokenizer_config.json
similarity index 100%
rename from test/test_models/vision-preprocessing/tokenizer_config.json
rename to test/models/phi3-v/tokenizer_config.json
diff --git a/test/test_models/pipeline-model/added_tokens.json b/test/models/pipeline-model/added_tokens.json
similarity index 100%
rename from test/test_models/pipeline-model/added_tokens.json
rename to test/models/pipeline-model/added_tokens.json
diff --git a/test/test_models/pipeline-model/genai_config.json b/test/models/pipeline-model/genai_config.json
similarity index 100%
rename from test/test_models/pipeline-model/genai_config.json
rename to test/models/pipeline-model/genai_config.json
diff --git a/test/test_models/pipeline-model/special_tokens_map.json b/test/models/pipeline-model/special_tokens_map.json
similarity index 100%
rename from test/test_models/pipeline-model/special_tokens_map.json
rename to test/models/pipeline-model/special_tokens_map.json
diff --git a/test/test_models/pipeline-model/tokenizer.json b/test/models/pipeline-model/tokenizer.json
similarity index 100%
rename from test/test_models/pipeline-model/tokenizer.json
rename to test/models/pipeline-model/tokenizer.json
diff --git a/test/test_models/pipeline-model/tokenizer_config.json b/test/models/pipeline-model/tokenizer_config.json
similarity index 100%
rename from test/test_models/pipeline-model/tokenizer_config.json
rename to test/models/pipeline-model/tokenizer_config.json
diff --git a/test/test_models/pipeline-model/vocab.json b/test/models/pipeline-model/vocab.json
similarity index 100%
rename from test/test_models/pipeline-model/vocab.json
rename to test/models/pipeline-model/vocab.json
diff --git a/test/test_models/qwen-vision-preprocessing/dummy_embedding.onnx b/test/models/qwen2-5-vl/dummy_embedding.onnx
similarity index 100%
rename from test/test_models/qwen-vision-preprocessing/dummy_embedding.onnx
rename to test/models/qwen2-5-vl/dummy_embedding.onnx
diff --git a/test/test_models/qwen-vision-preprocessing/dummy_text.onnx b/test/models/qwen2-5-vl/dummy_text.onnx
similarity index 100%
rename from test/test_models/qwen-vision-preprocessing/dummy_text.onnx
rename to test/models/qwen2-5-vl/dummy_text.onnx
diff --git a/test/test_models/qwen-vision-preprocessing/dummy_vision.onnx b/test/models/qwen2-5-vl/dummy_vision.onnx
similarity index 100%
rename from test/test_models/qwen-vision-preprocessing/dummy_vision.onnx
rename to test/models/qwen2-5-vl/dummy_vision.onnx
diff --git a/test/test_models/qwen-vision-preprocessing/genai_config.json b/test/models/qwen2-5-vl/genai_config.json
similarity index 100%
rename from test/test_models/qwen-vision-preprocessing/genai_config.json
rename to test/models/qwen2-5-vl/genai_config.json
diff --git a/test/test_models/qwen-vision-preprocessing/processor_config.json b/test/models/qwen2-5-vl/processor_config.json
similarity index 100%
rename from test/test_models/qwen-vision-preprocessing/processor_config.json
rename to test/models/qwen2-5-vl/processor_config.json
diff --git a/test/test_models/qwen-vision-preprocessing/special_tokens_map.json b/test/models/qwen2-5-vl/special_tokens_map.json
similarity index 100%
rename from test/test_models/qwen-vision-preprocessing/special_tokens_map.json
rename to test/models/qwen2-5-vl/special_tokens_map.json
diff --git a/test/test_models/qwen-vision-preprocessing/tokenizer.json b/test/models/qwen2-5-vl/tokenizer.json
similarity index 100%
rename from test/test_models/qwen-vision-preprocessing/tokenizer.json
rename to test/models/qwen2-5-vl/tokenizer.json
diff --git a/test/test_models/qwen-vision-preprocessing/tokenizer_config.json b/test/models/qwen2-5-vl/tokenizer_config.json
similarity index 100%
rename from test/test_models/qwen-vision-preprocessing/tokenizer_config.json
rename to test/models/qwen2-5-vl/tokenizer_config.json
diff --git a/test/test_models/qwen35-hybrid-preprocessing/dummy_embedding.onnx b/test/models/qwen3-5/dummy_embedding.onnx
similarity index 100%
rename from test/test_models/qwen35-hybrid-preprocessing/dummy_embedding.onnx
rename to test/models/qwen3-5/dummy_embedding.onnx
diff --git a/test/test_models/qwen35-hybrid-preprocessing/dummy_text.onnx b/test/models/qwen3-5/dummy_text.onnx
similarity index 100%
rename from test/test_models/qwen35-hybrid-preprocessing/dummy_text.onnx
rename to test/models/qwen3-5/dummy_text.onnx
diff --git a/test/test_models/qwen35-hybrid-preprocessing/dummy_vision.onnx b/test/models/qwen3-5/dummy_vision.onnx
similarity index 100%
rename from test/test_models/qwen35-hybrid-preprocessing/dummy_vision.onnx
rename to test/models/qwen3-5/dummy_vision.onnx
diff --git a/test/test_models/qwen35-hybrid-preprocessing/genai_config.json b/test/models/qwen3-5/genai_config.json
similarity index 100%
rename from test/test_models/qwen35-hybrid-preprocessing/genai_config.json
rename to test/models/qwen3-5/genai_config.json
diff --git a/test/test_models/qwen3-vl-vision-preprocessing/processor_config.json b/test/models/qwen3-5/processor_config.json
similarity index 100%
rename from test/test_models/qwen3-vl-vision-preprocessing/processor_config.json
rename to test/models/qwen3-5/processor_config.json
diff --git a/test/test_models/qwen3-vl-vision-preprocessing/special_tokens_map.json b/test/models/qwen3-5/special_tokens_map.json
similarity index 100%
rename from test/test_models/qwen3-vl-vision-preprocessing/special_tokens_map.json
rename to test/models/qwen3-5/special_tokens_map.json
diff --git a/test/test_models/qwen3-vl-vision-preprocessing/tokenizer.json b/test/models/qwen3-5/tokenizer.json
similarity index 100%
rename from test/test_models/qwen3-vl-vision-preprocessing/tokenizer.json
rename to test/models/qwen3-5/tokenizer.json
diff --git a/test/test_models/qwen3-vl-vision-preprocessing/tokenizer_config.json b/test/models/qwen3-5/tokenizer_config.json
similarity index 100%
rename from test/test_models/qwen3-vl-vision-preprocessing/tokenizer_config.json
rename to test/models/qwen3-5/tokenizer_config.json
diff --git a/test/test_models/qwen3-vl-vision-preprocessing/dummy_embedding.onnx b/test/models/qwen3-vl/dummy_embedding.onnx
similarity index 100%
rename from test/test_models/qwen3-vl-vision-preprocessing/dummy_embedding.onnx
rename to test/models/qwen3-vl/dummy_embedding.onnx
diff --git a/test/test_models/qwen3-vl-vision-preprocessing/dummy_text.onnx b/test/models/qwen3-vl/dummy_text.onnx
similarity index 100%
rename from test/test_models/qwen3-vl-vision-preprocessing/dummy_text.onnx
rename to test/models/qwen3-vl/dummy_text.onnx
diff --git a/test/test_models/qwen3-vl-vision-preprocessing/dummy_vision.onnx b/test/models/qwen3-vl/dummy_vision.onnx
similarity index 100%
rename from test/test_models/qwen3-vl-vision-preprocessing/dummy_vision.onnx
rename to test/models/qwen3-vl/dummy_vision.onnx
diff --git a/test/test_models/qwen3-vl-vision-preprocessing/genai_config.json b/test/models/qwen3-vl/genai_config.json
similarity index 100%
rename from test/test_models/qwen3-vl-vision-preprocessing/genai_config.json
rename to test/models/qwen3-vl/genai_config.json
diff --git a/test/test_models/qwen35-hybrid-preprocessing/processor_config.json b/test/models/qwen3-vl/processor_config.json
similarity index 100%
rename from test/test_models/qwen35-hybrid-preprocessing/processor_config.json
rename to test/models/qwen3-vl/processor_config.json
diff --git a/test/test_models/qwen35-hybrid-preprocessing/special_tokens_map.json b/test/models/qwen3-vl/special_tokens_map.json
similarity index 100%
rename from test/test_models/qwen35-hybrid-preprocessing/special_tokens_map.json
rename to test/models/qwen3-vl/special_tokens_map.json
diff --git a/test/test_models/qwen35-hybrid-preprocessing/tokenizer.json b/test/models/qwen3-vl/tokenizer.json
similarity index 100%
rename from test/test_models/qwen35-hybrid-preprocessing/tokenizer.json
rename to test/models/qwen3-vl/tokenizer.json
diff --git a/test/test_models/qwen35-hybrid-preprocessing/tokenizer_config.json b/test/models/qwen3-vl/tokenizer_config.json
similarity index 100%
rename from test/test_models/qwen35-hybrid-preprocessing/tokenizer_config.json
rename to test/models/qwen3-vl/tokenizer_config.json
diff --git a/test/test_models/audio-preprocessing/added_tokens.json b/test/models/whisper/added_tokens.json
similarity index 100%
rename from test/test_models/audio-preprocessing/added_tokens.json
rename to test/models/whisper/added_tokens.json
diff --git a/test/test_models/audio-preprocessing/audio_processor_config.json b/test/models/whisper/audio_processor_config.json
similarity index 100%
rename from test/test_models/audio-preprocessing/audio_processor_config.json
rename to test/models/whisper/audio_processor_config.json
diff --git a/test/test_models/audio-preprocessing/dummy_decoder.onnx b/test/models/whisper/dummy_decoder.onnx
similarity index 100%
rename from test/test_models/audio-preprocessing/dummy_decoder.onnx
rename to test/models/whisper/dummy_decoder.onnx
diff --git a/test/test_models/audio-preprocessing/dummy_encoder.onnx b/test/models/whisper/dummy_encoder.onnx
similarity index 100%
rename from test/test_models/audio-preprocessing/dummy_encoder.onnx
rename to test/models/whisper/dummy_encoder.onnx
diff --git a/test/test_models/audio-preprocessing/genai_config.json b/test/models/whisper/genai_config.json
similarity index 100%
rename from test/test_models/audio-preprocessing/genai_config.json
rename to test/models/whisper/genai_config.json
diff --git a/test/test_models/audio-preprocessing/special_tokens_map.json b/test/models/whisper/special_tokens_map.json
similarity index 100%
rename from test/test_models/audio-preprocessing/special_tokens_map.json
rename to test/models/whisper/special_tokens_map.json
diff --git a/test/test_models/audio-preprocessing/tokenizer.json b/test/models/whisper/tokenizer.json
similarity index 100%
rename from test/test_models/audio-preprocessing/tokenizer.json
rename to test/models/whisper/tokenizer.json
diff --git a/test/test_models/audio-preprocessing/tokenizer_config.json b/test/models/whisper/tokenizer_config.json
similarity index 100%
rename from test/test_models/audio-preprocessing/tokenizer_config.json
rename to test/models/whisper/tokenizer_config.json
diff --git a/test/test_models/audio-preprocessing/vocab.json b/test/models/whisper/vocab.json
similarity index 100%
rename from test/test_models/audio-preprocessing/vocab.json
rename to test/models/whisper/vocab.json
diff --git a/test/python/README.md b/test/python/README.md
index d316092d28..62cda8ab36 100644
--- a/test/python/README.md
+++ b/test/python/README.md
@@ -1,7 +1,7 @@
To run a test:
-python -m pytest -sv test_onnxruntime_genai_api.py -k "" --test_models ..\test_models
+python -m pytest -sv test_onnxruntime_genai_api.py -k "" --test_models ..\models
For example:
-python -m pytest -sv test_onnxruntime_genai_api.py -k "test_greedy_search" --test_models ..\test_models
+python -m pytest -sv test_onnxruntime_genai_api.py -k "test_greedy_search" --test_models ..\models
diff --git a/test/test_models/create_decoder_input_ids_test_models.py b/test/python/create/create_decoder_input_ids_test_models.py
similarity index 100%
rename from test/test_models/create_decoder_input_ids_test_models.py
rename to test/python/create/create_decoder_input_ids_test_models.py
diff --git a/test/test_models/create_dummy_lfm2_model.py b/test/python/create/create_dummy_lfm2_model.py
similarity index 100%
rename from test/test_models/create_dummy_lfm2_model.py
rename to test/python/create/create_dummy_lfm2_model.py
diff --git a/test/test_models/create_dummy_model.py b/test/python/create/create_dummy_model.py
similarity index 100%
rename from test/test_models/create_dummy_model.py
rename to test/python/create/create_dummy_model.py
diff --git a/test/test_models/qwen35-hybrid-preprocessing/create_dummy_models.py b/test/python/create/create_dummy_qwen_3.5_models.py
similarity index 98%
rename from test/test_models/qwen35-hybrid-preprocessing/create_dummy_models.py
rename to test/python/create/create_dummy_qwen_3.5_models.py
index fa447a458c..6f54dce9fe 100644
--- a/test/test_models/qwen35-hybrid-preprocessing/create_dummy_models.py
+++ b/test/python/create/create_dummy_qwen_3.5_models.py
@@ -11,7 +11,7 @@
shapes for testing the ort-genai runtime's auto-discovery and state management.
Usage:
- python create_qwen35_dummy_models.py --output test/test_models/qwen35-hybrid-preprocessing
+ python create_dummy_qwen_3.5_models.py --output test/models/qwen3-5
"""
import argparse
@@ -347,7 +347,7 @@ def main():
parser.add_argument(
"--output",
type=str,
- default="test/test_models/qwen35-hybrid-preprocessing",
+ default="test/models/qwen3-5",
help="Output directory for the dummy models",
)
args = parser.parse_args()
@@ -382,7 +382,7 @@ def main():
print(" Created genai_config.json")
# Copy tokenizer files from qwen3-vl test model if available
- src_dir = os.path.join(os.path.dirname(output_dir), "qwen3-vl-vision-preprocessing")
+ src_dir = os.path.join(os.path.dirname(output_dir), "qwen3-vl")
for fname in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "processor_config.json"]:
src = os.path.join(src_dir, fname)
if os.path.exists(src):
diff --git a/test/create_gqa_model.py b/test/python/create/create_gqa_model.py
similarity index 98%
rename from test/create_gqa_model.py
rename to test/python/create/create_gqa_model.py
index 6b1ad822bd..b61397535b 100644
--- a/test/create_gqa_model.py
+++ b/test/python/create/create_gqa_model.py
@@ -207,7 +207,7 @@ 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"))
+ "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}")
diff --git a/test/python/models/qwen_2.5_vl/run.sh b/test/python/models/qwen_2.5_vl/run.sh
deleted file mode 100644
index 6da708b743..0000000000
--- a/test/python/models/qwen_2.5_vl/run.sh
+++ /dev/null
@@ -1,65 +0,0 @@
-#!/bin/bash
-# -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT License. See License.txt in the project root for
-# license information.
-# --------------------------------------------------------------------------
-# This script builds and tests either an fp32, bf16 or fp16 Qwen2.5-VL-3B-Instruct model. Append -f to force export.
-# Usage: ./run.sh [fp32|bf16|fp16] [-f]
-
-# Exit immediately if a command fails
-set -e
-
-# 1. Validate Input
-if [ "$1" != "fp32" ] && [ "$1" != "bf16" ] && [ "$1" != "fp16" ]; then
- echo "Error: Invalid precision."
- echo "Usage: $0 fp32|bf16|fp16"
- exit 1
-fi
-
-# 2. Define variables based on input
-PRECISION=$1
-TEST_DIR="$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)"
-OUTPUT_DIR="${TEST_DIR}/qwen_${PRECISION}"
-ONNX_MODEL_PATH="${OUTPUT_DIR}/model.onnx"
-CACHE_DIR="${TEST_DIR}/cache"
-HF_MODEL="Qwen/Qwen2.5-VL-3B-Instruct"
-
-# Set the --bf16 or --fp16 flag for the test script
-TEST_FLAG=""
-if [ "$PRECISION" == "bf16" ]; then
- TEST_FLAG="--bf16"
-elif [ "$PRECISION" == "fp16" ]; then
- TEST_FLAG="--fp16"
-fi
-
-# 3. Remove output directory only if it exists and -f flag is provided.
-if [ "$2" == "-f" ] && [ -d "${OUTPUT_DIR}" ]; then
- echo "Removing existing directory: ${OUTPUT_DIR}"
- rm -rf "${OUTPUT_DIR}"
-fi
-
-BUILDER_DIR="$(cd ../../../../src/python/py/models && pwd)"
-
-# 4. Run the builder script if output directory does not exist.
-if ! [ -d "${OUTPUT_DIR}" ]; then
- echo "--- Building ${PRECISION} model ---"
- cd "${BUILDER_DIR}"
- python builder.py \
- -m ${HF_MODEL} \
- -p ${PRECISION} \
- -o ${OUTPUT_DIR} \
- -e cuda \
- -c ${CACHE_DIR}
-fi
-
-# 5. Run the parity test
-cd "${TEST_DIR}"
-echo "--- Testing ${PRECISION} model parity ---"
-python test_qwen_2.5_vl.py \
- --hf_model ${HF_MODEL} \
- --cache_dir ${CACHE_DIR} \
- --onnx_model ${ONNX_MODEL_PATH} \
- ${TEST_FLAG}
-
-echo "--- ${PRECISION} run complete ---"
\ No newline at end of file
diff --git a/test/python/test_decoder_state_input_ids.py b/test/python/models/test_decoder_state_input_ids.py
similarity index 98%
rename from test/python/test_decoder_state_input_ids.py
rename to test/python/models/test_decoder_state_input_ids.py
index aa00c8331f..65432e754f 100644
--- a/test/python/test_decoder_state_input_ids.py
+++ b/test/python/models/test_decoder_state_input_ids.py
@@ -11,7 +11,7 @@
Fix: use a decoder-only SessionInfo for the HasInput('input_ids') check.
-Two test model variants (in test/test_models/):
+Two test model variants (in test/models/):
- multimodal-decoder-no-input-ids/ Mistral3-like: embedding has input_ids,
decoder does NOT.
- multimodal-decoder-with-input-ids/ Gemma4-like: both embedding and decoder
diff --git a/test/python/test_gemma4_models.py b/test/python/models/test_gemma4_models.py
similarity index 98%
rename from test/python/test_gemma4_models.py
rename to test/python/models/test_gemma4_models.py
index 183ab83fb8..5bb2630c5f 100644
--- a/test/python/test_gemma4_models.py
+++ b/test/python/models/test_gemma4_models.py
@@ -6,8 +6,8 @@
Tests cover model loading, text-only processing, and image understanding.
This file can be used in two ways:
-1. As a pytest module: pytest test_gemma4_models.py --test_models=/path/to/test_models
-2. As a standalone runner: python test_gemma4_models.py --cwd test/python --test_models test/test_models
+1. As a pytest module: pytest test_gemma4_models.py --test_models=/path/to/models
+2. As a standalone runner: python test_gemma4_models.py --cwd test/python --test_models test/models
"""
import argparse
@@ -24,7 +24,7 @@
logging.basicConfig(format="%(asctime)s %(name)s [%(levelname)s] - %(message)s", level=logging.DEBUG)
log = logging.getLogger("gemma4-tests")
-GEMMA4_MODEL_NAME = "gemma4-vision-preprocessing"
+GEMMA4_MODEL_NAME = "gemma4"
def _get_gemma4_model_path(test_data_path):
@@ -294,7 +294,7 @@ def run_gemma4_vision_tests(
"-m",
"pytest",
"-sv",
- "test_gemma4_models.py",
+ os.path.abspath(__file__),
"--test_models",
test_models,
]
@@ -311,8 +311,8 @@ def parse_arguments():
)
parser.add_argument(
"--test_models",
- help="Path to the test_models directory",
- default=Path(__file__).parent.parent.resolve().absolute() / "test_models",
+ help="Path to the 'models' directory",
+ default=Path(__file__).parent.parent.resolve().absolute() / "models",
)
return parser.parse_args()
diff --git a/test/python/test_mistral3_preprocessor.py b/test/python/models/test_mistral3_preprocessor.py
similarity index 100%
rename from test/python/test_mistral3_preprocessor.py
rename to test/python/models/test_mistral3_preprocessor.py
diff --git a/test/python/test_mistral3_tokens.py b/test/python/models/test_mistral3_tokens.py
similarity index 98%
rename from test/python/test_mistral3_tokens.py
rename to test/python/models/test_mistral3_tokens.py
index b3bfc48852..50aac7c482 100644
--- a/test/python/test_mistral3_tokens.py
+++ b/test/python/models/test_mistral3_tokens.py
@@ -189,13 +189,13 @@ def test_genai_processor_token_counts(self, test_data_path):
"""Verify C++ processor produces correct token counts.
Requires onnxruntime_genai, the pre-exported Mistral3 model under
- test_data_path/mistral3-vision-preprocessing, and a test image.
+ test_data_path/mistral3, and a test image.
"""
from pathlib import Path
import numpy as np
- model_path = Path(test_data_path) / "mistral3-vision-preprocessing"
+ model_path = Path(test_data_path) / "mistral3"
image_path = Path(test_data_path) / "images" / "australia.jpg"
if not (model_path / "genai_config.json").is_file():
pytest.skip(f"Mistral3 model not found at {model_path} (missing genai_config.json)")
diff --git a/test/python/test_parakeet_tdt.py b/test/python/models/test_parakeet_tdt.py
similarity index 100%
rename from test/python/test_parakeet_tdt.py
rename to test/python/models/test_parakeet_tdt.py
diff --git a/test/python/test_quantized_model.py b/test/python/models/test_quantized_model.py
similarity index 100%
rename from test/python/test_quantized_model.py
rename to test/python/models/test_quantized_model.py
diff --git a/test/python/models/qwen_2.5_vl/test_qwen_2.5_vl.py b/test/python/models/test_qwen_2.5_vl.py
similarity index 100%
rename from test/python/models/qwen_2.5_vl/test_qwen_2.5_vl.py
rename to test/python/models/test_qwen_2.5_vl.py
diff --git a/test/python/test_qwen35_text_only.py b/test/python/models/test_qwen_3.5_text_only.py
similarity index 87%
rename from test/python/test_qwen35_text_only.py
rename to test/python/models/test_qwen_3.5_text_only.py
index e263e66637..b62fd308f4 100644
--- a/test/python/test_qwen35_text_only.py
+++ b/test/python/models/test_qwen_3.5_text_only.py
@@ -9,7 +9,7 @@
"qwen3_5_text"), which uses 2D position_ids and hybrid KV/recurrent state.
Usage:
- pytest test_qwen35_text_only.py --test_models=test/test_models
+ pytest test_qwen_3.5_text_only.py --test_models=test/models
"""
import os
@@ -18,7 +18,7 @@
import onnxruntime_genai as og
import pytest
-MODEL_DIR = "qwen35-text-only"
+MODEL_DIR = "qwen3-5-text-only"
def _model_path(test_data_path):
@@ -32,14 +32,14 @@ def _skip_if_missing(test_data_path):
return path
-def test_qwen35_text_only_model_loads(test_data_path):
+def test_qwen3_5_text_only_model_loads(test_data_path):
"""Test that a Qwen3.5 text-only model loads successfully."""
model_path = _skip_if_missing(test_data_path)
model = og.Model(model_path)
assert model is not None
-def test_qwen35_text_only_generator_creates(test_data_path):
+def test_qwen3_5_text_only_generator_creates(test_data_path):
"""Test that a Generator can be created for the text-only model.
Validates that hybrid state auto-discovery works with qwen3_5_text type."""
model_path = _skip_if_missing(test_data_path)
@@ -50,7 +50,7 @@ def test_qwen35_text_only_generator_creates(test_data_path):
assert generator is not None
-def test_qwen35_text_only_accepts_input_ids(test_data_path):
+def test_qwen3_5_text_only_accepts_input_ids(test_data_path):
"""Test that the text-only model accepts input_ids (not inputs_embeds).
The dummy model uses Identity pass-through which doesn't support KV cache
shape changes, so we only validate that the generator constructs and
diff --git a/test/python/test_qwen_fara_models.py b/test/python/models/test_qwen_fara_models.py
similarity index 83%
rename from test/python/test_qwen_fara_models.py
rename to test/python/models/test_qwen_fara_models.py
index 605f553d9c..ef20e7088b 100644
--- a/test/python/test_qwen_fara_models.py
+++ b/test/python/models/test_qwen_fara_models.py
@@ -6,8 +6,8 @@
Tests cover model loading, text generation, image understanding, and multimodal processing.
This file can be used in two ways:
-1. As a pytest module: pytest test_qwen_fara_models.py --test_models=/path/to/test_models
-2. As a standalone runner: python test_qwen_fara_models.py --cwd test/python --test_models test/test_models
+1. As a pytest module: pytest test_qwen_fara_models.py --test_models=/path/to/models
+2. As a standalone runner: python test_qwen_fara_models.py --cwd test/python --test_models test/models
"""
import argparse
@@ -27,9 +27,7 @@
log = logging.getLogger("qwen-fara-vision-tests")
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"])
def test_qwen_fara_vision_basic(test_data_path, relative_model_path, relative_image_path):
"""Test basic vision preprocessing for Qwen/Fara-style models."""
@@ -38,7 +36,7 @@ def test_qwen_fara_vision_basic(test_data_path, relative_model_path, relative_im
processor = model.create_multimodal_processor()
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
images = og.Images.open(image_path)
# Test with Qwen vision tokens
@@ -50,9 +48,7 @@ def test_qwen_fara_vision_basic(test_data_path, relative_model_path, relative_im
assert "pixel_values" in inputs
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "landscape.jpg"])
def test_qwen_fara_vision_load_from_bytes(test_data_path, relative_model_path, relative_image_path):
"""Test loading images from bytes for Qwen/Fara models."""
@@ -61,7 +57,7 @@ def test_qwen_fara_vision_load_from_bytes(test_data_path, relative_model_path, r
processor = model.create_multimodal_processor()
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
images = None
with open(image_path, "rb") as image:
image_bytes = image.read()
@@ -74,9 +70,7 @@ def test_qwen_fara_vision_load_from_bytes(test_data_path, relative_model_path, r
assert "pixel_values" in inputs
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize(
"relative_image_paths",
[[Path("images") / "australia.jpg", Path("images") / "landscape.jpg"]],
@@ -96,7 +90,7 @@ def test_qwen_fara_vision_multiple_images(test_data_path, relative_model_path, r
processor = model.create_multimodal_processor()
image_paths = [
- os.fspath(Path(test_data_path) / relative_image_path) for relative_image_path in relative_image_paths
+ os.fspath(Path(test_data_path).parent / relative_image_path) for relative_image_path in relative_image_paths
]
images = og.Images.open(*image_paths)
@@ -120,7 +114,9 @@ def test_qwen3_vl_vision_dynamic_grid_dim(test_data_path):
"""
onnx = pytest.importorskip("onnx")
- vision_path = os.path.join(test_data_path, "qwen3-vl-vision-preprocessing", "dummy_vision.onnx")
+ vision_path = os.path.join(
+ test_data_path, "qwen3-vl", "dummy_vision.onnx"
+ )
model = onnx.load(vision_path)
# Find image_grid_thw input
@@ -144,9 +140,7 @@ def test_qwen3_vl_vision_dynamic_grid_dim(test_data_path):
assert dim1.dim_value == 3, f"image_grid_thw dim-1 should be 3, got {dim1.dim_value}"
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
def test_qwen_fara_text_only_generation(test_data_path, relative_model_path):
"""Test text-only generation without images."""
model_path = os.fspath(Path(test_data_path) / relative_model_path)
@@ -163,9 +157,7 @@ def test_qwen_fara_text_only_generation(test_data_path, relative_model_path):
assert "input_ids" in inputs
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "sheet.png"])
def test_qwen_fara_vision_with_special_tokens(test_data_path, relative_model_path, relative_image_path):
"""Test vision processing with special tokens in prompt."""
@@ -174,7 +166,7 @@ def test_qwen_fara_vision_with_special_tokens(test_data_path, relative_model_pat
processor = model.create_multimodal_processor()
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
images = og.Images.open(image_path)
# Test with Qwen vision tokens
@@ -186,9 +178,7 @@ def test_qwen_fara_vision_with_special_tokens(test_data_path, relative_model_pat
assert "input_ids" in inputs
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "10809054.jpg"])
def test_qwen_fara_vision_different_image_formats(test_data_path, relative_model_path, relative_image_path):
"""Test processing different image formats."""
@@ -197,7 +187,7 @@ def test_qwen_fara_vision_different_image_formats(test_data_path, relative_model
processor = model.create_multimodal_processor()
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
images = og.Images.open(image_path)
prompt = "<|vision_start|><|image_pad|><|vision_end|>Analyze this image"
@@ -207,9 +197,7 @@ def test_qwen_fara_vision_different_image_formats(test_data_path, relative_model
assert "pixel_values" in inputs
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"])
def test_qwen_fara_accuracy_comparison(test_data_path, relative_model_path, relative_image_path):
"""
@@ -221,7 +209,7 @@ def test_qwen_fara_accuracy_comparison(test_data_path, relative_model_path, rela
pytest.skip("PyTorch or transformers not available for accuracy comparison")
model_path = os.fspath(Path(test_data_path) / relative_model_path)
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
# Load ONNX model
onnx_model = og.Model(model_path)
@@ -240,7 +228,7 @@ def test_qwen_fara_accuracy_comparison(test_data_path, relative_model_path, rela
onnx_pixel_values = onnx_inputs["pixel_values"]
- # For vision-preprocessing dummy model, we validate the preprocessing pipeline
+ # For vision dummy model, we validate the preprocessing pipeline
# In a real scenario, you would:
# 1. Load the same PyTorch model
# 2. Process the same image with PyTorch processor
@@ -271,9 +259,7 @@ def test_qwen_fara_accuracy_comparison(test_data_path, relative_model_path, rela
log.debug(f"ONNX pixel_values range: [{pixel_min:.4f}, {pixel_max:.4f}]")
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "sheet.png"])
def test_qwen_fara_preprocessing_consistency(test_data_path, relative_model_path, relative_image_path):
"""
@@ -281,7 +267,7 @@ def test_qwen_fara_preprocessing_consistency(test_data_path, relative_model_path
This validates deterministic behavior of the preprocessing pipeline.
"""
model_path = os.fspath(Path(test_data_path) / relative_model_path)
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
model = og.Model(model_path)
processor = model.create_multimodal_processor()
@@ -357,9 +343,9 @@ def test_qwen3_vl_model_type(test_data_path):
Test that the Qwen3-VL model loads with the correct model type (qwen3_vl).
Validates that the model type routing in model.cpp correctly handles the new type.
"""
- model_path = os.fspath(Path(test_data_path) / "qwen3-vl-vision-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-vl")
if not os.path.exists(model_path):
- pytest.skip("qwen3-vl-vision-preprocessing test model not found")
+ pytest.skip("qwen3-vl test model not found")
model = og.Model(model_path)
assert model is not None
@@ -375,14 +361,14 @@ def test_qwen3_vl_pixel_values_shape(test_data_path, relative_image_path):
Test that Qwen3-VL preprocesses images with patch_size=16 (not 14 like Qwen2.5-VL).
The patch_size difference affects the number of patches extracted from each image.
"""
- model_path = os.fspath(Path(test_data_path) / "qwen3-vl-vision-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-vl")
if not os.path.exists(model_path):
- pytest.skip("qwen3-vl-vision-preprocessing test model not found")
+ pytest.skip("qwen3-vl test model not found")
model = og.Model(model_path)
processor = model.create_multimodal_processor()
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
images = og.Images.open(image_path)
prompt = "<|vision_start|><|image_pad|><|vision_end|>Describe this image"
@@ -412,8 +398,8 @@ def test_qwen3_vl_pixel_values_shape(test_data_path, relative_image_path):
@pytest.mark.parametrize(
"model_name,expected_patch_dim",
[
- ("qwen-vision-preprocessing", 1176), # Qwen2.5-VL: patch_size=14, 14*14*3*2=1176
- ("qwen3-vl-vision-preprocessing", 1536), # Qwen3-VL: patch_size=16, 16*16*3*2=1536
+ ("qwen2-5-vl", 1176), # Qwen2.5-VL: patch_size=14, 14*14*3*2=1176
+ ("qwen3-vl", 1536), # Qwen3-VL: patch_size=16, 16*16*3*2=1536
],
)
@pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"])
@@ -429,7 +415,7 @@ def test_qwen_vl_family_patch_size_difference(test_data_path, model_name, expect
model = og.Model(model_path)
processor = model.create_multimodal_processor()
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
images = og.Images.open(image_path)
prompt = "<|vision_start|><|image_pad|><|vision_end|>Describe this image"
@@ -451,9 +437,7 @@ def test_qwen_vl_family_patch_size_difference(test_data_path, model_name, expect
log.debug(f"{model_name} pixel_values shape: {pixel_array.shape}")
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"])
def test_qwen_vl_preprocessing_output_completeness(test_data_path, relative_model_path, relative_image_path):
"""
@@ -468,7 +452,7 @@ def test_qwen_vl_preprocessing_output_completeness(test_data_path, relative_mode
model = og.Model(model_path)
processor = model.create_multimodal_processor()
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
images = og.Images.open(image_path)
prompt = "<|vision_start|><|image_pad|><|vision_end|>Describe this image"
@@ -510,9 +494,7 @@ def _to_numpy(tensor):
log.debug(f"{relative_model_path} output: pv={pv.shape}, grid={grid}, nit={nit}, ids={ids.shape}")
-@pytest.mark.parametrize(
- "relative_model_path", [Path("qwen-vision-preprocessing"), Path("qwen3-vl-vision-preprocessing")]
-)
+@pytest.mark.parametrize("relative_model_path", [Path("qwen2-5-vl"), Path("qwen3-vl")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"])
def test_qwen_vl_image_grid_thw_consistency(test_data_path, relative_model_path, relative_image_path):
"""
@@ -529,7 +511,7 @@ def test_qwen_vl_image_grid_thw_consistency(test_data_path, relative_model_path,
model = og.Model(model_path)
processor = model.create_multimodal_processor()
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
images = og.Images.open(image_path)
prompt = "<|vision_start|><|image_pad|><|vision_end|>Describe this image"
@@ -568,12 +550,12 @@ def test_qwen_vl_normalization_range_difference(test_data_path, relative_image_p
Qwen3-VL uses mean/std=[0.5, 0.5, 0.5] → pixel range [-1, 1].
Qwen2.5-VL uses OpenAI CLIP normalization → wider range.
"""
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
# Qwen3-VL: normalized with [0.5, 0.5, 0.5]
- q3_model_path = os.fspath(Path(test_data_path) / "qwen3-vl-vision-preprocessing")
+ q3_model_path = os.fspath(Path(test_data_path) / "qwen3-vl")
if not os.path.exists(q3_model_path):
- pytest.skip("qwen3-vl-vision-preprocessing test model not found")
+ pytest.skip("qwen3-vl test model not found")
q3_model = og.Model(q3_model_path)
q3_proc = q3_model.create_multimodal_processor()
@@ -582,9 +564,9 @@ def test_qwen_vl_normalization_range_difference(test_data_path, relative_image_p
q3_pv = q3_inputs["pixel_values"].as_numpy()
# Qwen2.5-VL: normalized with OpenAI CLIP constants
- q25_model_path = os.fspath(Path(test_data_path) / "qwen-vision-preprocessing")
+ q25_model_path = os.fspath(Path(test_data_path) / "qwen2-5-vl")
if not os.path.exists(q25_model_path):
- pytest.skip("qwen-vision-preprocessing test model not found")
+ pytest.skip("qwen2-5-vl test model not found")
q25_model = og.Model(q25_model_path)
q25_proc = q25_model.create_multimodal_processor()
@@ -613,44 +595,44 @@ def test_qwen_vl_normalization_range_difference(test_data_path, relative_image_p
# ---------------------------------------------------------------------------
-def test_qwen35_hybrid_model_loads(test_data_path):
+def test_qwen3_5_hybrid_model_loads(test_data_path):
"""Test that a Qwen3.5 hybrid model (with recurrent + KV states) loads successfully."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
model = og.Model(model_path)
assert model is not None
-def test_qwen35_hybrid_creates_processor(test_data_path):
+def test_qwen3_5_hybrid_creates_processor(test_data_path):
"""Test that the qwen3_5 model type routes to QwenImageProcessor."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
model = og.Model(model_path)
processor = model.create_multimodal_processor()
assert processor is not None
-def test_qwen35_hybrid_tokenizer(test_data_path):
+def test_qwen3_5_hybrid_tokenizer(test_data_path):
"""Test that tokenizer can be created for the qwen3_5 model type."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
model = og.Model(model_path)
tokenizer = og.Tokenizer(model)
assert tokenizer is not None
-def test_qwen35_hybrid_generator_creates(test_data_path):
+def test_qwen3_5_hybrid_generator_creates(test_data_path):
"""Test that a Generator can be created for the hybrid model.
This validates that RecurrentState and sparse KV cache auto-discovery don't crash."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
model = og.Model(model_path)
params = og.GeneratorParams(model)
@@ -659,14 +641,14 @@ def test_qwen35_hybrid_generator_creates(test_data_path):
assert generator is not None
-def test_qwen35_hybrid_text_generation(test_data_path):
+def test_qwen3_5_hybrid_text_generation(test_data_path):
"""Test basic text generation with the hybrid model.
The dummy model uses Identity pass-through which doesn't support KV cache
shape changes, so we only validate that the generator constructs and
the first forward pass (prefill) executes without errors on the recurrent state path."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
model = og.Model(model_path)
@@ -680,13 +662,13 @@ def test_qwen35_hybrid_text_generation(test_data_path):
@pytest.mark.parametrize("relative_image_path", [Path("images") / "australia.jpg"])
-def test_qwen35_hybrid_vision_preprocessing(test_data_path, relative_image_path):
+def test_qwen3_5_hybrid_vision_preprocessing(test_data_path, relative_image_path):
"""Test that the hybrid model processes images through the vision pipeline."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
- image_path = os.fspath(Path(test_data_path) / relative_image_path)
+ image_path = os.fspath(Path(test_data_path).parent / relative_image_path)
if not os.path.exists(image_path):
pytest.skip(f"Test image not found: {image_path}")
@@ -707,12 +689,12 @@ def test_qwen35_hybrid_vision_preprocessing(test_data_path, relative_image_path)
@pytest.mark.skipif(not og.is_cuda_available(), reason="CUDA EP not available")
-def test_qwen35_hybrid_generator_creates_cuda(test_data_path):
+def test_qwen3_5_hybrid_generator_creates_cuda(test_data_path):
"""Test that a Generator can be created for the hybrid model on CUDA.
Validates RecurrentState shared-buffer path on CUDA EP."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
config = og.Config(model_path)
config.clear_providers()
@@ -725,12 +707,12 @@ def test_qwen35_hybrid_generator_creates_cuda(test_data_path):
@pytest.mark.skipif(not og.is_cuda_available(), reason="CUDA EP not available")
-def test_qwen35_hybrid_text_generation_cuda(test_data_path):
+def test_qwen3_5_hybrid_text_generation_cuda(test_data_path):
"""Test that the hybrid model generator constructs and prefill executes on CUDA.
RecurrentState uses shared buffers (same tensor as input and output)."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
config = og.Config(model_path)
config.clear_providers()
@@ -757,13 +739,13 @@ def _is_webgpu_test_enabled():
@pytest.mark.skipif(not _is_webgpu_test_enabled(), reason="WebGPU EP not available or TEST_WEBGPU not set")
-def test_qwen35_hybrid_generator_creates_webgpu(test_data_path):
+def test_qwen3_5_hybrid_generator_creates_webgpu(test_data_path):
"""Test that a Generator can be created for the hybrid model on WebGPU.
Validates RecurrentState separate-buffer path (WebGPU cannot alias
input/output buffers in the same compute pass)."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
config = og.Config(model_path)
config.clear_providers()
@@ -776,13 +758,13 @@ def test_qwen35_hybrid_generator_creates_webgpu(test_data_path):
@pytest.mark.skipif(not _is_webgpu_test_enabled(), reason="WebGPU EP not available or TEST_WEBGPU not set")
-def test_qwen35_hybrid_text_generation_webgpu(test_data_path):
+def test_qwen3_5_hybrid_text_generation_webgpu(test_data_path):
"""Test that the hybrid model generator constructs and prefill executes on WebGPU.
RecurrentState uses separate past/present buffers to avoid the WebGPU
buffer aliasing restriction (Storage read-write | read-only conflict)."""
- model_path = os.fspath(Path(test_data_path) / "qwen35-hybrid-preprocessing")
+ model_path = os.fspath(Path(test_data_path) / "qwen3-5")
if not os.path.exists(model_path):
- pytest.skip("qwen35-hybrid-preprocessing test model not found")
+ pytest.skip("qwen3-5 test model not found")
config = og.Config(model_path)
config.clear_providers()
@@ -808,7 +790,7 @@ def run_qwen_fara_vision_tests(
"-m",
"pytest",
"-sv",
- "test_qwen_fara_models.py",
+ os.path.abspath(__file__),
"--test_models",
test_models,
]
@@ -825,8 +807,8 @@ def parse_arguments():
)
parser.add_argument(
"--test_models",
- help="Path to the test_models directory",
- default=pathlib.Path(__file__).parent.parent.resolve().absolute() / "test_models",
+ help="Path to the 'models' directory",
+ default=pathlib.Path(__file__).parent.parent.resolve().absolute() / "models",
)
return parser.parse_args()
diff --git a/test/python/test_yarn_rope_parity.py b/test/python/models/test_yarn_rope_parity.py
similarity index 100%
rename from test/python/test_yarn_rope_parity.py
rename to test/python/models/test_yarn_rope_parity.py
diff --git a/test/python/test_onnxruntime_genai.py b/test/python/test_onnxruntime_genai.py
index 0f77b5906c..3df4b13b69 100644
--- a/test/python/test_onnxruntime_genai.py
+++ b/test/python/test_onnxruntime_genai.py
@@ -9,8 +9,8 @@
import onnxruntime_genai as og
from _test_utils import download_models, run_subprocess
-from test_gemma4_models import run_gemma4_vision_tests
-from test_qwen_fara_models import run_qwen_fara_vision_tests
+from models.test_gemma4_models import run_gemma4_vision_tests
+from models.test_qwen_fara_models import run_qwen_fara_vision_tests
logging.basicConfig(format="%(asctime)s %(name)s [%(levelname)s] - %(message)s", level=logging.DEBUG)
log = logging.getLogger("onnxruntime-genai-tests")
@@ -60,8 +60,8 @@ def parse_arguments():
)
parser.add_argument(
"--test_models",
- help="Path to the test_models directory",
- default=pathlib.Path(__file__).parent.parent.resolve().absolute() / "test_models",
+ help="Path to the 'models' directory",
+ default=pathlib.Path(__file__).parent.parent.resolve().absolute() / "models",
)
parser.add_argument(
"--e2e",
diff --git a/test/python/test_onnxruntime_genai_api.py b/test/python/test_onnxruntime_genai_api.py
index 8f6721a137..1bd1686b0f 100644
--- a/test/python/test_onnxruntime_genai_api.py
+++ b/test/python/test_onnxruntime_genai_api.py
@@ -10,15 +10,11 @@
from pathlib import Path
import numpy as np
+import onnx
import onnxruntime
import onnxruntime_genai as og
import pytest
-if not sysconfig.get_platform().endswith("arm64"):
- # Skip importing onnx if running on ARM64
- # TODO(justinchuby): ONNX 1.18 supports arm64. Remove the condition when
- # there is a version bump
- import onnx
devices = ["cpu"]
@@ -38,8 +34,13 @@
devices.append("webgpu")
+@pytest.fixture
+def test_data_path(request):
+ return os.fspath(Path(request.config.getoption("--test_models")).parent)
+
+
def test_config(test_data_path):
- model_path = os.fspath(Path(test_data_path) / "hf-internal-testing" / "tiny-random-gpt2-fp32")
+ model_path = os.fspath(Path(test_data_path) / "models" / "hf-internal-testing" / "tiny-random-gpt2-fp32")
config = og.Config(model_path)
config.clear_providers()
config.append_provider("cuda")
@@ -59,7 +60,7 @@ def _log_callback(log: str):
og.set_log_options(enabled=True, generate_next_token=True)
og.set_log_callback(_log_callback)
- model_path = os.fspath(Path(test_data_path) / "hf-internal-testing" / "tiny-random-gpt2-fp32")
+ model_path = os.fspath(Path(test_data_path) / "models" / "hf-internal-testing" / "tiny-random-gpt2-fp32")
config = og.Config(model_path)
model = og.Model(config)
@@ -86,7 +87,7 @@ def _log_callback(log: str):
with tempfile.NamedTemporaryFile(mode="w+", suffix=".txt", delete=False) as log_file:
og.set_log_options(enabled=True, generate_next_token=True, filename=log_file.name)
- model_path = os.fspath(Path(test_data_path) / "hf-internal-testing" / "tiny-random-gpt2-fp32")
+ model_path = os.fspath(Path(test_data_path) / "models" / "hf-internal-testing" / "tiny-random-gpt2-fp32")
config = og.Config(model_path)
model = og.Model(config)
@@ -139,7 +140,7 @@ def test_NamedTensors():
),
)
def test_greedy_search(test_data_path, relative_model_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
config = og.Config(model_path) # Test using config vs path directly
model = og.Model(config)
@@ -189,7 +190,7 @@ def test_greedy_search(test_data_path, relative_model_path):
),
)
def test_rewind_cuda(test_data_path, relative_model_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
@@ -246,7 +247,7 @@ def test_rewind_cuda(test_data_path, relative_model_path):
([Path("hf-internal-testing") / "tiny-random-gpt2-fp32"]),
)
def test_rewind(test_data_path, relative_model_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
@@ -479,7 +480,7 @@ def test_load_model_from_memory(device, wrapper_bytes_function, phi2_for):
),
)
def test_model_device_type(test_data_path, relative_model_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path[0])
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path[0])
model = og.Model(model_path)
@@ -501,7 +502,7 @@ def test_model_device_type(test_data_path, relative_model_path):
),
)
def test_get_output(test_data_path, relative_model_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
@@ -619,10 +620,10 @@ def _split(onnx_model_path: os.PathLike, output_dir: os.PathLike):
_split(
Path(phi2_for("cuda")) / "model.onnx",
- Path(test_data_path) / relative_model_path,
+ Path(test_data_path) / "models" / relative_model_path,
)
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
tokenizer = og.Tokenizer(model)
@@ -656,10 +657,10 @@ def _split(onnx_model_path: os.PathLike, output_dir: os.PathLike):
assert equal
-@pytest.mark.parametrize("relative_model_path", [Path("vision-preprocessing")])
+@pytest.mark.parametrize("relative_model_path", [Path("phi3-v")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "sheet.png"])
-def test_vision_preprocessing(test_data_path, relative_model_path, relative_image_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+def test_phi3v_preprocessing(test_data_path, relative_model_path, relative_image_path):
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
processor = model.create_multimodal_processor()
@@ -671,10 +672,10 @@ def test_vision_preprocessing(test_data_path, relative_model_path, relative_imag
_ = processor(prompt, images=images)
-@pytest.mark.parametrize("relative_model_path", [Path("vision-preprocessing")])
+@pytest.mark.parametrize("relative_model_path", [Path("phi3-v")])
@pytest.mark.parametrize("relative_image_path", [Path("images") / "sheet.png"])
-def test_vision_preprocessing_load_image_from_bytes(test_data_path, relative_model_path, relative_image_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+def test_phi3v_preprocessing_load_image_from_bytes(test_data_path, relative_model_path, relative_image_path):
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
processor = model.create_multimodal_processor()
@@ -689,13 +690,13 @@ def test_vision_preprocessing_load_image_from_bytes(test_data_path, relative_mod
_ = processor(prompt, images=images)
-@pytest.mark.parametrize("relative_model_path", [Path("vision-preprocessing")])
+@pytest.mark.parametrize("relative_model_path", [Path("phi3-v")])
@pytest.mark.parametrize(
"relative_image_paths",
[[Path("images") / "australia.jpg", Path("images") / "sheet.png"]],
)
-def test_vision_preprocessing_multiple_images(test_data_path, relative_model_path, relative_image_paths):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+def test_phi3v_preprocessing_multiple_images(test_data_path, relative_model_path, relative_image_paths):
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
processor = model.create_multimodal_processor()
@@ -723,7 +724,7 @@ def test_adapters(test_data_path, device, multiple_adapters, phi2_for):
def _prepare_adapter_model(test_data_path):
phi2_model_path = phi2_for(device)
relative_model_path = "multiple_adapters" if multiple_adapters else "adapters"
- adapter_model_path = os.fspath(Path(test_data_path) / relative_model_path)
+ adapter_model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
if os.path.exists(adapter_model_path):
shutil.rmtree(adapter_model_path)
@@ -852,7 +853,7 @@ def test_preset_extra_inputs(test_data_path, device, phi2_for, extra_inputs):
def _prepare_model(test_data_path):
phi2_model_path = phi2_for(device)
relative_model_path = "preset_extra_inputs"
- extra_inputs_model_path = os.fspath(Path(test_data_path) / relative_model_path)
+ extra_inputs_model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
shutil.copytree(phi2_model_path, extra_inputs_model_path, dirs_exist_ok=True)
@@ -924,10 +925,10 @@ def _prepare_model(test_data_path):
generator.generate_next_token()
-@pytest.mark.parametrize("relative_model_path", [Path("audio-preprocessing")])
+@pytest.mark.parametrize("relative_model_path", [Path("whisper")])
@pytest.mark.parametrize("relative_audio_path", [Path("audios") / "1272-141231-0002.mp3"])
-def test_audio_preprocessing(test_data_path, relative_model_path, relative_audio_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+def test_whisper_preprocessing(test_data_path, relative_model_path, relative_audio_path):
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
processor = model.create_multimodal_processor()
@@ -941,10 +942,10 @@ def test_audio_preprocessing(test_data_path, relative_model_path, relative_audio
_ = processor(prompts, audios=audios)
-@pytest.mark.parametrize("relative_model_path", [Path("audio-preprocessing")])
+@pytest.mark.parametrize("relative_model_path", [Path("whisper")])
@pytest.mark.parametrize("relative_audio_path", [Path("audios") / "1272-141231-0002.mp3"])
-def test_audio_preprocessing_single_prompt(test_data_path, relative_model_path, relative_audio_path):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+def test_whisper_preprocessing_single_prompt(test_data_path, relative_model_path, relative_audio_path):
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
processor = model.create_multimodal_processor()
@@ -957,13 +958,13 @@ def test_audio_preprocessing_single_prompt(test_data_path, relative_model_path,
_ = processor(prompt, audios=audios)
-@pytest.mark.parametrize("relative_model_path", [Path("audio-preprocessing")])
+@pytest.mark.parametrize("relative_model_path", [Path("whisper")])
@pytest.mark.parametrize(
"relative_audio_paths",
[[Path("audios") / "1272-141231-0002.mp3"], [Path("audios") / "jfk.flac"]],
)
-def test_audio_preprocessing_multiple_audios(test_data_path, relative_model_path, relative_audio_paths):
- model_path = os.fspath(Path(test_data_path) / relative_model_path)
+def test_whisper_preprocessing_multiple_audios(test_data_path, relative_model_path, relative_audio_paths):
+ model_path = os.fspath(Path(test_data_path) / "models" / relative_model_path)
model = og.Model(model_path)
processor = model.create_multimodal_processor()
diff --git a/test/python/test_onnxruntime_genai_e2e.py b/test/python/test_onnxruntime_genai_e2e.py
index 199654aa19..1cf826fb79 100644
--- a/test/python/test_onnxruntime_genai_e2e.py
+++ b/test/python/test_onnxruntime_genai_e2e.py
@@ -47,13 +47,13 @@ def run_whisper():
num_beams = 5
(audio_path, expected_transcription) = (
- os.path.join(cwd, "..", "test_models", "audios", "1272-141231-0002.mp3"),
+ os.path.join(cwd, "..", "audios", "1272-141231-0002.mp3"),
"The cut on his chest is still dripping blood. The ache of his overstrained eyes. Even the soaring arena around him with thousands of spectators, retrievalidies not worth thinking about.",
)
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}")
+ built_model = os.path.join(cwd, "..", "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, enable_graph_capture=False)
@@ -97,7 +97,7 @@ def run_tool_calling():
for (model_name, tool_call_start, tool_call_end) in tool_call_models:
for (precision, execution_provider) in [("int4", "cpu")]: # TODO: add ("int4", "cuda"), ("int4", "dml") in CIs later
- model_path = os.path.join(cwd, "..", "test_models", model_name, precision, execution_provider)
+ model_path = os.path.join(cwd, "..", "models", model_name, precision, execution_provider)
if not os.path.exists(model_path): continue
# Run special_tokens.py to mark tool call token ids as special
@@ -126,7 +126,7 @@ def run_tool_calling():
"--response_format",
response_format,
"--tools_file",
- os.path.join(cwd, "..", "test_models", "tool-definitions", "weather.json"),
+ os.path.join(cwd, "..", "tool-definitions", "weather.json"),
"--tool_call_start",
tool_call_start,
"--tool_call_end",
@@ -151,7 +151,7 @@ def run_tool_calling():
"--response_format",
response_format,
"--tools_file",
- os.path.join(cwd, "..", "test_models", "tool-definitions", "weather.json"),
+ os.path.join(cwd, "..", "tool-definitions", "weather.json"),
"--tool_call_start",
tool_call_start,
"--tool_call_end",
@@ -169,15 +169,15 @@ def run_nemotron_speech():
"""Run Nemotron Speech Streaming ASR E2E test by invoking the nemotron_speech.py example."""
log.debug("Running Nemotron Speech Python E2E Test")
- # Look for nemotron speech model in test_models directory
+ # Look for nemotron speech model in "models" directory
cwd = os.path.dirname(os.path.abspath(__file__))
- model_path = os.path.join(cwd, "..", "test_models", "nemotron-speech-streaming")
+ model_path = os.path.join(cwd, "..", "models", "nemotron-speech-streaming")
if not os.path.exists(model_path):
log.info(f"Nemotron speech model not found at {model_path}, skipping E2E test.")
return
# Look for a test audio file
- audio_path = os.path.join(cwd, "..", "test_models", "audios", "1272-141231-0002.mp3")
+ audio_path = os.path.join(cwd, "..", "audios", "1272-141231-0002.mp3")
if not os.path.exists(audio_path):
log.info(f"Test audio file not found at {audio_path}, skipping E2E test.")
return
@@ -198,13 +198,13 @@ def run_parakeet_tdt():
log.debug("Running Parakeet TDT Python E2E Test")
cwd = os.path.dirname(os.path.abspath(__file__))
- model_path = os.path.join(cwd, "..", "test_models", "parakeet-tdt")
+ model_path = os.path.join(cwd, "..", "models", "parakeet-tdt")
if not os.path.exists(model_path):
log.info(f"Parakeet TDT model not found at {model_path}, skipping E2E test.")
return
for audio_filename in ("jfk.flac", "tedlium_long_120s.flac"):
- audio_path = os.path.join(cwd, "..", "test_models", "audios", audio_filename)
+ audio_path = os.path.join(cwd, "..", "audios", audio_filename)
if not os.path.exists(audio_path):
log.info(f"Test audio file not found at {audio_path}, skipping.")
continue
diff --git a/test/test_utils.h b/test/test_utils.h
index 522668fa05..3117b4a53c 100644
--- a/test/test_utils.h
+++ b/test/test_utils.h
@@ -10,7 +10,7 @@
// Our working directory is generators/build so one up puts us in the root directory:
#ifndef MODEL_PATH
-#define MODEL_PATH "../../test/test_models/"
+#define MODEL_PATH "../../test/models/"
#endif
namespace test_utils {
diff --git a/test/test_models/tool-definitions/filesystem.json b/test/tool-definitions/filesystem.json
similarity index 100%
rename from test/test_models/tool-definitions/filesystem.json
rename to test/tool-definitions/filesystem.json
diff --git a/test/test_models/tool-definitions/ocr.json b/test/tool-definitions/ocr.json
similarity index 100%
rename from test/test_models/tool-definitions/ocr.json
rename to test/tool-definitions/ocr.json
diff --git a/test/test_models/tool-definitions/weather.json b/test/tool-definitions/weather.json
similarity index 100%
rename from test/test_models/tool-definitions/weather.json
rename to test/tool-definitions/weather.json
diff --git a/tools/ci_build/github/apple/objectivec/assemble_objc_pod_package.py b/tools/ci_build/github/apple/objectivec/assemble_objc_pod_package.py
index 572bc4901f..110328f2ba 100755
--- a/tools/ci_build/github/apple/objectivec/assemble_objc_pod_package.py
+++ b/tools/ci_build/github/apple/objectivec/assemble_objc_pod_package.py
@@ -47,7 +47,7 @@
"objectivec/test/*.mm",
],
"test_resource_files": [
- "test/test_models/hf-internal-testing/tiny-random-gpt2-fp32",
+ "test/models/hf-internal-testing/tiny-random-gpt2-fp32",
],
}