在 M2 Ultra 上用 mlx-lm 加载 Llama-3.1-8B 的 4bit 版本,短提示下速度正常,上下文到 16k 后统一内存占用涨到 40GB,系统开始换页。

一开始怀疑是「对话历史没有截断,每轮都全量重算」,查了一圈,方向不对。

把日志和配置对完,能解释这个现象的基本就这几处。KV cache 没有量化,上下文越长占用越大。少数情况下是批次里把序列 padding 到最长。

不容易想到的是,统一内存被其他进程占用,挤压了可用额度。

上来先打印当前活跃内存与上下文长度

python -c "import mlx.core as mx;print(mx.metal.get_active_memory())"

限制最大上下文并复测内存

python -m mlx_lm.generate --model mlx-community/Llama-3.1-8B-4bit --max-tokens 512 --prompt-cache-size 4096

再往下按 token 截断历史后再请求

python chat.py --max-context 4096 --truncate oldest

剩下的就是关掉其他占内存进程复测

python -c "import psutil;print(psutil.virtual_memory().percent)"

隔了一天再看,现象没变,说明不是偶发。

别只看报错消失,验证要多做一步。16k 上下文下内存占用稳定在 20GB 以内,不再触发换页,生成速度维持每秒 20 token 以上。

这里还有个坑。4bit 只压权重不压缓存,长上下文场景要单独给 KV cache 留预算。改动前把原配置备份一份,出问题能马上回滚,这比任何技巧都实用。

作者 9823189