跳转到内容
123xiao | 无名键客

《大模型推理优化实战:从 KV Cache、量化到并发调度的性能提升路径》

字数: 0 阅读时长: 1 分钟

大模型推理优化实战:从 KV Cache、量化到并发调度的性能提升路径

做大模型推理优化,最容易遇到的情况是:模型能跑,但跑得不够快;吞吐还能接受,但延迟不稳定;显存勉强够,但并发一上来就炸。

很多团队一开始会直觉地“换更大的卡”或者“多加机器”,但实际工程里,推理性能往往不是一个点的问题,而是一条链路的问题:单 token 生成速度、显存占用、请求并发、批处理策略、调度方式,这些因素会互相拉扯。

这篇文章我想用一条比较实战的路径,把这几个关键优化点串起来:

  1. 先理解为什么推理会慢
  2. 搞清楚 KV Cache 到底省了什么
  3. 再看量化如何换取吞吐和显存空间
  4. 最后落到并发调度,解决“单请求快”和“系统整体快”的矛盾

文章会尽量避免“只讲原理不落地”,我会给出可运行的代码示例,并补上我自己在工程中比较常见的排查思路。


背景与问题

在大模型推理场景里,常见指标通常有这几类:

  • TTFT(Time To First Token):首 token 延迟
  • TPOT(Time Per Output Token):每个输出 token 的平均耗时
  • 吞吐量:单位时间处理多少 token / 请求
  • 显存占用:决定能否撑住更长上下文和更多并发
  • 稳定性:P95/P99 延迟是否抖动严重

如果把一次生成任务拆开看,大致可以分成两个阶段:

  1. Prefill 阶段:把整段输入 prompt 喂给模型
  2. Decode 阶段:每次生成一个 token,再继续下一个 token

对于长上下文任务,Prefill 很重;对于长输出任务,Decode 会成为主要瓶颈。

一个容易被忽略的事实

很多人只盯着“模型参数量”,但实际线上性能往往更受下面几件事影响:

  • 上下文长度是否很长
  • 是否启用 KV Cache
  • 是否做了权重量化
  • 调度器能不能把多个请求拼起来跑
  • 是否因为显存碎片、批处理不稳导致抖动

也就是说,推理优化不是单点优化,而是系统优化。


前置知识与环境准备

本文示例基于 Python,默认环境如下:

  • Python 3.10+
  • PyTorch 2.1+
  • transformers 4.38+
  • accelerate
  • bitsandbytes(用于 8bit/4bit 量化)
  • 一张支持 CUDA 的 GPU

安装示例:

pip install torch transformers accelerate bitsandbytes

如果你只是想先跑通逻辑,没有 GPU,也可以只看文中的简化版调度代码,核心思路是通用的。


先建立整体认知:推理链路长什么样

下面这张图可以帮助你先把优化对象放到一张图里看。

flowchart LR
    A[用户请求] --> B[Tokenizer]
    B --> C[Prefill<br/>计算整段输入]
    C --> D[KV Cache建立]
    D --> E[Decode循环<br/>逐token生成]
    E --> F[采样/贪心解码]
    F --> G[返回结果]

    C -.受上下文长度影响.-> H[首token延迟]
    D -.受显存影响.-> I[并发上限]
    E -.受单步算力与调度影响.-> J[吞吐与尾延迟]

你可以把三类优化对应到三个位置:

  • KV Cache:优化 Decode 阶段重复计算
  • 量化:减少模型权重占用,缓解显存和带宽压力
  • 并发调度:提升 GPU 利用率,平衡延迟和吞吐

核心原理

1. KV Cache:为什么它能显著提速

Transformer 自回归生成时,每生成一个 token,都要做注意力计算。
如果不做缓存,第 t 步生成时,会把前面 1...t-1 的内容反复参与大量计算。

KV Cache 的做法是:

  • 在前面步骤里,把每层 attention 的 Key/Value 结果缓存起来
  • 下一步生成新 token 时,只计算这个新 token 对历史 KV 的注意力
  • 避免重复构造历史 token 的 K/V

