针对如何处理“提前完成和延迟加入的请求”这个挑战，ORCA 给出的解决方案是用迭代级调度减少空闲时间，即以迭代为粒度（iteration-level）控制执行，而不是请求级粒度（request-level），并结合选择性批处理（selective batching）来进行优化。

3.1.1 迭代级调度

目标

迭代级调度的目标是：及时检测出推理完毕的请求，将其从 batch 中移出，以便新请求可以填补到旧请求的位置上，这样新请求和旧请求能接连不断组成新的 batch。

调度

前文提到，Orca 是第一篇提到迭代级别调度（Iteration-Level Schedule）的论文。具体来说就是：一个 batch 中的所有请求每做完 1 次 iteration（prefill 或者 decode），scheduler 就和 engine 交互一次，去检查 batch 中是否有做完推理的请求，以此决定是否要更新 batch。这样就可以在每次 GPU 推理的空隙，可以插入调度操作，实现 Batch 样本的增删和显存的动态分配释放。

下图给出了请求粒度的调度和迭代粒度调度的区别。前者需在整批请求全部完成前对调度批次进行多次迭代，而对于 ORCA，服务系统在调度任务时，每次只向 Execution Engine 提交一次迭代的计算，而非等到完成整个 Request 才能处理。这样 ORCA 就可以在每个迭代都动态更改要处理的请求，新请求只需等待单次迭代即可被处理，从而避免 early-finish 的请求等待其他请求的结束。通过迭代级调度，调度器能够完全控制每个迭代中处理哪些请求以及处理数量。

方案

下图展示了采用迭代级调度的 ORCA 系统架构和整体工作流程。ORCA 系统包括如下模块：

  • Endpoint（端点）。用于接收推理请求并发送响应。
  • Request Pool（请求池）。新到达的请求被放入请求池中，该组件负责管理系统中所有请求的生命周期。
  • Scheduler（调度器）。调度器监控请求池，负责以下任务：从池中选择一组请求，调度执行引擎对这些请求执行模型迭代；接收执行引擎返回的执行结果（即输出 token），并将每个输出 token 追加到对应请求中来更新请求池。
  • Execution Engine（执行引擎）。执行引擎是执行实际张量操作的抽象层，可在跨多个机器分布的多个 GPU 上并行化。

我们接下来看看下图中的工作流程，其中，虚线表示组件之间的交互，交互发生在执行引擎的每次迭代中。Xij 是第 i 个请求的第 j 个 token。阴影 token 表示从客户端接收到的输入 token，而非阴影 token 由 ORCA 生成。例如，请求 x 1 最初带有两个输入标记（x 11，x 12），到目前为止已经运行了两次迭代，其中第一次和第二次迭代分别生成了 x 13 和 x 14。另一方面，请求 x 3 只包含输入标记 x 31, x 32，请求 x4 包括 x41，x42，x43，因为它们还没有运行任何迭代。

工作流程分为如下几步：

  • 调度器与请求池交互，以决定下一步运行哪些请求。对应下图标号➀。
  • 调度器调用引擎为所选定的四个请求（x 1, x 2, x 3, x 4）执行一次迭代。此时，因为 x 3 和 x 4 还没有运行任何迭代，因此调度器为 x 3 移交 x 31，x 32 给执行引擎，为 x 4 移交 x 41，x42，x43 给执行引擎。对应下图标号➁。
  • 引擎对四个请求运行模型迭代，对应下图标号➂。
  • 引擎把生成的输出 token（x 15, x 23, x 33, x 44）返回给调度器，对应下图标号➃。调度器在每次引擎返回，接收该迭代的执行结果之后会检查请求是否完成。如果请求完成，请求池就会删除已完成的请求，并通知端点发送响应，返回给客户端。
  • 对于新到达的请求，在当前迭代执行完毕后，它有机会开始处理（即调度器可能选择新请求作为下一个执行对象）。因为新到达的请求只需等待一次迭代，从而显著减少了排队延迟。

ORCA 对于中止请求（Canceled Requests）并没有进行处理，实际上应该把这些请求会被及时从 Batch 中剔除并释放相应显存。

算法

下图详细描述如何在每次迭代中选择请求的算法。此算法中对 KV Cache 释放时机控制得不是很理想。在请求生成结束时就立即释放其 K/V Cache。在多轮对话场景中，这个机制会导致冗余计算，即“上一轮对话生成 K/V Cache → 释放 K/V Cache 显存 → 通过本轮对话的 Prompt 生成之前的 K/V Cache”。这样会恶化后续几轮对话的 First Token Time（产生第一个 Token 的时延）指标。

3.1.2 选择性批处理

ORCA 时代还没有 Paged Attention，因此需要 Selective Batching 来将注意力计算从 Batching 中解耦。即，为了提高计算效率，需要想办法让引擎能够以批处理方式处理任何选定的请求集。

问题

