Transformers documentation

连续批处理

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

连续批处理

连续批处理通过在每次生成步骤中动态重新调度批次,从而最大化 GPU 利用率。当某些请求完成时,新请求会立即加入,而无需等待整个批次完成。GPU 始终保持满载,吞吐量也维持在较高水平。

在生产部署中,请使用 transformers serve。它基于 ContinuousBatchingManager 构建,并暴露了一个兼容 OpenAI 的 HTTP 端点。

generate_batch

连续批处理通过 generate_batch() 提供支持。传入一个已分词的提示列表,待所有请求完成后即可获取全部结果。generate_batch 在内部处理调度,并阻塞直到所有请求完成。

对于服务提供和流式传输用例,请直接使用 ContinuousBatchingManager 来管理请求。

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

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")

prompts = [
    "What's up?",
    "Name a cat breed.",
    "Write a detailed history of quantum mechanics.",
]
inputs = [tokenizer.encode(p) for p in prompts]

generation_config = GenerationConfig(
    max_new_tokens=64,
    eos_token_id=tokenizer.eos_token_id,
)

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

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

ContinuousBatchingManager

ContinuousBatchingManager 会运行一个后台线程,允许您独立地提交请求并获取结果。在每次生成步骤中,它会检查已完成的请求,并调度新请求加入批次。这对于流式处理、实时服务或按到达顺序提交请求的场景非常有用。

使用 continuous_batching_context_manager() 来安全地启动和停止管理器。下面的示例包含变长输入。一旦最短的提示词生成完毕,它就会离开批次,而较长的提示词则继续生成。如果使用静态批处理,你必须将它们全部填充到相同长度。连续批处理会释放已完成的提示词,从而可以立即开始处理下一个提示词。

with model.continuous_batching_context_manager(generation_config=generation_config) as manager:
    manager.add_request(
        input_ids=tokenizer.encode("Write a detailed history of quantum mechanics."),
        request_id="long",
        max_new_tokens=512,
    )
    manager.add_request(
        input_ids=tokenizer.encode("What's up?"),
        request_id="short_0",
        max_new_tokens=32,
    )
    manager.add_request(
        input_ids=tokenizer.encode("Name a cat breed."),
        request_id="short_1",
        max_new_tokens=32,
    )

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

您也可以调用 init_continuous_batching() 来自行管理生命周期。

manager = model.init_continuous_batching(generation_config=generation_config)
manager.start()

# submit and retrieve requests...

关闭管理器

Manager 会运行一个后台线程并持有分布式资源。关闭过程分为两个阶段,以便您可以选择如何处理正在进行的任务。

调用 stop() 以停止后台线程。默认情况下,管理器会停止接受新的提交,并等待队列中及正在处理的请求完成后再退出线程。

manager.stop()

传入 hard_stop=True 可立即放弃待处理的工作。已排队和正在进行的请求将抛出 RuntimeError 而不会完成。

manager.stop(hard_stop=True)

调用 stop 后,add_request() 和 add_requests() 会丢弃新的提交并记录一条警告。您可以再次调用 start,使用同一个管理器运行另一个生成会话。

调用 destroy() 以释放分布式资源。如果 destroy 仍在运行,它会先停止管理器,且之后无法重启管理器。请在进程生命周期内完成连续批处理(continuous batching)后使用此方法。

manager.destroy()

continuous_batching_context_manager() 负责处理此流程。退出时它会调用 stop,并调用 destroy,除非你传入 persistent_manager=True 以将管理器缓存在模型上供下一个会话使用。

添加请求

add_request() 用于提交单个请求。您可以提供 request_id,也可以让管理器自动生成一个。

manager.add_request(input_ids=input_ids, request_id="my_request")

add_requests() 一次性提交一个批次。当启用块共享时,它会自动对输入进行排序,以最大化前缀缓存命中率。

manager.add_requests(inputs=inputs)

使用 cancel_request() 取消请求。

manager.cancel_request(request_id="my_request")

按请求的采样参数

启用 per_request_processors,以便在同一次前向传播中针对每个请求独立应用 temperature、top_k 和 top_p,从而允许不同请求使用不同的采样参数(例如,一个请求生成富有创意、高温度的输出,而另一个请求生成精确、低温度的输出)。

