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

in place

Browse Source
This commit is contained in:
Alex Barron
2024-10-31 12:22:36 -07:00
parent 1d53354b51
commit 8444ff0f6a
2 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
llms/mlx_lm/utils.py
View File

@@ -165,10 +165,10 @@ def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_
and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
and prompt_cache[0].offset > quantized_kv_start
):
return [
c.to_quantized(group_size=kv_group_size, bits=kv_bits) for c in prompt_cache
]
return prompt_cache
for i in range(len(prompt_cache)):
prompt_cache[i] = prompt_cache[i].to_quantized(
group_size=kv_group_size, bits=kv_bits
)
def generate_step(
@@ -290,7 +290,7 @@ def generate_step(
for processor in logits_processor:
logits = processor(tokens, logits)
prompt_cache = maybe_quantize_kv_cache(
maybe_quantize_kv_cache(
prompt_cache, quantized_kv_start, kv_group_size, kv_bits
)