Hugging Face's logo

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.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

连续批处理

本页介绍了连续批处理推理背后的类：提交提示词、配置调度和内存限制，以及获取结果。

有关使用示例，请参阅 连续批处理 指南；有关调度与内存的交互方式，请参阅 连续批处理架构 文档。

ContinuousMixin.generate_batch

transformers.ContinuousMixin.generate_batch

< >

( inputs: listgeneration_config: transformers.generation.configuration_utils.GenerationConfig | None = Nonecontinuous_batching_config: transformers.generation.configuration_utils.ContinuousBatchingConfig | None = Nonerecord_timestamps: bool = Falseprogress_bar: bool = Falsepersistent_manager: bool = Falsewarmup: bool = True**kwargs ) → dict[str, GenerationOutput]

参数

  • inputs — 输入 token 序列（提示词）列表
  • generation_config — 可选的生成配置
  • continuous_batching_config — 可选的连续批处理配置
  • record_timestamps — 如果设置为 true，请求将为每个生成的 token 记录时间戳
  • progress_bar — 如果设置为 true，将显示进度条
  • persistent_manager — 生成结束后是否保留管理器。默认值为 False。
  • warmup — 是否在处理请求之前预先捕获 CUDA 图。默认值为 True。

返回值

dict[str, GenerationOutput]

一个字典，将请求 ID 映射到 GenerationOutput 对象

使用连续批处理为一批提示生成序列。

ContinuousBatchingManager

class transformers.ContinuousBatchingManager

< >

( model: ProtoPretrainedModelgeneration_config: GenerationConfigcontinuous_batching_config: ContinuousBatchingConfigworkload_hints: transformers.generation.continuous_batching.utils.WorkloadHints | None = None )

用于处理生成请求连续批处理的 Manager。它提供了用于提交生成请求、获取结果以及管理后台生成线程的用户界面。不应直接创建此类，而应通过以下入口点之一（均为 ContinuousMixin 混入类的方法）进行创建：

  • init_continuous_batching
  • continuous_batching_context_manager
  • generate_batch

add_request

< >

( input_ids: listrequest_id: str | None = Nonemax_new_tokens: int | None = Nonestreaming: bool = Falserecord_timestamps: bool = Falseeos_token_id: int | list[int] | None = None**logit_processor_kwargs: Any ) → str | None

参数

  • input_ids — 用作提示的输入 token ID
  • request_id — 可选的自定义请求 ID（如果为 None，则自动生成）
  • max_new_tokens — 要生成的新 token 的最大数量
  • streaming — 是否在生成时流式传输 token
  • record_timestamps — 是否记录每个生成 token 的时间戳
  • eos_token_id — 序列结束标记的 ID
  • logit_processor_kwargs — logits 处理器的关键字参数。

返回值

str | None

如果进程是 TP 驱动程序，则为请求 ID，否则为 None。

向队列中添加一个新的生成请求。如果该进程不是 TP 驱动程序，则此操作无效。

add_requests

< >

( inputs: listmax_new_tokens: int | None = Nonestreaming: bool = Falserecord_timestamps: bool = False**logit_processor_kwargs: Any )

用于批量执行 add_request 并返回其 ID 的实用函数。有关更多详细信息，请查阅其文档。

cancel_request

< >

( request_id: str )

通过 ID 取消请求。如果此调用来自非 TP 驱动程序进程，则此操作无效：只有 TP 驱动程序进程会与管理器进行交互。

销毁

< >

( )

终止管理器并释放分布式资源。可以安全地多次调用。调用此方法后， 无法重新启动管理器。

get_result

< >

( request_id: str | None = Nonetimeout: float | None = None )

从输出队列中获取一个结果。如果提供了 ID，则返回第一个匹配的请求。如果提供了超时时间，则在超时（以秒为单位）后返回 None。

is_running

< >

( )

如果后台生成线程已启动且仍然存活，则返回 True。

join

< >

( stop_trigger_time: floattimeout: float | None = None )

等待后台线程结束。可以使用 timeout 参数（以秒为单位）来设置等待的超时上限。

暂停

< >

( )

一个上下文管理器，用于暂停生成循环，以便调用线程可以使用该模型，通常是为了就地更新它。只有当循环暂停后，线程才会进入此上下文，并在退出时恢复循环，同时保留其缓存和正在处理的请求：不会清空任何内容，也不会丢失任何请求。 多个线程可以同时持有暂停状态，当最后一个线程退出时，循环才会恢复。 如果启用了 TP，所有 rank 都必须进入此上下文，否则其他 rank 将永远挂起。