cb_config = ContinuousBatchingConfig(per_request_processors=True)

# each request gets its own sampling parameters
manager.add_request(input_ids=inputs_a, temperature=0.9, top_p=0.95)
manager.add_request(input_ids=inputs_b, temperature=0.1, top_k=10)

GenerationConfig 中的每个参数都必须设置为非默认值,才能在运行时创建对应的 logits 处理器。例如,将 temperature 设置为 None 或 1 以外的值,以支持按请求控制温度。之后仍然可以创建温度为 1 的请求。

获取结果

遍历管理器以在结果到达时接收它们。

for result in manager:
    print(tokenizer.decode(result.generated_tokens, skip_special_tokens=True))

get_result() 从输出队列中获取下一个结果。传入 request_id 可筛选特定请求。如果队列中的下一个结果不匹配,该结果会被重新入队,方法返回 None。

# next available result
result = manager.get_result()

# filter for a specific request
result = manager.get_result(request_id="my_request")

流式

在请求中设置 streaming=True,然后使用 request_id_iter() 在生成 token 时迭代获取部分输出。

from transformers.generation.continuous_batching import RequestStatus

manager.add_request(input_ids=input_ids, request_id="streamed", streaming=True)

for chunk in manager.request_id_iter(request_id="streamed"):
    token = tokenizer.decode(chunk.generated_tokens[-1:], skip_special_tokens=True)
    print(token, end="", flush=True)
    if chunk.status == RequestStatus.FINISHED:
        break

ContinuousBatchingConfig

ContinuousBatchingConfig 用于控制 KV 缓存、调度、CUDA 图、内存使用等。请通过 continuous_batching_config 参数将其传递给 generate_batch 或 init_continuous_batching(),并与 GenerationConfig 分开传递。这两个配置描述的内容不同:GenerationConfig 以模型为中心,包含采样和停止参数;而 ContinuousBatchingConfig 以硬件为中心,包含内存和调度参数。

在 GenerationConfig 上设置 continuous_batching_config 已被弃用,会发出 FutureWarning,并将在 v5.19 中移除。

默认情况下,max_batch_tokens 为 8192,受可用 GPU 内存限制,且不会低于 256,而 num_blocks 会填满剩余的内存。请使用下表来帮助您选择合适的功能。

功能内存吞吐量延迟
max_memory_percent / page_size✓ 控制 KV 预算
max_batch_tokens↑ 更大的输入缓冲区✓ 更大的预填充批次✓ 预填充受限时的 TTFT
scheduler✓ 调度策略✓ TTFT
CUDA 图↑ 图存储✓ 更低的调度开销✓
异步批处理↑ ~2× I/O 缓冲区✓ 重叠 CPU/GPU
编译↑ 仅预热时间✓ 更快的正向传播✓
快速解码路径↑ 每个请求的块表✓ 更快的仅解码步骤✓
CPU 卸载↑ 固定 CPU 内存✓ 跳过部分重新预填充
前缀缓存↓ 共享 KV 块✓ 跳过冗余的预填充✓ TTFT
分页注意力↓ 无碎片✓ 动态批次成员
滑动窗口↓ 每层有界的 KV
按请求处理的处理器✓ 每个批次使用不同的采样参数
max_requests_per_batch↓ 限制 logits 缓冲区✓ 限制批处理大小
safety_margin✓ 保护解码延迟
from transformers.generation import ContinuousBatchingConfig

cb_config = ContinuousBatchingConfig(
    max_memory_percent=0.8,
    page_size=256,           # KV cache page size in tokens
    scheduler_type="fifo",        # "fifo" or "prefill_first"
)

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

KV 缓存页大小

page_size 设置单个注意力层中一个完整 KV 页面所容纳的 token 数量。该值必须至少为 4,否则初始化时会抛出 ValueError。默认值 256 适用于大多数工作负载。

block_size 参数已弃用,并由 page_size 替代。

