MLX-LM 長文本聊天實戰：Context、KV Cache 與記憶體取捨

一. 前言：模型支援 128K，不代表你的 Mac 喜歡 128K

#

本地 LLM 剛跑起來時，大家通常只在意兩件事：

• 模型能不能載入；
• 每秒可以吐幾個 token。 等你把一份長文件、幾十輪聊天紀錄，或整包程式碼塞進 prompt，問題就變了。 第一個 token 等很久、記憶體一路上升、系統開始 swap，最後模型還忘記開頭說過什麼。 模型卡上寫著很大的 context window，也救不了錯誤的工作流。 這篇不再介紹 MLX-LM 的基本載入與聊天介面。 如果你第一次使用，先看MLX-LM 本機模型入門。 今天拍拍君要專心處理長文本推論的五個問題：

1. prompt 到底用了多少 token；
2. prefill 和 generation 為什麼是兩種不同成本；
3. KV cache 如何隨序列長度長大；
4. 哪些旋鈕在省峰值記憶體，哪些是在丟掉歷史；
5. 怎麼用同一套案例比較速度、記憶體和答案品質。 先講結論：長 context 沒有免費午餐。 你是在「保留多少歷史、多久拿到第一個 token、占多少 unified memory」之間做選擇。

二. 安裝與建立測試環境

#

建立一個獨立專案，避免拿日常聊天環境直接亂試：

uv init mlx-long-context-lab
cd mlx-long-context-lab
uv add mlx-lm
mkdir -p prompts results

先確認目前 CLI 真的有哪些選項：

uv run mlx_lm.generate --help
uv run mlx_lm.cache_prompt --help

MLX-LM 更新速度很快。 文章中的參數是目前版本提供的介面，但你的模型可能有自訂 cache 實作。 所以不要只確認指令接受參數，還要用後面的量測確認它真的產生預期效果。 先指定一個你實際跑得動的 instruct model：

export MLX_MODEL_ID="mlx-community/Qwen3-4B-Instruct-2507-4bit"

模型名稱只是範例，不是硬性推薦。 請確認模型頁面的授權、chat template、context 長度，以及你的機器是否有足夠記憶體。

三. 先分清楚三種「長」

#

長文本推論常把三件事混在一起：

四. 建立 Token 預算，不要用字數猜

#

中文一個字不固定等於一個 token，程式碼、JSON、URL 更難靠肉眼估算。 最可靠的方法，是用該模型自己的 tokenizer。 建立 token_budget.py：

from pathlib import Path
from mlx_lm import load
MODEL = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
MODEL_CONTEXT_LIMIT = 262_144
WORKING_CONTEXT_LIMIT = 32_768
OUTPUT_RESERVE = 1_024
SAFETY_MARGIN = 512
model, tokenizer = load(MODEL)
document = Path("prompts/document.md").read_text(encoding="utf-8")
messages = [
    {"role": "system", "content": "只根據文件回答，使用繁體中文。"},
    {"role": "user", "content": f"文件：\n{document}\n\n請整理五個重點。"},
]
prompt = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=False,
)
prompt_tokens = len(tokenizer.encode(prompt))
usable = min(MODEL_CONTEXT_LIMIT, WORKING_CONTEXT_LIMIT) - OUTPUT_RESERVE - SAFETY_MARGIN
print(f"prompt tokens: {prompt_tokens:,}")
print(f"usable input budget: {usable:,}")
print(f"remaining input budget: {usable - prompt_tokens:,}")
if prompt_tokens > usable:
    raise SystemExit("prompt 超過預算：請切片、檢索或摘要")

這裡故意同時保留輸出空間和安全邊界。 不要把 context window 塞到剛好等於上限，再期待模型生成完整答案。 MODEL_CONTEXT_LIMIT 必須來自模型設定或模型卡；WORKING_CONTEXT_LIMIT 則是你依延遲與記憶體訂出的保守上限。 換模型時一起檢查，不要把它們寫成全世界共用的魔法常數。

五. Context Window 是上限，不是品質保證

#

模型宣稱支援長 context，只表示架構和設定允許處理那麼長的序列。 它不保證：

• 遠距離資訊一定找得準；
• 中間段落不會被忽略；
• 相同問題在 4K 和 64K prompt 裡品質一樣；
• 你的本機可以在合理時間與記憶體內完成。 因此，長文本測試不能只問「有沒有成功輸出」。 至少準備三種題目：