register_result_handler

< >

( request_id: strcallback: Callable )

参数

  • request_id (str) — 用于接收输出的请求 ID。
  • callback（callable）— 针对每个结果，使用 GenerationOutput 进行调用。

注册用于结果传递（流式或非流式）的回调。

每次为该请求生成结果时，回调都会通过 call_soon_threadsafe 在事件循环中被调用。对于流式请求，这会在每个 token 上发生；对于非流式请求，仅在完成时发生。请求结束时，处理程序会自动清理。

request_id_iter

< >

( request_id: str )

迭代与特定请求 ID 匹配的结果（阻塞）。

使用带有重新入队功能的共享输出队列。对于高并发服务，请改用 register_result_handler。

开始

< >

( )

启动后台生成线程。

停止

< >

( block: bool = Truetimeout: float | None = Nonekeep_for_next_session: bool = Falsehard_stop: bool = False )

停止后台生成线程。如果 block 标志设置为 True，则此方法会等待线程停止，最长等待时间为 timeout 秒（None 表示无超时）。如果 keep_for_next_session 标志设置为 True，则管理器会被缓存在模型上以供后续使用。如果设置了 hard_stop 标志，后台生成线程将立即停止，且待处理请求将失败。

switch_to_cb_friendly_attn

< >

( model: ProtoPretrainedModelauto_switch_to_flash: bool = True )

将 attn 实现切换为与 CB 兼容的实现。如果 auto_switch_to_flash 为 True，且 attn 实现为 SDPA 或 eager，则在支持且可用的情况下也切换到 flash。

预热

< >

( )

通过运行虚拟批次，为变长和解码路径预捕获 CUDA 图。如果尚未初始化，则初始化批次处理器。

连续批处理配置

class transformers.ContinuousBatchingConfig

< >

( page_size: int = 256num_blocks: int | None = Nonemax_batch_tokens: int | None = Nonemax_memory_percent: float | None = Nonemax_requests_per_batch: int | None = Nonemax_blocks_per_request: int | None = Noneallow_block_sharing: bool = Trueuse_async_batching: bool | None = Noneuse_cuda_graph: bool | tuple[bool, bool] | None = Noneq_padding_interval_size: int = 0kv_padding_interval_size: int = 0varlen_compile_config: transformers.generation.configuration_utils.CompileConfig | None = Nonedecode_compile_config: transformers.generation.configuration_utils.CompileConfig | None = Nonedefault_compile_level: int = 0scheduler_type: str = 'fifo'safety_margin: float | None = Nonereturn_logprobs: bool = Falseseed: int | None = Noneauto_switch_to_flash: bool = Truecpu_offload_space: float | None = 0.0cpu_offload_space_safety_threshold: float = 0.8max_queue_size: int = 0per_request_processors: bool = Falsedrop_unsupported_processors: bool = Truedisable_nccl_graph_mixing: bool = Truecpu_group_timeout: float | None = 300.0use_default_compile_configs: bool | None = Nonemax_cached_graphs: int | None = Noneblock_size: int | None = None )

