Add cuda support when loading local onnx model (#249)

Author: jiafatom

setup.py

@@ -33,8 +33,6 @@
         "invoke>=2.0.0",
         "onnx>=1.11.0",
         "onnxmltools==1.10.0",
-        "onnxruntime >=1.10.1;platform_system=='Linux'",
-        "onnxruntime-directml>=1.19.0;platform_system=='Windows'",
         "torch>=1.12.1",
         "pyyaml>=5.4",
         "typeguard>=2.3.13",
@@ -46,6 +44,10 @@
         "fasteners",
         "GitPython>=3.1.40",
         "psutil",
+        # Conditional dependencies for ONNXRuntime backends
+        "onnxruntime >=1.10.1;platform_system=='Linux' and extra != 'llm-oga-cuda'",
+        "onnxruntime-directml >=1.19.0;platform_system=='Windows' and extra != 'llm-oga-cuda'",
+        "onnxruntime-gpu >=1.19.1;extra == 'llm-oga-cuda'",
     ],
     extras_require={
         "llm": [
@@ -71,6 +73,18 @@
             "fastapi",
             "uvicorn[standard]",
         ],
+        "llm-oga-cuda": [
+            "onnxruntime-genai-cuda==0.4.0",
+            "tqdm",
+            "torch>=2.0.0,<2.4",
+            "transformers<4.45.0",
+            "accelerate",
+            "py-cpuinfo",
+            "sentencepiece",
+            "datasets",
+            "fastapi",
+            "uvicorn[standard]",
+        ],
         "llm-oga-npu": [
             "transformers",
             "torch",

src/turnkeyml/llm/tools/ort_genai/oga.py

@@ -35,7 +35,7 @@
 oga_model_builder_cache_path = "model_builder"
 
 # Mapping from processor to executiion provider, used in pathnames and by model_builder
-execution_providers = {"cpu": "cpu", "npu": "npu", "igpu": "dml"}
+execution_providers = {"cpu": "cpu", "npu": "npu", "igpu": "dml", "cuda": "cuda"}
 
 
 class OrtGenaiTokenizer(TokenizerAdapter):
@@ -248,7 +248,7 @@ def parser(add_help: bool = True) -> argparse.ArgumentParser:
         parser.add_argument(
             "-d",
             "--device",
-            choices=["igpu", "npu", "cpu"],
+            choices=["igpu", "npu", "cpu", "cuda"],
             default="igpu",
             help="Which device to load the model on to (default: igpu)",
         )
@@ -312,6 +312,7 @@ def run(
             "cpu": {"int4": "*/*", "fp32": "*/*"},
             "igpu": {"int4": "*/*", "fp16": "*/*"},
             "npu": {"int4": "amd/**-onnx-ryzen-strix"},
+            "cuda": {"int4": "*/*", "fp16": "*/*"},
         }
         hf_supported = (
             device in hf_supported_models