1. 答案在文件開頭；
2. 答案在文件中間；
3. 答案要整合開頭與結尾。 再放入一題文件裡根本沒有答案的問題，檢查模型會不會硬掰。 這比丟一篇長文叫它「摘要一下」更容易發現退化。

六. Prefill 與 Decode：兩張不同的帳單

#

推論可以粗略拆成兩段：

prompt tokens -> prefill -> 第一個輸出 token -> decode -> 完整答案

prefill 一次讀入既有 prompt，長文件通常把時間花在這裡。 decode 則逐 token 生成答案，速度常用 generation tokens/s 表示。 因此你可能遇到：

• prompt throughput 很高，但文件太長，第一個 token 還是慢；
• 第一個 token 很快，但長答案生成很久；
• 重複問同一份文件時，每次又把固定前綴 prefill 一遍。 MLX-LM 的 verbose 統計會分開顯示 prompt 與 generation throughput， Python 回傳資訊也包含 prompt tokens、generation tokens、速度與 peak memory。 請分開記錄，不要只留一個「總共跑了幾秒」。

七. 用 Prefill Step 壓低讀入時的峰值

#

長 prompt 會分段讀入。 目前預設 prefill step 是 2048 tokens；較小的 step 通常降低 prefill 峰值記憶體，代價是處理速度可能下降。 先跑基準：

uv run mlx_lm.generate \
  --model "$MLX_MODEL_ID" \
  --prompt "$(< prompts/long-question.txt)" \
  --max-tokens 256 \
  --prefill-step-size 2048 \
  2>&1 | tee results/prefill-2048.txt

再比較 1024 與 512：

for step in 1024 512; do
  uv run mlx_lm.generate \
    --model "$MLX_MODEL_ID" \
    --prompt "$(< prompts/long-question.txt)" \
    --max-tokens 256 \
    --prefill-step-size "$step" \
    2>&1 | tee "results/prefill-${step}.txt"
done

固定模型、prompt、輸出上限與背景程式。 只改一個參數，才知道差異從哪裡來。 如果 2048 不會造成記憶體壓力，沒有必要為了「數字比較小」改成 512。 這個旋鈕的目標是穿過峰值，不是競賽誰切得最碎。

八. KV Cache 到底存了什麼？

#

Transformer 在生成下一個 token 時，需要回頭注意前面的 token。 KV cache 保存各層已計算的 key 和 value，避免每一步把完整歷史全部重算。 好處是 decode 快很多；代價是 cache 會跟序列長度一起成長。 對常見 GQA 模型，可以用這個粗估式建立直覺：

KV bytes ≈ 2 × layers × kv_heads × head_dim × tokens × bytes_per_value

前面的 2 代表 key 與 value。 注意這是估算，不包含 allocator、對齊、量化 scale、其他 activation 與程式本身。 建立 estimate_kv.py：

def kv_gib(
    *,
    layers: int,
    kv_heads: int,
    head_dim: int,
    tokens: int,
    bits: int = 16,
) -> float:
    values = 2 * layers * kv_heads * head_dim * tokens
    return values * (bits / 8) / (1024**3)
for length in (4_096, 8_192, 16_384, 32_768):
    size = kv_gib(
        layers=36,
        kv_heads=8,
        head_dim=128,
        tokens=length,
    )
    print(f"{length:>6,} tokens -> 約 {size:.2f} GiB")

請從模型的 config.json 取得 layer、KV heads 與 head dimension。 有些模型使用 sliding-window、hybrid attention 或自訂 cache，不能套公式後就當成實測值。 公式負責幫你預判，峰值記憶體仍以實際執行為準。

九. Rotating KV Cache：固定記憶體，交換歷史

#

--max-kv-size 會要求使用固定大小的 rotating KV cache。 舊項目會被覆寫，而最前面的少量 sink tokens 會保留。 例如只保留 4096 個 cache positions：

uv run mlx_lm.generate \
  --model "$MLX_MODEL_ID" \
  --prompt "$(< prompts/long-question.txt)" \
  --max-tokens 512 \
  --max-kv-size 4096