直觉理解

不使用 KV Cache:

  • 第 1 步看 1 个 token
  • 第 2 步重算前 2 个 token
  • 第 3 步重算前 3 个 token
  • 越往后越贵

使用 KV Cache:

  • 历史 token 的 K/V 已经算好
  • 每一步只补新 token 的增量计算

KV Cache 的代价

它不是白赚的,代价是显存

缓存大小大致与这些因素有关:

  • 层数
  • 注意力头数 / hidden size
  • 序列长度
  • batch size
  • 数据类型(FP16、BF16、INT8 等)

简单理解:上下文越长、并发越高,KV Cache 越大。


2. 量化:为什么量化后常常更快

量化通常有两种主要目标:

  1. 降显存
  2. 提吞吐

最常见的是权重量化,比如:

  • FP16 / BF16
  • INT8
  • 4-bit(如 NF4)

它快在哪里

对于推理来说,量化的收益通常来自:

  • 权重更小,显存占用下降
  • 从显存搬运到计算单元的数据更少
  • 更容易放下更大 batch 或更长上下文
  • 给 KV Cache 和并发调度腾出空间

但量化不总是“越低越快”

我自己踩过一个坑:
有些卡上 4-bit 的确更省显存,但由于内核支持、反量化开销、算子实现差异,实际延迟不一定比 8-bit 更好

所以量化要看三件事:

  • 你的 GPU 架构
  • 你的推理框架
  • 你的目标是低延迟还是高吞吐

工程上很常见的选择是:

  • 先尝试 8-bit:兼顾稳定与收益
  • 显存特别紧张再试 4-bit
  • 对输出质量敏感时做回归评估

3. 并发调度:单请求快,不代表系统快

如果 GPU 上同时来了很多请求,最糟糕的情况是:

  • 每个请求都各自单独跑
  • batch 很小
  • GPU 时而空闲,时而拥堵
  • 尾延迟很高

这时就需要调度器

调度器的核心工作通常有两类:

  • 动态批处理(dynamic batching)
  • 连续批处理(continuous batching)

动态批处理

在一个短时间窗口里收集请求,拼成 batch 一起跑。

优点:

  • 容易实现
  • 提升吞吐明显

缺点:

  • 请求要等待凑批
  • 可能拉高 TTFT

连续批处理

已经在跑的 batch 中,随着某些请求结束,新请求可以补进来,不必等整批结束。

优点:

  • 更充分利用 GPU
  • 更适合长度差异大的生成任务

缺点:

  • 实现复杂
  • KV Cache 管理更麻烦

下面这张时序图可以看出两者差异。

sequenceDiagram
    participant U1 as 请求A
    participant U2 as 请求B
    participant U3 as 请求C
    participant S as 调度器
    participant G as GPU

    U1->>S: 到达
    U2->>S: 到达
    S->>G: 动态批处理 A+B
    G-->>S: 生成中
    U3->>S: 到达
    Note over S,G: 连续批处理中,C可在空槽位补入
    S->>G: 补入C继续decode
    G-->>S: 返回A/B/C结果

4. 三者之间的关系:不要孤立看优化项

真正实战时,这三个优化点往往不是独立的。

flowchart TD
    A[量化降低权重显存] --> B[腾出更多显存]
    B --> C[允许更大KV Cache]
    B --> D[允许更高并发]
    C --> E[更长上下文/更多会话]
    D --> F[连续批处理更有效]
    F --> G[总体吞吐提升]
    C --> H[Decode阶段更平稳]
    H --> G

一句话概括就是:

量化提供空间,KV Cache提升单请求生成效率,并发调度把这些空间转化为系统吞吐。


实战代码(可运行)

下面我们做两个层次的实战:

  1. 用 Hugging Face 跑一个可执行的 KV Cache + 量化示例
  2. 用纯 Python 写一个简化版调度器,帮助你理解并发调度

实战一:启用 KV Cache 与量化推理

下面示例使用 transformersbitsandbytes
为了保证大多数环境可改造,这里写成比较通用的形式。

