refactor topk to use mlx (#175)

Blaizzy · Jan 8, 2025 · 1e4579a · 1e4579a
1 parent 71a3f66
commit 1e4579a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/mlx_vlm/models/deepseek_vl_v2/language.py b/mlx_vlm/models/deepseek_vl_v2/language.py
@@ -408,9 +408,9 @@ def __call__(self, x):
 
             # Calculate group scores using top-2 sum per group
             scores_reshaped = scores_for_choice.reshape(bsz * seq_len, self.n_group, -1)
-            k = 2
-            group_scores_topk = mx.sort(scores_reshaped, axis=-1)[..., -k:]
-            group_scores = group_scores_topk.sum(axis=-1)
+
+            # Get top 2 scores per group
+            group_scores = mx.topk(scores_reshaped, 2, axis=-1).sum(axis=-1)
 
             # Get top groups
             k = self.n_group - self.topk_group