dynamic sliding window

Blaizzy · Jan 18, 2025 · b3aae82 · b3aae82
1 parent 6a98147
commit b3aae82
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/mlx_vlm/models/llava/language.py b/mlx_vlm/models/llava/language.py
@@ -22,6 +22,7 @@ class TextConfig:
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = False
+    sliding_window: int = None
 
     @classmethod
     def from_dict(cls, params):
@@ -51,9 +52,9 @@ def __init__(self, config: TextConfig):
         super().__init__()
 
         dim = config.hidden_size
+        self.config = config
         self.n_heads = n_heads = config.num_attention_heads
         self.n_kv_heads = n_kv_heads = config.num_key_value_heads
-
         self.repeats = n_heads // n_kv_heads
 
         head_dim = config.hidden_size // n_heads
@@ -102,6 +103,12 @@ def __call__(
             queries = self.rope(queries, offset=cache.offset)
             keys = self.rope(keys, offset=cache.offset)
             keys, values = cache.update_and_fetch(keys, values)
+            if self.config.sliding_window:
+                print(self.config.sliding_window)
+                keys = keys[:, :, self.config.sliding_window :, :]
+                values = values[:, :, self.config.sliding_window :, :]
+                if mask is not None:
+                    mask = mask[:, self.config.sliding_window :]
         else:
             queries = self.rope(queries)
             keys = self.rope(keys)

diff --git a/mlx_vlm/models/llava/llava.py b/mlx_vlm/models/llava/llava.py
@@ -164,6 +164,7 @@ def __call__(
             inputs_embeds = self.prefill(
                 inputs_embeds, cache=cache, prefill_step_size=prefill_step_size
             )
+            self.config.text_config.sliding_window = 4096
 
         logits = self.language_model(
             input_ids, cache=cache, inputs_embeds=inputs_embeds