較小的上限會省記憶體，但模型能直接注意到的舊內容也變少，品質可能下降。 這不是無損壓縮，也不是把 32K context 免費塞進 4K 空間。 它比較適合：

• 長時間續寫，近期內容比遠古細節重要；
• 聊天內容允許舊輪次逐漸淡出；
• 你已用摘要或檢索保存真正重要的資訊。 它不適合需要精確引用文件開頭的任務。 而且某些模型會提供自己的 cache 實作；請確認參數是否真的生效。

十. KV Cache Quantization：省的是動態 Cache，不是模型權重

#

模型權重量化與 KV cache 量化是兩件事。 前者已在MLX-LM 模型轉換與量化完整討論； 這裡只處理長序列產生的動態 cache。 可以單獨比較 8-bit KV cache：

uv run mlx_lm.generate \
  --model "$MLX_MODEL_ID" \
  --prompt "$(< prompts/long-question.txt)" \
  --max-tokens 512 \
  --kv-bits 8 \
  --kv-group-size 64 \
  --quantized-kv-start 4096

--quantized-kv-start 讓前段先使用原精度 cache，到指定位置後才開始量化。 這可以避免短 prompt 也支付不必要的量化成本。 測試時至少比較：

• 不量化；
• 8-bit KV；
• 你的真實長度，而不是只有 200 tokens 的玩具 prompt；
• 開頭、中間與跨段問題的答案正確率。 不要一開始就把 --max-kv-size 和 --kv-bits 全部疊上去。 不同 cache 類型與模型架構的支援並不一致；先分開確認，再測組合。

十一. Prompt Cache：同一份長前綴，不必每次重讀

#

如果你要對同一份規格書問十個不同問題，最浪費的做法是每次都重新 prefill 全文。 MLX-LM 可以把預先計算的 prompt cache 存成檔案。 先建立穩定前綴：

uv run mlx_lm.cache_prompt \
  --model "$MLX_MODEL_ID" \
  --prompt - \
  --prompt-cache-file prompts/manual-cache.safetensors \
  < prompts/manual.txt

再針對這個前綴提問：

uv run mlx_lm.generate \
  --prompt-cache-file prompts/manual-cache.safetensors \
  --prompt $'\n請列出文件中的三個失敗處理原則。' \
  --max-tokens 300

使用 cache 時，新 prompt 會接在 cached prefix 後面。 模型資訊也會跟 cache 一起保存，因此命令不必重複指定 model。 這和MLX-LM API Server的自動 prefix cache 不同： 那篇處理多 request 的 HTTP 服務；這裡是你明確建立、保存與重用一份長前綴。

十二. Prompt Cache 的正確邊界

#

適合 cache 的內容：

• 固定 system instruction；
• 同一份長文件；
• 不常改的工具說明；
• 多個問題共用的背景資料。 不適合混進 cache 的內容：
• 每次都變的日期與 request ID；
• 不同使用者的私人資料；
• 尚未確認版本的臨時文件；
• 你無法追蹤來源的對話垃圾。 cache 檔案本身可能反映敏感 prompt 的中間狀態，請把它視為本機資料產物管理。 不要隨手 commit、上傳或跨使用者共用。 當模型、tokenizer、chat template 或固定文件改變時，重新建立 cache。 最簡單的做法是把來源檔 checksum 寫進旁邊的 metadata 檔名或 manifest。

十三. 多輪聊天：不要只做「留下最近 N 輪」

#

單純保留最近六輪很可靠，但長期對話容易丟掉早期決策。 比較實用的 prompt 分層是：

固定規則
長期摘要
與本題相關的檢索片段
最近幾輪原始對話
這次的新問題

每一層都要有 token 預算。 例如：

BUDGET = {
    "system": 800,
    "summary": 1_500,
    "retrieval": 6_000,
    "recent_chat": 4_000,
    "answer": 1_000,
    "margin": 500,
}
total = sum(BUDGET.values())
print(f"planned context: {total:,} tokens")

當 recent chat 超額，就把較舊內容合併進摘要。 當 retrieval 超額，就提高檢索門檻或減少片段，而不是直接截掉答案的一半。 需要從大量文件找相關內容時，先用MLX Embeddings 語意搜尋取回片段， 通常比把整個資料夾硬塞進 context 更便宜、更可控。

十四. 建立可重現的長文本 Benchmark

