概述

Transformers 提供了多种推理优化技术，使模型运行更快、成本更低且更易获取。可选方案包括用于降低内存访问量的替代注意力机制、用于加速执行的代码编译，以及用于提升吞吐量的优化内核。将这些技术叠加使用，可获得最佳性能。

有关设备设置和特定于硬件的配置，请参阅硬件指南。在应用以下优化之前，请检查您的硬件支持哪些技术。

Agent prompt

帮助用户为其 Transformers 推理工作负载选择合适的优化方案。首先，在加载模型、处理长输入、生成长序列或处理并发请求时，识别内存或速度是否成为瓶颈。使用下方对应的指南，并在建议更改之前，检查该优化方案对用户模型和硬件的支持情况。

请询问模型及修订版本、库版本、GPU 及可用显存、CPU 内存、输入长度、生成上限、批处理大小以及并发请求数。当缺失的细节可能会影响推荐结果时，请提出后续问题，并说明你的假设。
根据问题发生的位置来选择排查方向。针对加载时的内存，检查模型权重和加载设置；针对生成序列增长时的内存，检查 KV 缓存；针对长输入的内存，检查注意力机制的内存分配；针对并发请求的内存，检查每个请求的状态和调度。将权重量化和 KV 缓存量化视为两个独立的设置。
每次只建议一项更改。先检查模型与硬件的兼容性，然后在同一具有代表性的工作负载上测量峰值 GPU 内存、延迟或吞吐量以及输出质量。将输入长度、生成长度或并发数的任何减少解释为工作负载的权衡。
报告所提议的配置、支持该配置的文档、预期的权衡、未解决的兼容性问题以及验证流程。将实测数据与估算值区分开来。如果无法运行该工作负载，请说明适配性和性能尚未得到验证。

内存和速度密切相关，但并不相同。缩小内存占用会使模型“更快”，因为需要移动的数据更少。纯粹的速度优化并不总能减少内存，有时甚至会增加内存使用量。请根据您的使用场景和硬件选择合适的优化方式。

使用下表选择一种优化技术。

技术	速度	内存
编译	✅
注意力后端	✅	✅
内核	✅	✅
量化	✅	✅
缓存	✅	✅
并行	✅
连续批处理	✅

本指南将带你快速入门 Transformers 优化。

编译

torch.compile 能够减少 Python 开销、融合算子，并针对您的张量形状和硬件生成调优后的内核。首次运行会进行预热，后续运行则会使用更快的编译路径。

将固定大小的缓存传递给 generate() 以自动触发 torch.compile。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", dtype=torch.float16, device_map="auto")
input = tokenizer("The French Bread Law states", return_tensors="pt").to(model.device)

output = model.generate(**input, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(output, skip_special_tokens=True)[0]

避免在 generate() 之外调用 torch.compile(model)，以防止模型在每一步都重新编译。

注意力后端

替代的注意力后端可以降低内存流量。例如，FlashAttention 通过分块处理注意力计算并避免生成大型中间张量，从而减少内存占用。

在 from_pretrained() 中设置 attn_implementation 以加载优化的注意力后端。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", attn_implementation="flash_attention_2")

内核

Kernels 通过融合操作来提升吞吐量并降低内存占用。Kernels 库以灵活且版本安全的方式从 Hub 加载优化后的计算内核。

以下示例在不安装该包的情况下加载了优化后的 FlashAttention-2 内核。

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", attn_implementation="kernels-community/flash-attn2"
)

量化

量化会缩小每个参数的尺寸，从而降低内存占用并提升速度，因为可以执行更多操作。

将量化配置传递给 from_pretrained() 中的 quantization_config 参数。每个量化后端都有不同的配置和不同的参数。下面的示例使用 bitsandbytes 后端将模型量化为 4 位，并配置计算的数据类型。

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)

model = AutoModelForCausalLM.from_pretrained(
    "allenai/Olmo-3-7B-Think", quantization_config=bnb_config
)

缓存

缓存通过复用过去的键和值，而不是为每个 token 重新计算，从而加速生成过程。为了抵消并降低存储过去键和值所带来的内存开销，Transformers 支持将缓存卸载到 CPU。只有当前层保留在 GPU 上。

使用 generate() 中的 cache_implementation 参数来设置缓存策略。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", attn_implementation="kernels-community/flash-attn2"
)
inputs = tokenizer("The Le Décret Pain states that a baguette must,", return_tensors="pt")
outputs = model.generate(**inputs, do_sample=False, max_new_tokens=50, cache_implementation="offloaded")

并行

并行将模型分布到多个设备上，使超出单个设备容量的模型也能快速运行。由于分片开销以及同步结果所需的通信，这种方法会占用更多内存。

张量并行将模型层拆分到多个设备上。使用 DistributedConfig(tp_size=N) 设置设备数量，并将其传递给 from_pretrained() 以启用该功能。

from transformers import AutoModelForCausalLM, DistributedConfig

distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    distributed_config=distributed_config,
)
print(model._tp_plan)

连续批处理

连续批处理通过动态调度和分块预填充让 GPU 保持忙碌，从而最大化吞吐量。服务应用使用它来并发处理多个传入请求。

使用 generate_batch() 以启用连续批处理。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="sdpa",
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

prompts = [
    "The Le Décret Pain states that a baguette must",
    "Explain gravity in one sentence.",
    "Name the capital of France.",
]
inputs = [tokenizer.encode(p) for p in prompts]

generation_config = GenerationConfig(
    max_new_tokens=32,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
    do_sample=False,
)
continuous_batching_config = ContinuousBatchingConfig(max_batch_tokens=512)

outputs = model.generate_batch(
    inputs=inputs,
    generation_config=generation_config,
    continuous_batching_config=continuous_batching_config,
)

for request_id, output in outputs.items():
    text = tokenizer.decode(output.generated_tokens, skip_special_tokens=True)
    print(f"[{request_id}] {text}")

Update on GitHub