KaiSpace
tech

Chunked Prefill详解

LLM 推理通常包含两个阶段:

  • Prefill:一次处理大量 prompt token,矩阵规模大,通常偏 compute-bound。
  • Decode:每个 request 每轮只产生一个 token,计算规模小,但需要读取模型权重和大量 KV Cache,通常偏 memory-bound。

这会造成一个很直接的问题:

Prefill 很容易把 GPU 算力吃满,而 decode 即使延迟很重要,却经常无法充分利用 GPU。

Chunked Prefill 的核心思想,就是把一个很长的 prefill 拆成多个小块,然后让这些 prefill chunk 和正在运行的 decode request 一起执行,尽量增加GPU的利用率。


1. Chunked Prefill 要解决什么问题?

假设来了一个 4096-token prompt。

传统做法可能是:

iteration 1: decode
iteration 2: decode
iteration 3: [------------ 4096-token Prefill ------------]
iteration 4: decode
iteration 5: decode

这个大 prefill 会占用 GPU 很长时间。

正在 decode 的 request 必须等待,因此会出现明显的 decode stall,Inter-Token Latency(ITL)恶化。

另一方面,如果为了保护 decode latency,始终优先执行 decode:

decode
decode
decode
decode
...

又会产生另一个问题:

decode 本身通常没有充分利用 GPU 的计算资源。

因此 Chunked Prefill 把:

40964096

拆成例如:

512+512+512+512+512+512+\cdots

然后变成:

iteration 1:
    decode requests + 512 prefill tokens

iteration 2:
    decode requests + 512 prefill tokens

iteration 3:
    decode requests + 512 prefill tokens

这样既避免一个巨大 prefill 长时间阻塞 decode,也可以利用 decode iteration 中没有充分使用的 GPU 资源。

这也是 Sarathi 最早强调的 “piggyback decodes with chunked prefills” 思路。


2. Scheduler 是怎么做的?

以目前 vLLM 的策略为例,每个 iteration 有一个 token budget:

B=max_num_batched_tokensB=\texttt{max\_num\_batched\_tokens}

scheduler 首先放入需要 decode 的 token。

例如现在有 32 个 decode request:(如果没有32个decode request,那么下面一个任务也按照现在有多少request发射,不会等待一定要32个request)

32×1=3232\times1=32

假设:

B=512B=512

那么还剩:

51232=480512-32=480

于是 scheduler 可以从某个 prefill request 中取出 480 个 token:

Batch:

32 decode tokens
+
480 prefill tokens
----------------
512 tokens

如果 prefill 还有 3000 个 token 没处理,则留到后面的 iteration。

目前 vLLM 的 chunked-prefill scheduler 也是类似逻辑:优先调度 decode,然后使用剩余 token budget 调度 prefill;放不下的 prefill 会被自动切 chunk。

所以 Chunked Prefill 首先是一个 scheduling mechanism

把一个 request 的 prefill 从“一次必须全部完成”,变成“允许跨多个 iteration 增量完成”。


3. Decode 和 Prefill 为什么能放进同一个 batch?

假设当前有:

Request A: decode       1 token
Request B: decode       1 token
Request C: prefill    128 tokens

虽然这是三个 request,但对于很多 Transformer 算子来说,可以直接把 token 拼起来:

X=[XAXBXC]R130×dX= \begin{bmatrix} X_A\\ X_B\\ X_C \end{bmatrix} \in\mathbb{R}^{130\times d}

例如 Q projection:

Q=XWQQ=XW_Q

可以直接执行一次:

(130,d)(d,h)(130,d)(d,h)

而不需要分别执行:

XAWQ,XBWQ,XCWQX_AW_Q,\quad X_BW_Q,\quad X_CW_Q

这样节省了每次读WQW_Q的开销。


4. Attention 怎么处理不同 request?

multi-head-attention.png

Attention 稍微复杂一些。

三个 request 的 Q 可以放在一起,但它们对应不同的 KV Cache:

Q:

[A][B][--------- C ---------]
 │  │             │
 ▼  ▼             ▼
KV_A KV_B         KV_C

例如:

A:
q_len  = 1
kv_len = 1001

B:
q_len  = 1
kv_len = 501

C:
q_len  = 128
kv_len = 384

因此 attention kernel 还需要 metadata,例如:

request_id
q_offset
q_len
kv_offset
kv_len
page_table

根据这些metadata,GPU会给每个request分配block,满足以下几个约束(下一节会仔细讨论):

  1. 一个block最多处理QblockQ_{block}个token
  2. 一个block最多处理lenkvlen_{kv}的KV Cache
  3. 一个block处理一个head