更大的页面大小可以减少簿记开销,但当序列在页面中间停止时会浪费空间。较小的页面浪费的空间更少,但需要更多的簿记。通常无需调整 page_size。如果您的工作负载包含许多短请求,可以尝试较小的值。如果包含少数非常长的请求,可以尝试较大的值。有关页面、块和块如何协同工作,请参阅 内存管理。

cb_config = ContinuousBatchingConfig(page_size=128)

预填充批处理大小

max_batch_tokens 用于设置单次前向传播的 token 预算,即模型一次处理的最大查询 token 数量。更大的预算可以将更多的 prompt token 打包到每次 prefill 中,从而提升 prefill 吞吐量,并降低以 prompt 为主的工作负载的首 token 延迟。调度器会将超出预算的 prompt 拆分到多个步骤中处理。有关如何处理超大 prompt,请参阅 分块 prefill。

默认情况下,max_batch_tokens 为 8192。该值受可用 GPU 显存的限制,以确保每个批次的输入张量不会挤占 KV 缓存的空间,并且其值不会低于 256。在空闲显存较少的 GPU 上,该限制会将其值降低到 8192 以下。

max_batch_tokens 和 num_blocks 共享同一内存预算,因此二者相互制约。较大的 token 预算会分配更大的输入缓冲区,从而为并发或长请求保留更少的 KV cache 块。当可用内存更充足时,可提高 max_batch_tokens 以提升 prefill 吞吐量;反之,可降低该值以释放内存用于更多 KV 块,或避免内存溢出错误。

cb_config = ContinuousBatchingConfig(max_batch_tokens=16384)

批处理和调度限制

max_requests_per_batch 限制了单次前向传播中运行的请求数量。模型会生成一个 logits 张量,其大小由批处理中的 token 数量和词表大小决定,并且为了进行采样,该张量会被转换为 fp32,从而使内存占用翻倍。由于每个请求只需要一个下一个 token 的预测,因此设置较小的上限可以保持该张量更小,并避免在包含大量预填充(prefill)且词表较大的批处理中发生内存溢出错误。默认情况下,该值设置为提交的提示(prompt)数量,若未指定则回退为 1024,随后会根据 max_batch_tokens 和 num_blocks 进行上限约束以适配资源。

cb_config = ContinuousBatchingConfig(max_requests_per_batch=256)

safety_margin 为正在处理的请求预留一部分空闲缓存容量。当空闲容量低于 safety_margin(即总缓存容量的一定比例)时,调度器将停止接纳新的 prefill 请求,仅继续处理已在进行中的 decode 请求。这一机制优先保证完成当前活跃的工作,而非启动新请求,从而保护 decode 延迟并推迟 offloading。该值必须介于 0 和 1 之间,其中 0 表示禁用该预留机制。

cb_config = ContinuousBatchingConfig(safety_margin=0.15)

默认值取决于调度器。对于 FIFO 调度器,该值为 0.15;对于 prefill_first,该值为 0.0。有关 margin 如何与调度交互,请参阅 Admission。

对数概率

当 return_logprobs=True 时,ContinuousBatchingConfig 会返回每个生成 token 的对数概率。这在强化学习中非常有用,因为 logprobs 是某些训练循环的输入。

cb_config = ContinuousBatchingConfig(return_logprobs=True)

# generate_batch()

for request_id, output in outputs.items():
    for token_id, log_prob in zip(output.generated_tokens, output.logprobs):
        token = tokenizer.decode([token_id])
        print(f"{token} | logprob: {log_prob}")

CUDA 图

CUDA 图通过一次性记录 GPU 执行图,并在形状匹配的批次上重放该图,从而消除 CPU 调度开销。请显式使用 use_cuda_graph=True 来启用它们。

cb_config = ContinuousBatchingConfig(use_cuda_graph=True)

启用后,管理器会将查询和 KV 的长度填充到固定的间隔,从而使形状重复出现并实现图复用。q_padding_interval_size 和 kv_padding_interval_size 的值越小,填充造成的计算浪费越少,但这也意味着图需要记录和存储更多独特的形状,从而占用更多内存。

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    q_padding_interval_size=64,
    kv_padding_interval_size=16384,
)

异步批处理

