diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 6e2786624..0339386e6 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -796,7 +796,7 @@ def test_ep_cuda_passes_through(self, tmp_path): data = json.load(f) provider_opts = data["model"]["decoder"]["session_options"]["provider_options"] assert len(provider_opts) == 1 - assert "CUDAExecutionProvider" in provider_opts[0] + assert "cuda" in provider_opts[0] def test_raises_when_pkg_config_is_none(self, tmp_path): """ValueError is raised when pkg.config is None.""" diff --git a/src/mobius/integrations/ort_genai/ep_config.py b/src/mobius/integrations/ort_genai/ep_config.py index b96a817ec..b5d8066fb 100644 --- a/src/mobius/integrations/ort_genai/ep_config.py +++ b/src/mobius/integrations/ort_genai/ep_config.py @@ -17,12 +17,14 @@ from mobius._execution_providers import ep_registry -# ORT GenAI provider name mapping (internal name → ORT GenAI provider string) +# ORT GenAI provider name mapping (internal name → ORT GenAI provider string). +# GenAI expects short lowercase names (e.g. "cuda", "dml"), not the full +# ORT EP class names (e.g. "CUDAExecutionProvider"). _ORT_PROVIDER_NAMES: dict[str, str] = { - "cpu": "CPUExecutionProvider", - "cuda": "CUDAExecutionProvider", - "dml": "DmlExecutionProvider", - "webgpu": "WebGpuExecutionProvider", + "cpu": "cpu", + "cuda": "cuda", + "dml": "dml", + "webgpu": "webgpu", "trt-rtx": "NvTensorRtRtx", } diff --git a/src/mobius/integrations/ort_genai/ep_config_test.py b/src/mobius/integrations/ort_genai/ep_config_test.py index 7de962830..31a10af52 100644 --- a/src/mobius/integrations/ort_genai/ep_config_test.py +++ b/src/mobius/integrations/ort_genai/ep_config_test.py @@ -20,26 +20,26 @@ def test_cpu_returns_empty(self): def test_cuda_default(self): result = make_provider_options("cuda") assert len(result) == 1 - assert "CUDAExecutionProvider" in result[0] - assert result[0]["CUDAExecutionProvider"]["enable_cuda_graph"] == "0" + assert "cuda" in result[0] + assert result[0]["cuda"]["enable_cuda_graph"] == "0" def test_cuda_with_graph(self): result = make_provider_options("cuda", enable_cuda_graph=True) - assert result[0]["CUDAExecutionProvider"]["enable_cuda_graph"] == "1" + assert result[0]["cuda"]["enable_cuda_graph"] == "1" def test_dml(self): result = make_provider_options("dml") assert len(result) == 1 - assert "DmlExecutionProvider" in result[0] + assert "dml" in result[0] def test_webgpu_default(self): result = make_provider_options("webgpu") - assert result[0]["WebGpuExecutionProvider"]["enableGraphCapture"] == "0" - assert result[0]["WebGpuExecutionProvider"]["validationMode"] == "basic" + assert result[0]["webgpu"]["enableGraphCapture"] == "0" + assert result[0]["webgpu"]["validationMode"] == "basic" def test_webgpu_with_graph(self): result = make_provider_options("webgpu", enable_webgpu_graph=True) - opts = result[0]["WebGpuExecutionProvider"] + opts = result[0]["webgpu"] assert opts["enableGraphCapture"] == "1" assert opts["validationMode"] == "disabled" diff --git a/src/mobius/integrations/ort_genai/genai_config_test.py b/src/mobius/integrations/ort_genai/genai_config_test.py index a96c984ef..033b67e3d 100644 --- a/src/mobius/integrations/ort_genai/genai_config_test.py +++ b/src/mobius/integrations/ort_genai/genai_config_test.py @@ -692,22 +692,22 @@ def test_cpu_has_empty_provider_options(self): assert opts["provider_options"] == [] def test_cuda_has_cuda_provider_options(self): - """CUDA EP produces a provider_options entry for CUDAExecutionProvider.""" + """CUDA EP produces a provider_options entry for cuda.""" from mobius.integrations.ort_genai.genai_config import _make_session_options opts = _make_session_options("cuda") assert opts["log_id"] == "onnxruntime-genai" assert len(opts["provider_options"]) == 1 - assert "CUDAExecutionProvider" in opts["provider_options"][0] + assert "cuda" in opts["provider_options"][0] def test_dml_has_dml_provider_options(self): - """DML EP produces a provider_options entry for DmlExecutionProvider.""" + """DML EP produces a provider_options entry for dml.""" from mobius.integrations.ort_genai.genai_config import _make_session_options opts = _make_session_options("dml") assert opts["log_id"] == "onnxruntime-genai" assert len(opts["provider_options"]) == 1 - assert "DmlExecutionProvider" in opts["provider_options"][0] + assert "dml" in opts["provider_options"][0] class TestGenaiConfigGeneratorEp: @@ -735,7 +735,7 @@ def test_cuda_ep_decoder_has_cuda_provider_options(self): config = self._gen("cuda").generate() opts = config["model"]["decoder"]["session_options"]["provider_options"] assert len(opts) == 1 - assert "CUDAExecutionProvider" in opts[0] + assert "cuda" in opts[0] def test_cuda_ep_all_blocks_have_cuda_session_options(self): """CUDA EP applied to all 4 session blocks (decoder, vision, embedding, audio).""" @@ -761,6 +761,4 @@ def test_cuda_ep_all_blocks_have_cuda_session_options(self): session_opts = config["model"][block]["session_options"] provider_options = session_opts["provider_options"] assert len(provider_options) == 1, f"{block} missing CUDA provider options" - assert "CUDAExecutionProvider" in provider_options[0], ( - f"{block} has wrong EP in provider_options" - ) + assert "cuda" in provider_options[0], f"{block} has wrong EP in provider_options"