Transformer 的自回归解码在每一步dou要把前面的 Key/Value 缓存下来这块儿叫 KV‑Cache。若不加以管理，显存会像泄漏的水桶一样快速被填满。

Transformer 解码:
- 每个 token 需要保存 Key 与 Value
- batch_size= , seq_len= , hidden= -
KV Cache:  ×  ×  ×  ×  2 字节 ≈ 1GB

Ru果一次性为所有序列预留完整内存，往往会导致 OOM；相反，Ru果按需分配，则Ke以让显卡保持高利用率。

PagedAttention 把 KV‑Cache 切成若干「页」并共享内存池，只在需要时才申请新页，从而把原本几 GB 的占用压到几百 MB。

pip install vllm   # 一键装好 PagedAttention 实现

不同位宽的量化方案在速度、显存和精度之间各有取舍：

下面演示如何在 vLLM 中打开 KV‑Cache 的 FP8 量化：

from vllm import LLM
llm = LLM(
    model="meta-llama/Llama-7b-hf",
    kv_cache_dtype="fp8",          # 将缓存压到 FP8
)

思路是先让一个体积geng小的“草稿”模型快速生成候选 token，再让完整的大模型并行校验，通过这种“抢跑—审查”模式Ke以把吞吐提升近两倍。

from transformers import AutoModelForCausalLM
target = AutoModelForCausalLM.from_pretrained
draft  = AutoModelForCausalLM.from_pretrained
outputs = target.generate(
    inputs,
    assistant_model=draft,
    max_new_tokens=128,
)
print

将网络层划分到不同 GPU 上，每张卡只负责自己那段计算，再通过高速互联传递激活值。适合层数hen多、单卡显存不足的千亿级模型。

from transformers import AutoModelForCausalLM
import deepspeed
model = AutoModelForCausalLM.from_pretrained
model = deepspeed.init_inference(
    model,
    mp_size=4,                 # 四卡张量并行
    dtype=torch.half,
    replace_with_kernel_inject=True,
)

A100 或 H100 的 NVLink 带宽足够支撑将同一层的大矩阵切分到多块显卡上执行，常配合 Flash‑Attention 使用效果geng佳。

标准 Self‑Attention 要保存 $O$ 大小的注意力矩阵，而 Flash‑Attention 利用块状乘法只保留必要的数据流，从而把显存占用降至原来的三分之一左右，同时算子实现了 Tensor Core 加速。

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-7b-hf",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",   # 开启 Flash‑Attention
)

传统批处理等所有请求凑齐后才一起算，一旦有长序列拖慢整体，GPU 利用率会瞬间掉到冰点。连续批处理则是“刚来就塞”，完成的序列立刻释放槽位，后面的请求立即填补空缺，实现了近乎满负荷运行。

from vllm import LLM, SamplingParams
llm = LLM(
    model="meta-llama/Llama-7b-chat-hf",
    tensor_parallel_size=2,
)
params = SamplingParams
prompts = 
outputs = llm.generate
for out in outputs:
    print

from fastapi import FastAPI
from pydantic import BaseModel
from vllm import LLM, SamplingParams
app = FastAPI
# 启动时加载模型，仅一次 I/O 开销
llm = LLM(
    model="meta-llama/Llama-7b-chat-hf",
    tensor_parallel_size=1,
    dtype="float16",
)
class GenerateRequest:
    prompt: str
    max_tokens: int = 128
    temperature: float = 0.8
@app.post
async def generate:
    sp = SamplingParams(temperature=req.temperature,
                        max_tokens=req.max_tokens)
    res = llm.generate
    return {"text": res.outputs.text}
# 启动方式： uvicorn main:app --host 0.0.0.0 --port 8000

apiVersion: apps/v1
kind: Deployment
metadata:
  name: llama-service
spec:
  replicas: 2
  selector:
      matchLabels:
          app: llama
  template:
      metadata:
          labels:
              app: llama
      spec:
          containers:
          - name: vllm
            image: ghcr.io/huggingface/text-generation-inference:latest
            args:
              - --model-id meta-llama/Llama-7b-chat-hf
              - --num-shard 2               # 两卡张量并行
              - --max-batch-prefill-tokens 2048
            resources:
              limits:
                nvidia.com/gpu: "2"
            ports:
              - containerPort: 80
---
apiVersion: v1
kind: Service
metadata:
  name: llama-svc
spec:
  selector:
      app: llama
  ports:
      - port: 80
        targetPort: 80
        protocol: TCP
  type: LoadBalancer   # 云厂商自动创建公网 IP 

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from vllm import LLM, SamplingParams
app = FastAPI
# 同一台机器上同时加载多个模型，内部共用显存池  
models = {
   "chat": LLM,
   "codellama": LLM,
}
class GenReq:
   model_name:str          # chat / codellama ...
   prompt:str             
   max_tokens:int=128     
   temperature float=0.9  
@app.post
async def gen:
   if model_name not in models:
       raise HTTPException
   sp = SamplingParams(temperature=req.temperature ,
                         max_tokens=req.max_tokens )
   out = models.generate
   return {"text": out.outputs.text}
# 启动方式同上，只要改端口即可对外提供多种服务。

​​​ ​​​

​​​ ​​​ ​​​​​​​

注：以上数据均为个人实验室环境测得，仅供参考。实际效果仍受硬件驱动版本和网络带宽影响。

十、一句话

想让千亿参数的大语言模型像闪电一样跑完，是技术细节 + 心里暗暗祈祷这两件事共同作用的结果！🚀💡🌈 🚧 🤖 📈







本文所使用代码均Yi在公开仓库测试通过如有疑问请在评论区留言

©2026 AI技术社区 · 保留所有权利 如需转载，请注明出处及原文链接 作者： 专家

标签： 详解

---

• 上一篇： 让智能助手从得心应手迈向心悦诚服 —— Swift 技术周报解析
• 下一篇： 校招C的20并发系列08，有没有用并发容器替代STL的实战案例？