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 在整个过程中始终保持占用。其结果是吞吐量显著提高,平均延迟降低。

请求生命周期

请求从提交到完成会经历四个状态。

       WAIT IN QUEUE          LOAD PROMPT INTO KV       STREAM OUTPUT           DONE
    ┌────────────────┐    ┌─────────────────────┐    ┌─────────────────┐    ┌────────┐
    │    PENDING     │───▶│     PREFILLING      │───▶│    DECODING     │───▶│FINISHED│
    │  ○ ○ ○ · · ·   │    │ ████████░░░░░░░░░░  │    │  → → → → · ·    │    │   ✓    │
    └────────────────┘    └─────────────────────┘    └─────────────────┘    └────────┘
         others ahead            prompt chunks              +1 token/step
  1. Pending — 请求已加入队列,正在等待调度。
  2. 预填充 — 提示词标记在一次前向传播中处理。
  3. 解码 — 输出 token 逐个生成。
  4. 已完成 — 生成完毕,结果已可用。

当调度器找到足够的 token 预算和缓存空间以接纳请求时,该请求会从 pending 状态移至 prefilling 状态。如果提示词过长,无法在单个步骤中容纳,调度器会将其拆分到多个前向传播中(chunk prefill)。

分块预填充

在单个步骤中处理长提示词成本较高,因为它会阻塞其他请求的生成。分块预填充通过将大型预填充拆分到多个步骤中来解决此问题。

Blocking (long prefill blocks the batch)
  Step:   1        2        3        4
  Req A  [▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓]   ← one huge prefill
  Req B  ················ wait ················
  Req C  ················ wait ················

Chunked (same prompt split and interleave with decode)
  Step:   1          2          3          4
  Req A  [▓▓ pre]    [▓▓ pre]   [▓▓ pre]   [→ dec]
  Req B  [→ dec]     [→ dec]    [▓▓ pre]   [→ dec]
  Req C  [→ dec]     · idle ·   [→ dec]    · idle ·

当新请求的提示词超出可用的 token 预算(由 ContinuousBatchingConfig 中的 max_batch_tokens 设定)时,调度器会处理尽可能多的 token,并保留其余部分。在后续步骤中,它会从上次中断处继续,将 prefill 工作与正在进行的 decode 步骤交错执行。这使批次保持高效,并降低了其他请求的首 token 延迟。

调度器

调度器根据两个预算和一个请求上限,决定哪些请求加入每次前向传播。

  • Token 预算 — 单次前向传播中处理的最大查询 token 数量,由 ContinuousBatchingConfig 中的 max_batch_tokens 设置。
  • 缓存预算 — 单次读取中可读取的 KV 页面总数,其上限由缓存总大小决定。
  • 请求上限 — 单次前向传播中请求的最大数量,由 max_requests_per_batch 设置。由于 logits 张量的大小随词表大小变化,因此该上限限制了大词表模型的峰值内存。

FIFO

FIFOScheduler 是默认调度器。它按优先级顺序填充批次:首先处理处于活跃状态的解码请求,然后是处于活跃状态的预填充请求,最后按到达顺序处理新的等待请求。这种策略通过优先处理现有请求再引入新请求,从而最大化吞吐量。

PrefillFirst

PrefillFirstScheduler 会在恢复解码之前完成分块预填充操作。这可以减少以长提示词为主的工作负载中的碎片化。

在 ContinuousBatchingConfig 中设置调度器。

cb_config = ContinuousBatchingConfig(scheduler_type="prefill_first")

内存管理

KV 缓存采用分页机制,使得不同长度的请求能够共享一个固定的内存池,而不会产生内存碎片。一个由 CachePool 实现的共享内存块池,会向每种注意力类型的分配器提供这些内存块,分配器再将它们划分为块和页。

Shared pool (first two chunks reserved for padding / never allocated)

  chunk N  [ ======================================== ]
               |                            |
    full attn: [ FULL BLOCK 0 | FULL BLOCK 1 | … ]
               [ page L0 | page L0 | … ]   (one page per full layer)

    sliding:   [ ----------- SLIDING BLOCK 0 ----------- ]
               [ page L1 | page L2 | page L3 | … ]  (one page per sliding layer)

  page  = page_size tokens of KV for one layer
  block = one page per layer of that attention type
  chunk = pool unit sized so full and sliding blocks tile the same bytes

连续批处理仅支持两种注意力类型:full_attention(FullAttentionCacheAllocator)和 sliding_attention(SlidingAttentionCacheAllocator)。混合模型会获得每种类型各一个分配器。注意力类型由模型(layer_types / sliding_window)决定,而非来自 ContinuousBatchingConfig。

滑动层通过复用少量页面,仅保留最近 token 的固定窗口。它们不支持内核块表或块共享。

缓存会在数据块之前预留两个块,用于处理请求中从未使用的填充读取和写入。这些块永远不会分配给任何请求。

缓存大小

管理器在启动时根据空闲的 GPU 内存推断可容纳的块数。它会考虑 KV 张量、注意力掩码、激活值以及簿记索引,然后将池的大小调整为可用内存的 max_memory_percent 以内。如果未设置,则会根据硬件和连续批处理参数采用默认值。在 CPU 和 GPU 共享内存的 MPS 环境下,会为缓存预留较少的空间,以避免交换。当显式设置 num_blocks 时,它表示容量为 page_size 个 token 的完整模型块数量。

您可以固定 ContinuousBatchingConfig 中的值。