从 GPU 的角度看,它并不是:

request A
request B
request C

而是进一步 materialize 成很多 block-level task:

Request A
 ├─ block
 ├─ block
 └─ block

Request B
 ├─ block
 ├─ block
 └─ block

Request C
 ├─ block
 ├─ block
 ├─ block
 ├─ ...
 └─ block

最后形成一个大的 CUDA grid:

[A tasks][B tasks][------------- C tasks -------------]

GPU 再把这些 blocks 动态分配给各个 SM。

所以:

逻辑上存在多个 request,但并不意味着物理上必须启动多个 kernel。


5. 每个 Request 会产生多少 Block?

主要取决于 attention kernel 的 tiling。

对于 prefill,可以粗略理解为:

NblocksqlenQblock×NheadsN_{\text{blocks}} \approx \left\lceil \frac{q_{\text{len}}}{Q_{\text{block}}} \right\rceil \times N_{\text{heads}}

因为 prefill 有很多 Q,因此天然可以沿 Q sequence 切分。

例如:

q_len = 128
Q_tile = 16
heads = 32

blocks ≈ 8 × 32 = 256

而 decode:

qlen=1q_{\text{len}}=1

没有足够的 Q 可以切。

因此长 context 下还可能沿 KV 方向 split:

NblocksNheads×NKV splitsN_{\text{blocks}} \approx N_{\text{heads}} \times N_{\text{KV splits}}

例如:

Q = 1 token

KV:
[0:4096]
[4096:8192]
[8192:12288]
...

多个 block 分别算 partial attention,最后再 reduction。这两种切分是正交的,可以同时apply。

因此 mixed batch 最终可以理解成:

             Scheduler
                 │
                 ▼
        requests + token budget
                 │
       ┌─────────┴─────────┐
       │                   │
    Decode              Prefill
   q_len = 1          q_len > 1
       │                   │
       └─────────┬─────────┘
                 ▼
          block-level tasks
                 │
                 ▼
     [D][D][D][P][P][P][P]...
                 │
                 ▼
                GPU

不过这里有一个问题:对于一个prefill来说,Q本身就切分了很多block,occupancy已经非常高了,如果再切KV就没有意义了。这里有一个问题就是每个Q都要load一遍KV Cache,这个的确是一个问题,所以要找到一个 BqB_q 的sweet point让compute也能打满。


6. Chunked Prefill 当前最大的痛点

Chunked Prefill 解决了“能不能混合”的问题,但并没有解决“怎么混合才最优”。

痛点一:Chunk Size 很难选

chunk 太大:

decode + [------------- huge prefill -------------]

prefill 仍然会显著拉长整个 iteration,decode ITL 变差。

chunk 太小:

decode + [prefill]
decode + [prefill]
decode + [prefill]

虽然 decode latency 好,但:

  • GPU资源打不满
  • TTFT 可能上升。
  • 每次不同seq还要再load KV。虽然本身Q_block就要多次load KV,这两个的关系比较复杂。

因此存在明显的:

TBTTTFT / throughput\text{TBT} \leftrightarrow \text{TTFT / throughput}

trade-off。

另外KV Cache length不同我们prefill和decode刚好对齐的配比也不同。

vLLM 官方也明确指出,较小的 max_num_batched_tokens 通常有利于 TBT,而较大的值更有利于 TTFT 和吞吐。


痛点二:Prefill 和 Decode 并不是完美互补

我们经常简单地说:

Decode memory-bound,Prefill compute-bound,所以把它们放一起就好了。

但真实 GPU 没这么简单。

两者仍然会竞争:

  • SM
  • registers
  • shared memory
  • L2
  • HBM bandwidth
  • Tensor Core
  • instruction issue slots

例如 decode attention 正在大量访问 KV Cache,而 prefill attention 同样需要访问 KV。

因此:

decode unused compute\text{decode unused compute}

并不等于:

prefill 可以免费使用的 compute\text{prefill 可以免费使用的 compute}

Chunked Prefill 是一种 co-location,不是严格意义上的资源隔离或资源填洞。


痛点三:需要每层对齐

prefill和decode其实走的是同一个路径所以不存在两者各做各的,每一次attention和ffn都要进行对齐。所以在做PD multiplexing的时候实际上他们是用layerwise chunk做的。这个好处是通过prefill和decode发射不同的kernel,让prefill和decode完全解耦。当然我们也可以让PD multiplexing用seqlen chunk,只要两个发射不同kernel就可以了。

Comments

No comments yet.