在前面分析中，我们其实做了一个简化的假设，即所有请求序列具有相同的长度。这是因为 GPU 的特殊性，如果想批量执行多个请求，每个请求的执行应该包含相同的操作，且消耗形状相同的输入张量。然而，在现实中，请求序列的长度是不同的。诚然 Padding+Masking 的方法可以解决，但严重浪费算力和显存，对于算力和显存均有限的推理 GPU 是不利的。

当使用迭代级别调度时，上述挑战会愈发加剧。因为：

  • 请求池中的请求可能具有不同的特征。
  • Prefill 和 decode 的计算方式不同。
    • Prefill 过程是长序列并行计算的，decode 过程是 token by token 的。
    • Prefill 过程不需要读取 KV cache，decode 过程需要读取 KV cache。
    • 对于 prefill，各个请求的 prompt 长度是不一致的。
    • 对于 decode，不同请求的 decode token 的 index 不一样，意味着它们计算 attention 的 mask 矩阵也不一样。
  • 迭代级调度方法可能导致同一个批处理中的不同请求的处理进度不一样，即输入张量的形状会因为已处理的 token 数量不同而不一致。

我们用上面架构图作为例子来进行分析，来看看即使对于一对请求 (xi, xj)，也不能保证它们的下一次迭代可以合并、替换为批处理版本。有三种情况导致请求对不能合并批处理：

  • 两个请求都处于初始化阶段，但输入 token 数量不同（如下图中的 x 3 和 x 4）或者说输入张量的“长度”维度不相等，因此无法将两个请求进行批处理。

  • 两个请求都处于增量阶段，但各自处理的 token 索引不同（如 x 1 和 x 2）。由于每个请求处理的 token 索引不同，导致注意力键和值的张量形状不同，因此也不能合并批处理。

  • 请求处于不同阶段：有的处于初始化阶段，有的处于增量阶段（如 x 1 和 x 3）。由于不同阶段的迭代输入 token 数量不同（初始化阶段迭代并行处理所有输入 token 以提高效率，而增量阶段每次迭代仅处理一个 token），因此无法合并批处理。

对于第二种情况，在典型的批处理机制中，每次迭代时，Transformer 层都会接收一个由批中多个[L, H]请求输入张量拼接而成的三维输入张量[B, L, H]，其中 B 是批大小，L 是一起处理的 token 数，H 是模型的隐藏大小。

但是在下图中，“iter 1”（初始化阶段）接收形状为[2,2, H]的输入张量，而“iter 2”（增量阶段）接收形状为[2,1, H]的张量。然而，在上图中，当调度器决定对批 (x 1, x 2, x 3, x 4)执行迭代时，初始化阶段请求 (x 3: [2, H]和 x 4: [3, H])的输入无法合并成一个形状为[B, L, H]的单一张量，因为 x 3 和 x 4 的输入 token 数不同，分别为2和3。

上述关于批处理的主要问题在于，前述三种情况对应于形状不规则的输入（或状态）张量，这些张量无法合并成一个大的张量并输入到批处理操作中。因此，并非所有的请求都能在任意 Iteration 被 Batching 到一起。仅当两个选定请求处于同一阶段，且（在初始化阶段）具有相同数量的输入 token 或（在增量阶段）具有相同的 token 索引时，批处理才适用。这一限制大大降低了在实际工作负载中执行批处理的可能性，因为调度器需要同时找到两个符合批处理条件的请求。

思路

解决这些问题的一个好思路是：尽量找到这些请求计算时的共同之处，使得计算能最大化合并。对于有差异的部分再单独处理。我们先以一个 transformer decode block 为例，回顾一下序列要经过哪些计算。下图是 decoder block 的各种计算类型。可以看到，Transformer decoder block 在计算上可以看做六个操作的总和：pre-proj，attn，post-proj，ffn_ln 1，ffn_ln 2，others（比如 layer normalization，activation functions，residual connection）。Transformer 输出一个形状为 [B, L, H] 的张量。其中 B 是 batch size，L 是 input tokens length，H 是模型的 embedding size。每个 token 的 KV Cache 大小均为 [1, H]。

我们把上面的介绍稍作提炼，得到如下重要信息：Transformer 层中的操作可以分为两种类型: Attention 和 non-Attention，这两种模块的算子特点不同。

  • preproj/postproj/FFN1/FFN2：这几个模块中主要是 Add、Linear、GeLU 等算子，这些算子的特点是：
    • 不需要区分 token 来自于哪个请求。因此，虽然它们是 token-wise 的，但可以使用批处理实现。
    • 和输入序列长度无关。这意味着我们可以把一个 batch 中所有的 tokens 都展平成一行进行计算（维护好各自的位置向量就好），这样不同长度的输入也可以组成 batch，从而进行计算。例如，上述 x 3 和 x 4 的输入张量可以组合成一个二维张量[ΣL, H] = [5, H]，而不需要明确的批处理维度。
    • 需要从显存读取模型权重。读取模型权重意味着我们应该尽量增大 batch size，使得一次读取能就可以造福更多请求，以此减少 IO 次数。
  • attn： 该模块的特点是：
    • 由于计算受各个序列的差异性影响（例如不同序列的 mask 矩阵不同、是否需要读取 KV cache），因此需要将序列拆分开独立处理，即 batch 维度是重要的。而由于 attn 部分本身不涉及到权重读取，因此你把序列拆分开处理，也不会在这一方面上带来额外的 IO 开销。
    • 对于注意力操作, 无论是 token-wise 还是 request-wise 的 batching 都无法执行，这是因为 Attention 操作需要一个请求的概念来计算同一请求的 token 之间的注意力。
    • 不对 Attention 层进行批处理对效率的影响很小，因为 Attention 层的操作不涉及到模型参数的重复使用，无法通过批处理来减少 GPU 内存读取。