提示:模型名请替换成你有权限访问、且适合本机显存的因果语言模型。

import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

MODEL_NAME = "facebook/opt-1.3b"  # 可替换为其他 causal LM

def load_model(load_in_8bit=True):
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)

    quant_config = None
    if load_in_8bit:
        quant_config = BitsAndBytesConfig(load_in_8bit=True)

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="auto",
        quantization_config=quant_config,
        torch_dtype=torch.float16,
    )
    model.eval()

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    return tokenizer, model

def generate_once(tokenizer, model, prompt, max_new_tokens=64, use_cache=True):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    torch.cuda.synchronize()
    start = time.perf_counter()

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=use_cache,
            pad_token_id=tokenizer.eos_token_id,
        )

    torch.cuda.synchronize()
    end = time.perf_counter()

    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return text, end - start

def main():
    tokenizer, model = load_model(load_in_8bit=True)

    prompt = (
        "请用简洁方式解释什么是大模型推理中的KV Cache,"
        "并给出它在长文本生成中的作用。"
    )

    text_cache, t_cache = generate_once(tokenizer, model, prompt, use_cache=True)
    text_no_cache, t_no_cache = generate_once(tokenizer, model, prompt, use_cache=False)

    print("=" * 80)
    print("启用 KV Cache 耗时:", round(t_cache, 3), "")
    print("关闭 KV Cache 耗时:", round(t_no_cache, 3), "")
    print("=" * 80)
    print(text_cache[:500])

if __name__ == "__main__":
    main()

运行后你应该关注什么

不要只看“能不能出结果”,重点看:

  • use_cache=TrueFalse 的耗时差异
  • 长 prompt、长输出下差异是否更明显
  • 8bit 模型是否明显降低了显存压力
  • 不同 batch size 下显存和延迟的变化

一个更靠谱的测试姿势

做 benchmark 时,不要只跑一次。
第一次运行通常会受到以下因素影响:

  • CUDA kernel 初始化
  • 显存分配
  • tokenizer / 权重缓存
  • 编译或 lazy loading

可以这样写一个更稳定的测试:

import time
import torch
from statistics import mean

def benchmark_generate(tokenizer, model, prompt, repeats=5, warmup=2, use_cache=True):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # warmup
    for _ in range(warmup):
        with torch.no_grad():
            _ = model.generate(
                **inputs,
                max_new_tokens=64,
                do_sample=False,
                use_cache=use_cache,
                pad_token_id=tokenizer.eos_token_id,
            )

    latencies = []
    for _ in range(repeats):
        torch.cuda.synchronize()
        start = time.perf_counter()

        with torch.no_grad():
            _ = model.generate(
                **inputs,
                max_new_tokens=64,
                do_sample=False,
                use_cache=use_cache,
                pad_token_id=tokenizer.eos_token_id,
            )

        torch.cuda.synchronize()
        latencies.append(time.perf_counter() - start)

    return {
        "avg_sec": round(mean(latencies), 4),
        "runs": [round(x, 4) for x in latencies]
    }

实战二:查看显存占用,验证量化收益

很多优化讨论都停留在“理论上”,但工程上一定要做量化前后的显存对比。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

MODEL_NAME = "facebook/opt-1.3b"

def print_cuda_memory(tag=""):
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"[{tag}] allocated={allocated:.2f}MB reserved={reserved:.2f}MB")

def load_fp16():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float16,
        device_map="auto",
    )
    return tokenizer, model

def load_int8():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    quant_config = BitsAndBytesConfig(load_in_8bit=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=quant_config,
        device_map="auto",
        torch_dtype=torch.float16,
    )
    return tokenizer, model

def main():
    torch.cuda.empty_cache()
    print_cuda_memory("start")

    tokenizer, model = load_int8()
    print_cuda_memory("after_int8_load")

    prompt = "请解释量化推理为什么能提升吞吐。"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        _ = model.generate(**inputs, max_new_tokens=32, use_cache=True)

    print_cuda_memory("after_generate")