异步批处理(Async batching)将下一批次的 CPU 调度与当前批次的 GPU 计算重叠执行。它需要 CUDA graphs,并且会使输入张量占用的显存大约翻倍。

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    use_async_batching=True,
)

编译

default_compile_level 会对模型的前向传播应用 torch.compile。编译以一次性的预热开销换取生成阶段更快的前向传播速度。更高的级别会执行更激进的优化,从而提升吞吐量,但也会延长预热时间。在测试和基准测试迭代期间,请保持较低的级别;对于长时间运行的服务负载,由于预热成本会在大量请求中分摊,因此可以提高该级别。

级别范围从 0 到 3。级别 0 为默认值,将完全跳过编译。

级别modedynamic权衡
0——无编译(默认),启动速度最快
1defaultTrue速度提升有限,预热时间短
2max-autotune-no-cudagraphsTrue更快的加速效果,更长的预热时间
3max-autotune-no-cudagraphsFalse最佳吞吐量,最长的预热时间
cb_config = ContinuousBatchingConfig(default_compile_level=1)

该级别为 varlen 和 decode 执行路径提供默认的 CompileConfig。它仅适用于未显式配置的路径,因此当 varlen_compile_config 和 decode_compile_config 被设置时,它们具有更高的优先级。在 FlashAttention 下,由于 max_length_k 会触发频繁的重新编译,varlen 路径会跳过编译,因此在这种情况下,该级别仅影响 decode 路径。

快速解码路径

当批次中仅包含解码请求(每个序列一个查询 token)时,管理器可以调度到 flash_attn_with_kvcache 内核,而不是变长内核。这比 varlen 路径更快,因为该内核通过块表直接就地读写分页 KV 缓存,而无需经过手动更新。有关内核级别的详细信息,请参阅 Paged attention。

快速路径的大小由 max_blocks_per_request 决定,该参数用于确定每个请求的块表维度。默认情况下,此值会自动推断。如果管理器上设置了 max_prompt_length 和 max_generated_length,则块表的大小将调整为可容纳最大序列长度。否则,将使用回退默认值(每个请求 32 个块)。

将 max_blocks_per_request 设置为特定值,以显式指定块表的大小。当您已知每个请求的最大序列长度并希望限制块表的内存开销时,此设置非常有用。

cb_config = ContinuousBatchingConfig(max_blocks_per_request=64)

将 max_blocks_per_request=0 设置为 0 可禁用快速路径,强制所有批次通过 varlen 内核处理。这将恢复默认行为之前的状态,在快速路径不适用于您的注意力实现时非常有用(当底层内核无法使用时,管理器也会自动禁用它)。

cb_config = ContinuousBatchingConfig(max_blocks_per_request=0)

快速路径依赖于 flash_attn_with_kvcache 内核,该内核支持两种设备和注意力实现组合。

设备attn_implementation
CUDAflash_attention_2 或 flash_attention_3
XPUflash_attention_2

对于其他任何组合,或者当无法导入内核时,管理器会回退到 varlen 路径。仅当显式设置了 max_blocks_per_request 时,才会记录警告。

Full attention 支持 kernel block tables,但 sliding attention 不支持。如果模型包含任何 sliding attention 层(包括同时使用 full attention 的混合模型),初始化会将 max_blocks_per_request 强制设为 0,您设置的任何非零值都会被覆盖。随后,管理器会记录日志,指出由于模型的部分 attention 类型不支持 kernel block tables,因此 decode 快速路径不可用。

CPU 卸载

CPU offloading 在 GPU KV 缓存已满时,将已驱逐的 KV 缓存块复制到预先分配的固定 CPU 缓冲区中。当缓存空间可用后,管理器会将这些块复制回 GPU,并恢复该请求,而无需重新计算其提示词和已生成的 token。

将 cpu_offload_space 设置为以 GiB 为单位的 CPU 交换空间。默认值 0.0 表示禁用 CPU 卸载。

cb_config = ContinuousBatchingConfig(cpu_offload_space=8.0)

默认情况下,当安装了 psutil 时,cpu_offload_space_safety_threshold=0.8 会将请求的空间限制为可用系统内存的 80%。将 cpu_offload_space=None 设置为根据安全阈值来确定交换池的大小。

