先把当时的情况摆出来。用 TGI 3.0 在 A10 24GB 上部署 Llama-3.1-8B-Instruct,按文档加了 --speculate 2 做投机解码,并发过 30 就报 CUDA out of memory,去掉这个参数能撑到 60 并发。

这类报错背后常见的原因有几处,逐个排除就行。投机解码要额外加载草稿头,权重与 KV 都多一份。也可能是--max-batch-total-tokens 没跟着下调,KV cache 仍按原预算预留。

还有一种情况,草稿模型与原模型共用一个显存池,峰值翻倍。另外,并发上来后每条序列的 KV 都在涨,投机窗口再乘一层。

动手时先看启动日志给 KV cache 分了多少显存

docker logs tgi 2>&1 | grep -i kv_cache

接着把批总量预算压到一半再开投机

docker run --gpus all -p 8080:80 ghcr.io/huggingface/text-generation-inference:3.0 --model-id meta-llama/Llama-3.1-8B-Instruct --speculate 2 --max-batch-total-tokens 4096

然后用固定并发压一轮看是否还 OOM

python bench_tgi.py --url http://127.0.0.1:8080 --concurrency 30

剩下的就是把并发上限限到实测值

docker exec tgi env | grep MAX_CONCURRENT

恢复的判定其实很直接。30 并发压测全部返回 200,日志里不再出现 OOM,单 token 延迟比不开投机低两成才值得保留这个参数。

补一句。投机解码的收益取决于草稿命中率,短问答基本不赚,长文续写才明显,先压测再决定开不开。监控比日志更早给出信号,事后回看经常能提前半小时发现异常。

作者 9823189