if __name__ == "__main__":
    main()

如果你想做更严谨的 A/B 对比:

  • 进程隔离运行 FP16 和 INT8
  • 记录加载后显存
  • 记录生成后的峰值显存
  • 记录平均延迟与 P95 延迟

实战三:写一个简化版并发调度器

真实生产环境里,你可能会使用 vLLM、TGI、TensorRT-LLM 一类框架。
但如果你不先理解调度本质,很多配置项只会“背参数”而不会调。

下面这个例子不用 GPU,也能帮助你理解:

  • 请求不断到来
  • 调度器按小窗口凑批
  • batch 内任务一起执行
  • 不同请求长度不同,产生排队和资源占用差异
import asyncio
import random
import time
from dataclasses import dataclass, field
from typing import List

@dataclass
class Request:
    req_id: int
    prompt_len: int
    output_len: int
    created_at: float = field(default_factory=time.time)

class DynamicBatchScheduler:
    def __init__(self, max_batch_size=4, batch_wait_ms=50):
        self.queue = asyncio.Queue()
        self.max_batch_size = max_batch_size
        self.batch_wait_ms = batch_wait_ms / 1000.0
        self.running = True

    async def submit(self, req: Request):
        await self.queue.put(req)

    async def fetch_batch(self) -> List[Request]:
        batch = []
        first = await self.queue.get()
        batch.append(first)

        start = time.time()
        while len(batch) < self.max_batch_size:
            remain = self.batch_wait_ms - (time.time() - start)
            if remain <= 0:
                break
            try:
                req = await asyncio.wait_for(self.queue.get(), timeout=remain)
                batch.append(req)
            except asyncio.TimeoutError:
                break
        return batch

    async def run(self):
        while self.running:
            batch = await self.fetch_batch()
            await self.process_batch(batch)

    async def process_batch(self, batch: List[Request]):
        batch_ids = [r.req_id for r in batch]
        max_prompt = max(r.prompt_len for r in batch)
        max_output = max(r.output_len for r in batch)

        # 用 sleep 模拟 Prefill + Decode 成本
        prefill_cost = max_prompt * 0.002
        decode_cost = max_output * 0.01

        print(f"[Scheduler] batch={batch_ids}, size={len(batch)}, "
              f"prefill={prefill_cost:.3f}s, decode={decode_cost:.3f}s")

        await asyncio.sleep(prefill_cost + decode_cost)

        now = time.time()
        for r in batch:
            latency = now - r.created_at
            print(f"  -> req={r.req_id}, total_latency={latency:.3f}s")

async def producer(scheduler: DynamicBatchScheduler, total=12):
    for i in range(total):
        req = Request(
            req_id=i,
            prompt_len=random.randint(20, 200),
            output_len=random.randint(16, 128),
        )
        await scheduler.submit(req)
        await asyncio.sleep(random.uniform(0.01, 0.08))

async def main():
    scheduler = DynamicBatchScheduler(max_batch_size=4, batch_wait_ms=40)
    worker = asyncio.create_task(scheduler.run())
    await producer(scheduler, total=12)

    await asyncio.sleep(5)
    scheduler.running = False
    worker.cancel()

if __name__ == "__main__":
    asyncio.run(main())

这个例子能说明什么

虽然它不是真正的 GPU 推理,但足够说明几个关键事实:

  • batch 变大后,总体吞吐通常更好
  • 但凑批等待时间会拉高部分请求延迟
  • 一个超长请求会拖慢整个 batch
  • 因此生产环境常常要做长度分桶请求分类

进一步:为什么长度分桶很重要

如果短请求和长请求混在一起,通常会发生“木桶效应”:
batch 的执行时间常常接近其中最重的那个请求。

flowchart LR
    A[短请求] --> D[同一批次]
    B[中请求] --> D
    C[长请求] --> D
    D --> E[批次耗时接近最长请求]
    E --> F[短请求被长请求拖慢]

    G[长度分桶] --> H[短中长分开调度]
    H --> I[延迟更稳定]