#

至少準備同一份文件的四種長度：

prompts/doc-4k.txt
prompts/doc-8k.txt
prompts/doc-16k.txt
prompts/doc-32k.txt

每份都放相同類型的針對性問題，並保存標準答案。 測試矩陣可以是：

• prompt tokens；
• prompt tokens/s；
• time to first token；
• generation tokens/s；
• peak memory；
• 答案是否命中標準；
• 系統是否開始 swap。 不要把不同 prompt 長度的數字塞在同一欄直接比較。 長度本身就是最重要的自變數。

十五. macOS 上的實際觀察方式

#

CLI 的 peak memory 是第一手指標之一，也可以用系統工具補充：

/usr/bin/time -l uv run mlx_lm.generate \
  --model "$MLX_MODEL_ID" \
  --prompt "$(< prompts/long-question.txt)" \
  --max-tokens 256 \
  > results/baseline.txt 2>&1

同時開啟 Activity Monitor，觀察 Memory Pressure 與 Swap Used。 unified memory 很方便，但 CPU、GPU、模型權重、KV cache 和其他 App 都在用同一池資源。 「程式沒有 OOM」不等於系統狀態健康。 如果 swap 持續增加、整台機器明顯卡頓，先降低模型大小或 context 工作量。 不要用更激進的 cache 技巧把一個根本塞不下的配置硬撐成可啟動。

十六. 一套務實的調校順序

#

拍拍君建議照這個順序做：

1. 用 tokenizer 算出真實輸入；
2. 保留輸出與安全 margin；
3. 建立不加任何 cache 限制的品質基準；
4. 若 prefill 峰值過高，降低 prefill step；
5. 若多次使用相同長前綴，建立 prompt cache；
6. 若 KV 成長仍是主要壓力，單獨測 8-bit KV；
7. 若任務容許遺忘遠端歷史，再測 rotating cache；
8. 最後才組合參數，並重跑品質題。 這個順序會先拿到可信答案，再逐步降低成本。 如果一開始同時改五個旋鈕，就算速度變快，你也不知道是誰幫了忙、誰偷走了品質。

十七. 常見誤區

#

1. 模型標示 128K，所以每次都塞滿

#

上限不是建議工作長度。 先找任務真正需要的證據，再決定 context。

2. 降低 max tokens 就能解決長 prompt

#

它只限制輸出。 長輸入的 prefill 與既有 KV 成本仍然存在。

3. Prompt cache 會讓所有問題更快

#

只有相同 token prefix 能重用。 在前綴最前面加入時間戳，就會破壞大段命中。

4. Rotating cache 是無損長期記憶

#

它會覆寫舊 KV。 重要資訊應由摘要、資料庫或檢索保存。

5. KV 量化一定沒有品質損失

#

量化改變數值表示。 請用遠距引用與跨段整合題實測，不要只看聊天語氣正常。

6. 只看 tokens/s

#

使用者先感受到的是第一個 token 等多久，系統管理者先看到的可能是 swap。 速度、記憶體和正確率要一起看。

結語：保留真正需要的 Context

#

長文本推論的核心，不是追求最大的 context 數字。 而是讓每一段資訊都有清楚用途，讓每一份記憶體成本都能被量測。 先用 tokenizer 建立預算，再區分 prefill 和 decode。 prefill step 負責降低讀入峰值；prompt cache 負責重用穩定前綴；KV quantization 負責壓低動態 cache；rotating cache 則用有限歷史換固定上限。 它們不是四個同義的「省記憶體按鈕」。 每個工具都解不同問題，也都有自己的品質代價。 當文件真的很大，勇敢切片、摘要或檢索。 能少放一萬個無關 token，往往比任何神奇參數都有效。 把 benchmark 留下來，換模型、升級 MLX-LM 或修改 prompt 時再跑一次。 這樣你的本地長文本工作流不是「今天剛好能動」，而是真的可維護。🔭

延伸閱讀

#

• MLX-LM 實戰：在 Apple Silicon 上跑本地模型推論
• MLX-LM 本機 API Server：OpenAI 相容介面、Prompt Cache 與 Tool Calling
• MLX-LM 模型轉換與量化：4/8-bit、Mixed Quant 與品質評測
• MLX-LM 官方文件：Long Prompts and Generations