cb_config = ContinuousBatchingConfig(
    page_size=256,
    num_blocks=4096,
    max_batch_tokens=512,
    max_memory_percent=0.8,
)

准入

在请求加入批次之前,调度器会检查每种注意力类型分配器是否都有足够的空闲容量。如果任何分配器的容量不足,该请求将被拒绝,且不会进行任何分配。

为避免缓存被占满以至于只能进行offloading,调度器会强制执行safety_margin。当空闲缓存容量低于safety_margin时,新的 prefill 请求将被暂缓,仅继续处理正在进行的 decode。FIFO 调度器的默认值为0.15(预留 15% 的容量),而prefill_first的默认值为0.0。可在ContinuousBatchingConfig中设置safety_margin以进行调整,其中0表示禁用该预留。

前缀缓存

当两个请求共享相同的提示前缀时,它们可以共享存储该前缀 KV 的块。每个已完成的块都会进行内容哈希。后续具有匹配前缀的请求会复用该块,并跳过对这些 token 的预填充(prefill)。共享块采用引用计数机制,只有当所有使用它们的请求都完成后,才会返回到空闲池。

allow_block_sharing 的默认值为 True。仅当所有分配器都支持共享时,前缀共享才会启用,因此只有纯全注意力模型会使用它。纯滑动窗口模型和混合模型会将其禁用。对于短提示词、长生成的工作负载,由于簿记开销超过了节省的成本,请设置 allow_block_sharing=False。

驱逐

释放的块不会被擦除,它们保持“已初始化”状态,其内容和哈希值完整保留,以便后续的匹配前缀能够复用它们。当未初始化池中的块不足时,管理器会延迟地将最近的一个已初始化块降级回未初始化状态,并丢弃其哈希值。前缀缓存是一个尽力而为的层,分配器会将其丢弃,以确保能够继续处理新请求。

软重置

软重置是 offloading 的回退路径。当管理器无法将某个请求的 KV 缓存复制到 CPU 交换池时(例如 CPU offloading 已禁用或交换池已满),它会改为对该请求执行软重置。具体操作是:将该请求已生成的 token 追加到其原始 prompt 中,释放其缓存块,并将其重新加入队列。Offloading 部分介绍了如何选择需要软重置的请求。

当缓存再次有可用空间时,请求将从中断处恢复,其生成历史已编码在提示中。

CUDA 图

每个生成步骤都涉及 CPU 开销,从组装批次到调度 GPU 内核以及读取结果。CUDA graphs 通过一次性记录完整的 GPU 执行序列,并在形状匹配的批次上重放该序列,从而消除这些开销。

由于批处理形状在每一步都会发生变化,连续批处理系统通过填充和缓存来处理这一问题。

  1. 查询长度会被填充到 q_padding_interval_size 的最近整数倍。
  2. KV 长度会被填充到 kv_padding_interval_size 的最近整数倍。
  3. 记录的图会存储在缓存中。

当某个批次的填充形状与已缓存的图匹配时,该图将直接重放,无需任何 CPU 调度。新的形状会触发一次新的图捕获。

CUDA 图要求静态控制流,并且与注意力掩码不兼容。默认情况下会自动检测,当条件不满足时会被禁用。

异步批处理

异步批处理使用两个 I/O 缓冲对和两个 CUDA 流,以重叠 CPU 和 GPU 的工作。

Sequential
  CPU  ── [prep N] ·········idle········· [prep N+1] ·········idle·········
  GPU  ·········idle········· [compute N] ·········idle········· [compute N+1]

Async
  CPU  ── [prep N] [prep N+1] [prep N+2] ──
  GPU  ── ········ [compute N] [compute N+1] ──

在 GPU 计算第 N 个批次(batch N)的同时,CPU 会准备第 N+1 个批次(batch N+1)。然而,将这两者重叠执行大约需要两倍的显存(VRAM)。

异步批处理需要启用 CUDA graphs,因为图重放提供了流重叠正确性所需的稳定张量地址。

Offloading

生成较长输出的请求可能会不断增长,直到 KV cache 没有空间来分配下一个块。为了避免崩溃,管理器会将处于活动状态的请求卸载回等待队列,从而为剩余请求释放块。

Offloading 是由需求驱动的。调度器会报告 stalled 请求,即那些未能分配到所需块的活动请求,以及每个请求所需的块数。管理器会卸载最少数量的请求,以释放足够的块来满足该需求。它会优先卸载最新的请求,因此已调度到当前批次中的请求可以保留其缓存,且最后一个剩余的活动请求永远不会被卸载。

两种情况会触发卸载。

  • 一个批次已被调度,但部分活跃请求发生了停滞。管理器将足够多的停滞请求卸载,以便下一个批次能够为其余请求分配块。
  • 缓存已满,无法调度任何批次。管理器会卸载请求,并在同一步骤内循环重试调度,直到某个批次能够装入,或者没有更多请求可以卸载。它不再等待下一次前向传播来恢复。

被卸载的请求会走两条路径之一。当 cpu_offload_space 大于 0.0 时,管理器会将该请求的 KV 缓存块复制到预先分配的固定 CPU 缓冲区中,并将所有能放入交换池的请求批量合并为一次传输。待 GPU 缓存空间可用后,管理器会将这些块复制回 GPU,并恢复该请求,而无需重新计算其提示词和已生成的 token。对于无法放入交换池的请求,或者在禁用 CPU 卸载时的所有被卸载请求,则会回退到 soft reset。

后续步骤

Update on GitHub