todo for speculative model

57215041a0edd70307736fa3a8246b92e46a2245 · 2024-09-01 23:16:10 +0100 · Alex Cheema

Files touched

M exo/inference/mlx/sharded_model.py

Diff

commit 57215041a0edd70307736fa3a8246b92e46a2245
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Sep 1 23:16:10 2024 +0100

    todo for speculative model
---
 exo/inference/mlx/sharded_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exo/inference/mlx/sharded_model.py b/exo/inference/mlx/sharded_model.py
index c4570fbf..46f06b3d 100644
--- a/exo/inference/mlx/sharded_model.py
+++ b/exo/inference/mlx/sharded_model.py
@@ -8,7 +8,7 @@ from mlx_lm.sample_utils import top_p_sampling
 
 from ..shard import Shard
 
-
+# TODO: support a speculative model so we can parallelise compute across devices
 class StatefulShardedModel:
   def __init__(self, shard: Shard, model: nn.Module, max_kv_size: int = 1024, max_caches: int = 2):
     self.shard = shard

← dc3b2bde use NousResearch/Meta-Llama-3.1-70B-Instruct as tinygrad lla · back to Exo · cleaner discovery baf6efd3 →