逐步验证清单

如果你准备把优化从实验环境推进到线上,我建议按下面顺序验证,不要一上来全开。

第一步:先验证 KV Cache 是否真正生效

检查点:

  • generate(..., use_cache=True) 是否开启
  • 模型配置里 config.use_cache 是否被意外关闭
  • 长输出场景下,TPOT 是否明显下降
  • 显存占用是否随上下文增长而上升

第二步:再验证量化收益是否真实

检查点:

  • 模型加载后显存下降多少
  • 平均延迟、P95 延迟是否改善
  • 输出质量是否退化
  • 某些任务是否出现奇怪重复、事实性下降

第三步:最后做并发调度调参

重点观察:

  • batch size 增大后,TTFT 是否劣化过多
  • 短请求是否被长请求拖慢
  • GPU 利用率是否上升
  • OOM 是否更频繁出现

常见坑与排查

下面这些问题,我基本都见过,很多还真不是看一眼日志就能发现。

1. 开了 KV Cache,速度却没明显提升

可能原因

  • 输出很短,缓存收益还没体现出来
  • Prefill 占比太大,Decode 提升被掩盖
  • batch 太小或负载太轻,GPU 原本就不忙
  • 某些框架路径没有真正走到缓存优化

排查建议

  • max_new_tokens 拉长到 128、256 再测
  • 区分记录 TTFT 和 TPOT,不要只看总耗时
  • 打印模型配置中的 use_cache
  • 使用 profiler 看 decode 每步耗时是否下降

2. 量化后更慢了

可能原因

  • 当前 GPU 对低比特算子支持一般
  • 量化/反量化开销抵消了收益
  • 框架没有调用高效 kernel
  • 任务偏低延迟场景,不是带宽瓶颈

排查建议

  • 先比较 FP16 vs INT8,再比较 INT8 vs 4bit
  • 固定 batch size、输入长度,做同口径测试
  • 看吞吐和延迟,不要只看其中一个
  • 检查框架版本、CUDA 版本、bitsandbytes 版本兼容性

3. 并发一上来就 OOM

可能原因

  • KV Cache 随上下文和并发叠加膨胀
  • 动态 batch 过大
  • 显存碎片严重
  • 长上下文请求没有限流

排查建议

  • 限制 max_batch_size
  • 限制最大输入长度和最大输出长度
  • 对超长请求单独队列处理
  • 监控峰值显存,而不是只看平均值

4. 延迟抖动非常大

可能原因

  • 请求长度差异大
  • 调度窗口过长,凑批等待明显
  • 机器上混跑其他任务
  • 显存频繁回收或碎片化

排查建议

  • 做长度分桶
  • 缩短 batch wait 时间
  • 隔离推理服务节点
  • 观察 P95/P99,而不是只看平均值

安全/性能最佳实践

这里的“安全”,我更偏向线上可控性与资源安全,因为推理服务最怕的是一两个异常请求把整台机器拖垮。

1. 给上下文长度设硬限制

不要完全相信客户端传入的 prompt 长度。
超长输入会直接抬高:

  • Prefill 时间
  • KV Cache 占用
  • OOM 风险

建议:

  • 对输入 token 数设置上限
  • 超限请求直接拒绝或截断
  • 把长上下文请求路由到专门实例

2. 把“吞吐最优”和“低延迟最优”分开配置

同一套参数不可能同时最优满足:

  • 在线问答的低 TTFT
  • 离线批处理的高吞吐

建议拆成两套服务配置:

  • 低延迟服务:小 batch、短等待窗口
  • 高吞吐服务:更大 batch、允许更长凑批

3. 量化前先做任务级回归

不要只看 benchmark 分数。
真正要回归的是你的业务任务,比如:

  • 摘要质量
  • 代码补全正确率
  • 问答事实性
  • 格式遵从率

建议:

  • 先抽一组固定样本集
  • 比较 FP16 与量化版本的输出差异
  • 对关键场景做人审

4. 监控一定要细到 token 维度

