Add Llama.cpp Support

exo/inference/inference_engine.py

@@ -1,3 +1,4 @@
+from llama_cpp.llama import Llama
 import numpy as np
 import os
 
@@ -27,5 +28,9 @@ def get_inference_engine(inference_engine_name: str, shard_downloader: 'ShardDow
     tinygrad.helpers.DEBUG.value = int(os.getenv("TINYGRAD_DEBUG", default="0"))
 
     return TinygradDynamicShardInferenceEngine(shard_downloader)
+  elif inference_engine_name == "llama_cpp":
+    from exo.inference.llama_cpp.inference import LlamaCppDynamicShardEngine
+    return LlamaCppDynamicShardEngine(shard_downloader)
+
   else:
     raise ValueError(f"Inference engine {inference_engine_name} not supported")

exo/inference/llama_cpp/__init__.py

@@ -0,0 +1 @@
+

exo/inference/llama_cpp/inference.py

@@ -0,0 +1,24 @@
+import numpy as np
+from typing import Optional, Tuple
+from llama_cpp import Llama
+from exo.inference.shard import Shard
+from exo.inference.inference_engine import InferenceEngine
+from exo.download.shard_download import ShardDownloader
+
+class LlamaCppDynamicShardEngine(InferenceEngine):
+    def __init__(self, shard_downloader: ShardDownloader):
+        self.shard = None
+        self.shard_downloader = shard_downloader
+
+    async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, image_str: Optional[str] = None, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
+        await self.ensure_shard(shard)
+
+    async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> Tuple[np.ndarray, str, bool]:
+        await self.ensure_shard(shard)
+
+    async def ensure_shard(self, shard: Shard):
+        if self.shard == shard:
+            return
+
+        model_path = await self.shard_downloader.ensure_shard(shard)
+        self.shard = shard

exo/models.py

@@ -5,6 +5,8 @@
   "llama-3.1-8b": {
     "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=32),
     "TinygradDynamicShardInferenceEngine": Shard(model_id="mlabonne/Meta-Llama-3.1-8B-Instruct-abliterated", start_layer=0, end_layer=0, n_layers=32),
+    "LlamaCppDynamicShardEngine": Shard(model_id="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF", start_layer=0, end_layer=0, n_layers=32)
+
   },
   "llama-3.1-70b": {
     "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=80),
@@ -17,7 +19,7 @@
   },
   "llama-3-70b": {
     "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3-70B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=80),
-    "TinygradDynamicShardInferenceEngine": Shard(model_id="TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-70B-R", start_layer=0, end_layer=0, n_layers=80),
+    "TinygradDynamicShardInferenceEngine": Shard(model_id="", start_layer=0, end_layer=0, n_layers=80),
   },
   ### mistral
   "mistral-nemo": {"MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Mistral-Nemo-Instruct-2407-4bit", start_layer=0, end_layer=0, n_layers=40),},

setup.py

@@ -29,6 +29,7 @@
   "tqdm==4.66.4",
   "transformers==4.43.3",
   "uuid==1.30",
+  "llama-cpp-python==0.2.89",
   "tinygrad @ git+https://github.com/tinygrad/tinygrad.git@639af3f823cf242a1945dc24183e52a9df0af2b7",
 ]