参数

  • page_size (int，可选，默认值为 256) — 每个（全注意力）页内为每一层存储的 token 数量。存储 N 层缓存的块包含 N 个页（每层一个），每个页为单层存储 page_size 个 token 的缓存。默认值为 256。
  • num_blocks (int, 可选) — KV 缓存中的块数。当为 None 时，会根据 GPU 内存自动推断。
  • max_batch_tokens (int, 可选) — 批次中的最大 token 数。当为 None 时，会根据 GPU 内存自动推断。
  • max_memory_percent (float, 可选) — 用于 KV 缓存的 GPU 空闲内存最大百分比（在模型加载之后）。当 None 时，运行时解析为：如果没有 logit 处理则为 0.9，如果有则为 0.8，以便为词表大小的临时张量预留空间。
  • max_requests_per_batch (int, 可选) — 每个批次中的最大请求数。当为 None 时，会根据工作负载提示自动推断，默认回退值为 1024。
  • max_blocks_per_request (int, 可选) — 每个请求的最大块数，用于 flash_attn_with_kvcache 快速解码路径中确定块表的维度。将其设置为 0 可禁用快速解码路径。默认值为 None（自动推断）。
  • allow_block_sharing (bool，可选，默认为 True) — 是否允许为前缀缓存启用块共享。块共享只能被允许，而不能被强制启用，因为某些模型不支持该功能。如果您拥有少量短提示词但生成长度较长，请禁用此功能。
  • use_async_batching (bool, 可选) — 是否启用异步双缓冲，该功能以 VRAM 使用量翻倍为代价，消除连续批处理循环中的 CPU 开销。当值为 None 时自动检测。
  • use_cuda_graph (bool 或 tuple[bool, bool]，可选) — 是否启用 CUDA graphs。该参数可以是一个布尔值元组（分别对应 varlen 路径和 decode 快速路径），一个同时应用于两条路径的布尔值，或 None（自动推断）。在调用 decide_use_cuda_graphs 之后，该属性将变为一个布尔值元组。默认值为 None（自动推断）。
  • q_padding_interval_size (int，可选，默认值为 0) — 用于 CUDA graphs 的查询填充粒度（以 token 为单位）。当设置为 0 时，使用 continuous_api.py 中的预设值。
  • kv_padding_interval_size (int，可选，默认值为 0) — 用于 CUDA graphs 的 KV padding 粒度（以 token 为单位）。当设置为 0 时，使用 continuous_api.py 中的预设值。
  • varlen_compile_config (CompileConfig，可选) — 用于 varlen（prefill）路径的 CompileConfig。默认值为 None（使用 generation_config 作为回退）。 varlen 路径处理具有不同 query 和 KV 长度的批次，通常受益于 dynamic=True。
  • decode_compile_config (CompileConfig, 可选) — 用于解码（快速）路径的 CompileConfig。默认值为 None（使用 generation_config 作为回退）。 解码路径处理的批次没有动态的 KV 长度，因此静态形状更为合适。
  • default_compile_level (int，可选，默认为 0) — 如果该值 >0 且未为 varlen 或 decode 路径提供编译配置，则会自动提供一个默认的编译配置。级别最高可达 3，级别越高，性能越好，但预热时间也越长。
  • scheduler_type (str，可选，默认为 "fifo") — 要使用的调度器类型。
  • safety_margin (float, 可选) — 用于限制卸载量的安全边际。默认为 None（使用类默认值）。
  • return_logprobs (bool，可选，默认为 False) — 是否随生成的 token 一起返回对数概率。
  • seed (int | None，可选) — 用于生成的可选随机种子。如果未指定，内部种子将被设置为一个随机值。
  • auto_switch_to_flash (bool，可选，默认为 True) — 当实现为 SDPA 或 eager 时，是否自动切换到 flash attention（如果可用）。 默认值为 True，因为 flash 对连续批处理（continuous batching）性能更好。
  • cpu_offload_space (float, 可选，默认为 0.0) — 用于 KV cache 卸载的 CPU 交换空间，单位为 GiB。初始化时会创建一个该大小的预分配 pinned CPU 缓冲区。当 GPU 缓存已满时，被驱逐请求的 KV cache 会被复制到此缓冲区，而不是直接丢弃。0 表示禁用卸载（默认值）。
  • cpu_offload_space_safety_threshold (float, 可选，默认值为 0.8) — 如果 cpu_offload_space 超过系统总内存的该比例，则会被限制以避免主机 OOM。 设置为 1.0 可禁用此安全上限。当 psutil 不可用时，此参数将被忽略。
  • max_queue_size (int，可选，默认为 0) — 用于服务的最大请求队列大小。0 表示无限制。
  • per_request_processors（bool，可选，默认为 False）— 启用针对每个请求的 logits 处理器参数。默认值为 False。
  • drop_unsupported_processors (bool，可选，默认为 True) — 移除不受支持的 logits 处理器，而不是抛出错误。默认值为 True。
  • disable_nccl_graph_mixing (bool，可选，默认为 True) — 禁用 NCCL 针对并行图捕获通信的安全机制。在 CB 中此情况不会发生，且能为 TP 带来性能提升。
  • cpu_group_timeout (float，可选，默认值为 300.0) — CPU 通信超时并导致进程崩溃的时间（以秒为单位）。设为 None 表示不超时。默认值为 300 秒。
  • use_default_compile_configs (bool | None，可选) — 5.11 中已弃用：请改用 default_compile_level。
  • max_cached_graphs (int, 可选) — 5.13 中已弃用：图的最大数量不再是一个问题。
  • block_size (int | None，可选) — 5.17 版本中已弃用：现在改用 page_size。

用于保存与连续批处理相关的参数的类，适用于通过 generate_batch 方法或 continuous_batching_context_manager 上下文管理器使用连续批处理的情况。

Update on GitHub

←配置 数据整理器→