zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-01 12:49:50 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

simplify

Browse Source
This commit is contained in:
Alex Barron
2024-10-28 16:03:43 -07:00
parent 48655a7f83
commit 37a3723823
6 changed files with 197 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
llms/mlx_lm/models/llama.py
View File

@@ -7,7 +7,7 @@ import mlx.core as mx
import mlx.nn as nn
from .base import BaseModelArgs, create_attention_mask
from .cache import QuantizedKVCache
from .cache import QuantizedKVCache, quantized_scaled_dot_product_attention
@dataclass
@@ -192,10 +192,10 @@ class Attention(nn.Module):
keys = self.rope(keys)
if isinstance(cache, QuantizedKVCache):
output = mx.fast.quantized_scaled_dot_product_attention(
output = quantized_scaled_dot_product_attention(
queries,
*keys,
*values,
keys,
values,
scale=self.scale,
mask=mask,
group_size=cache.group_size,