方案

总结上述思路：Transformer Layer 里，并非所有的算子都要求组成批次的输入具有相同的形状。基于上述思路，Orca 提出了第二点核心技术： Selective batching（选择性批处理），它不是对构成模型的所有张量操作（注意力和非注意力）都进行批处理, 而是有选择地将批处理仅应用于少数非注意力操作，即对于不同类型的请求应用于不同类型的操作来解决问题，具体如下：

  • 单独处理每个注意力操作。即对于必须有相同 Shape 才能 Batching 的算子（例如 Attention）会对不同的输入进行单独的计算。
  • 对其他层（例如 MLP 层）则进行批处理。

上图展示了处理图 4 中描述的请求批 (x 1, x 2, x 3, x 4)的选择性批处理机制。示意图描述的是，系统正在生成请求 x 3、x 4 的第 1 个 Token（论文称为 Initiation Phase），同时生成请求 x 1、x 2 的后续 Token（论文称为 Increment Phase）。因为产生第 1 个 Token 后才会有 K/V Cache，所以从 Attention K/V Manager 那里看有没有 request 对应的 K/V cache，就能分辨出每个 request 处于哪个 phase。具体机制如下：

  • 应用批处理。
    • 该批有7个输入 token 要处理，因此我们将输入张量的形状设置为[7,H]，[7,H]的意义是：给定形状为 [(s1, h), (s2, h), ...] 的输入，我们将它们堆叠成一个形状为 (sum(si), h) 的大矩阵。
    • 对这个堆叠矩阵应用 Linear 这个非注意力操作。Linear 的计算不涉及 Token 之间的交互，因此将输入的 Batch 维度和 Seq 维度 Reshape 成 1 个维度，即可完成 Linear 的 Batch 计算。
  • 将密集层的结果分割回 [(s1, h), (s2, h), ...]。在注意力操作之前，我们插入一个拆分（split）操作，目的是把 batch 中的每个请求的张量区分出来。（注：有的读者可能会疑惑为什么 Linear 的输入和输出维度分别是[7, H]和[7,3 H]，为什么升维了？这是因为通过矩阵拼接可以实现用一个 Linear 同时计算出 QKV，即输出包含了有 H 维的 Q、H 维的 K 和 H 维的 V，合计3H 维）。
  • 对拆分后的张量中的每个请求分别运行注意力操作。Attention 的计算涉及 Token 之间的交互，OCRA 的操作是：将输入按 Batch 维 Split，每个样本分别计算 Attention。
  • 然后，通过合并操作将注意力操作的输出合并回形状为[7, H]的张量，以恢复对其他操作的批处理功能。

虽然 ORCA 是把 Prefill 和 decode 阶段合并在一起处理，但其实它们的计算模式差很多，分开做能够更好地做优化（例如计算 Linear 时，Initiation Phase 的计算 Kernel 更接近 GEMM，Increment Phase 的计算 Kernel 更接近 GEMV）

后续有的工作就是将 Initiation Phase 的请求批量做完计算后，再这些请求与已在 decode 阶段的请求合并做批量计算，有的框架通过超参数来管理这个问题：等待已服务比和等待结束序列标记的请求比（waiting_served_ratio）。

3.1.3 流水线并行

另外，ORCA 的调度器实现了引擎中工作线程对多个批次的流水线执行。

ORCA 的调度器使引擎中 worker 的执行能够跨多个批次进行流水线处理。若当前已调度批次数 n_scheduled 达到工作线程数 n_workers（算法1的第9-10行）时，调度器就不再等待。通过这种方式，调度器保持引擎中并发运行的批次数为 n_workers，这意味着引擎中的每个工作线程都在处理一个批次，而不会处于空闲状态。

下图展示了 3 个 ORCA 工作线程的执行流水线，最大批次大小为 2。假设请求 A 先于 B 到达，B 先于 C 到达，以此类推。首先，调度器根据到达时间选择请求 A 和 B，并调度引擎来处理包含 A 和 B 的批次（我们称之为 AB 批次），其中 Worker 1、Worker 2 和 Worker 3 依次处理该批次。只有在调度器注入另外两个批次 CD 和 EF 之后，调度器才等待 AB 批次返回。一旦 AB 批次返回，请求 A 和 B 会再次被选中并调度，因为它们是请求池中最早到达的请求。