如果监控只看“请求成功率”和“平均响应时间”,很多问题会被埋掉。

建议至少监控:

  • TTFT
  • TPOT
  • tokens/s
  • batch size 分布
  • 输入长度 / 输出长度分布
  • 显存占用与峰值
  • OOM 次数
  • P95/P99 延迟

5. 预留显存余量,不要卡极限

我见过不少线上事故,本质都是“实验环境刚好能跑”就直接上线。
但线上有波动、有尖峰、有偶发超长请求。

经验上建议:

  • 不要把显存打到 95% 以上常态运行
  • 给 KV Cache 波动预留空间
  • 给框架内部 buffer 和碎片预留空间

6. 优先采用成熟推理框架

如果你已经进入生产环境,通常不建议从零手写调度和缓存系统。
成熟框架在这些点上已经做了很多优化:

  • Paged KV Cache
  • Continuous batching
  • 高性能 attention kernel
  • 更细粒度的显存管理

常见选择包括:

  • vLLM:在吞吐和连续批处理上非常常见
  • Text Generation Inference (TGI):部署生态较成熟
  • TensorRT-LLM:在特定 NVIDIA 环境下性能很强

当然,框架不是银弹。
你仍然需要理解这篇文章讲的原理,否则调优时只会“碰运气”。


一个实战上的推荐优化路径

如果你现在手上有一个“能跑但不够快”的大模型推理服务,我建议按下面顺序做:

路径一:先拿到基线

先测出:

  • FP16/BF16
  • 无特殊调度
  • 固定输入长度和输出长度
  • 单并发 / 多并发

产出基线指标:

  • TTFT
  • TPOT
  • 吞吐
  • 显存峰值
  • P95/P99

路径二:启用 KV Cache

目标:

  • 看 Decode 是否明显加速
  • 看显存增加是否可接受

如果你的任务以生成长回答为主,这一步通常收益最大。

路径三:引入 8-bit 量化

目标:

  • 先释放显存
  • 为更大 batch 或更长上下文创造空间

我一般建议从 8-bit 开始,而不是直接冲 4-bit。

路径四:做动态批处理,再尝试连续批处理

目标:

  • 提升 GPU 利用率
  • 提高整体吞吐
  • 控制尾延迟

如果请求长度差异大,再加上长度分桶。

路径五:回到业务目标调参

最后别忘了,优化不是为了跑分,而是为了业务:

  • 如果你做的是在线客服,优先压低 TTFT
  • 如果你做的是离线生成,优先拉高吞吐
  • 如果你做的是长文档问答,重点看 Prefill 和 KV Cache

总结

大模型推理优化,最核心的不是记住一堆术语,而是建立一个稳定的判断框架:

  • KV Cache:减少 Decode 阶段重复计算,提升单请求生成效率
  • 量化:降低权重显存和带宽压力,为更高并发、更长上下文腾空间
  • 并发调度:把释放出来的资源转化成真正的系统吞吐

如果只让我给三个最实用的建议,我会说:

  1. 先开 KV Cache,再测长输出场景的 TPOT
  2. 量化先从 8-bit 开始,不要盲目追求更低 bit
  3. 调度一定配合长度分桶和显存监控做,不然很容易吞吐上去了,尾延迟和 OOM 也上去了

最后给一个边界判断:

  • 如果你的请求很短、并发很低,复杂调度收益可能有限
  • 如果你的上下文很长、输出也长,KV Cache 基本是必选项
  • 如果你的显存已经成为并发瓶颈,量化往往是最直接的杠杆
  • 如果你的 GPU 利用率长期不高,并发调度通常比“继续换大卡”更值得先做

真正好的推理优化,不是某个指标单点爆表,而是在成本、延迟、吞吐、稳定性之间拿到一个能长期运行的平衡点


分享到:

上一篇
《从源码到部署:基于开源项目 MinIO 搭建高可用对象存储服务的实战指南-220》
下一篇
《Docker 多阶段构建与镜像瘦身实战:面向中级开发者的构建提速、体积优化与安全加固指南》