在 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 留预算。改动前把原配置备份一份,出问题能马上回滚,这比任何技巧都实用。