张量并行超时

在张量并行下,管理器会创建一个 CPU 通信组,用于跨 rank 协调请求的提交、取消和关闭。cpu_group_timeout 限制了该组上的集合通信在导致进程崩溃前可以阻塞的最长时间。如果某个 rank 停滞,该超时机制可防止其他 rank 无限期等待。

对于发出较少集合通信操作的工作负载,请设置更长的超时时间,或传入 None 以禁用该功能。

cb_config = ContinuousBatchingConfig(cpu_group_timeout=600.0)

前缀缓存

当多个请求共享相同的前缀(例如系统提示词)时,管理器会复用它们的 KV 缓存块,而不是重新计算。allow_block_sharing 的默认值为 True。仅当所有注意力类型的分配器都支持共享时,前缀共享才会生效,因此只有纯全注意力模型会使用它。纯滑动窗口模型和混合模型即使 allow_block_sharing 为 True,也会将其关闭。

cb_config = ContinuousBatchingConfig(
    allow_block_sharing=True,  # default
)

选择注意力后端

Continuous batching 支持 FlashAttention、SDPA 以及 "paged|eager"。加载模型时,请设置 attn_implementation。

后端attn_implementation要求
FlashAttention 2“flash_attention_2”flash-attn 或 kernels 包
FlashAttention 3“flash_attention_3”flash-attn-3 或 kernels 软件包
SDPA(PyTorch 原生)“sdpa”无
Eager“paged|eager”无
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)

在启用连续批处理(Continuous batching)时,使用 FlashAttention 比使用 eager 或 SDPA 更快,这主要是因为 FlashAttention 不需要注意力掩码(attention mask)。如果模型支持 FlashAttention,那么以 "eager" 或 "sdpa" 方式加载的模型会切换到 FlashAttention 3;如果未安装 FlashAttention 3,则会切换到 FlashAttention 2,并记录一条警告。如果两个版本均未安装,则后端保持不变。

例如,在测试时若要保留 eager 或 SDPA,请在 ContinuousBatchingConfig 中设置 auto_switch_to_flash=False。只要后端保持为 "eager",连续批处理就会运行 "paged|eager"。

张量并行

对于无法在单个 GPU 上容纳的模型,请使用张量并行将权重分片到多个设备上。通过 DistributedConfig(tp_size=N) 设置设备数量。连续批处理会从模型中读取张量并行大小,以按分片确定分页 KV 缓存的大小。有关支持的架构列表及分片机制,请参阅 张量并行。

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

distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-32B",
    attn_implementation="flash_attention_2",
    distributed_config=distributed_config,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-32B")

inputs = [tokenizer.encode(p) for p in ["What's up?", "Name a cat breed."]]
generation_config = GenerationConfig(max_new_tokens=64, eos_token_id=tokenizer.eos_token_id)

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

使用 torchrun 启动脚本,并将 --nproc-per-node 设置为您希望进行分片的 GPU 数量。

torchrun --nproc-per-node 4 cb_tp.py

张量并行大小必须能整除模型的 num_key_value_heads(请检查模型配置)。否则,分页缓存会在启动时抛出错误,因此请选择合适的 --nproc-per-node。

不要同时设置 device_map 和 distributed_config。这两者会相互冲突,因为 device_map 会将整个模块放置在特定的 GPU 上,而张量并行则会将这些相同的参数分片到所有 GPU 上。

滑动窗口注意力

采用滑动窗口注意力机制的模型(Mistral、Gemma 2)支持连续批处理。滑动层使用 SlidingAttentionCacheAllocator,通过复用少量页面来仅保留最近 token 的固定窗口。注意力类型取自模型的 layer_types 或 sliding_window,而非 ContinuousBatchingConfig。

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("google/gemma-2-2b")
config.sliding_window = 4096

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b",
    config=config,
    attn_implementation="sdpa",
    device_map="auto",
    dtype=torch.bfloat16,
)

当存在滑动注意力时,前缀缓存和解码快速路径均会被禁用,包括同时包含完整注意力层的混合模型。

后续步骤

Update on GitHub