Chunked Prefill 与 Preemption

Chunked Prefill：改变调度粒度，不改变 causal 结果

一次性 prefill 让长度 $T$ 的 Q/K/V 同时进入 causal attention。Chunked prefill 把 Query positions 切成 chunks；处理 [a,b) 时，K/V cache 已包含 0:b，每个 query 仍只能看到绝对位置不大于自身的位置：

M_{i,j}=\begin{cases}0,&j\le i\\-\infty,&j>i\end{cases}, \quad i\in[a,b),\ j\in[0,b)

只要 position ids、RoPE、mask 与 KV append 正确，拼接所有 chunk outputs 应等于 full prefill。它允许 scheduler 在 chunks 之间插入 decode 或其他请求，改善公平性/TPOT，却增加更多 launches、调度边界，且后续 chunks 要读取更长历史。

Preemption：释放哪种状态，未来付什么成本

当 Continuous Batching 的 KV/token budget 不足，runtime 可以暂停低优先级请求：

Recompute：释放 GPU KV blocks；恢复时从 token ids 重新 prefill 已处理 prefix。立即回收容量，未来支付计算与 TTFT/完成延迟。
Swap/Offload：把 KV payload 搬到 CPU/其他层级；恢复时搬回。保存计算，支付 PCIe/内存带宽、CPU 容量和同步。
Reject/queue：尚未开始的请求不分配状态；不是 active-request preemption。

可运行等价性与成本核算

import math,platform,torch
torch.manual_seed(127)
DEVICE=torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE=torch.float32

def dense_causal(q,k,v):
    t=q.shape[-2]; scores=q@k.transpose(-2,-1)/math.sqrt(q.shape[-1])
    mask=torch.triu(torch.full((t,t),-torch.inf,device=DEVICE),1)
    return torch.softmax(scores+mask,-1)@v

def chunked_causal(q,k,v,chunk_size):
    outputs=[]; k_cache=v_cache=None; trace=[]; total=q.shape[-2]
    for start in range(0,total,chunk_size):
        end=min(start+chunk_size,total)
        k_cache=k[...,start:end,:] if k_cache is None else torch.cat((k_cache,k[...,start:end,:]),-2)
        v_cache=v[...,start:end,:] if v_cache is None else torch.cat((v_cache,v[...,start:end,:]),-2)
        scores=q[...,start:end,:]@k_cache.transpose(-2,-1)/math.sqrt(q.shape[-1])
        qi=torch.arange(start,end,device=DEVICE)[:,None]
        kj=torch.arange(end,device=DEVICE)[None,:]
        scores=scores.masked_fill(kj>qi,-torch.inf)
        outputs.append(torch.softmax(scores,-1)@v_cache)
        trace.append((start,end,tuple(k_cache.shape)))
    return torch.cat(outputs,-2),trace

def preemption_accounting(cached_tokens,block_size,kv_bytes_per_token,cpu_bandwidth_gbps):
    blocks=(cached_tokens+block_size-1)//block_size
    payload=cached_tokens*kv_bytes_per_token
    return {"released_blocks":blocks,"recompute_tokens":cached_tokens,
            "swap_roundtrip_ms":2*payload/(cpu_bandwidth_gbps*1e9)*1e3}

def main():
    q=torch.randn(1,2,7,4,device=DEVICE,dtype=DTYPE)
    k=torch.randn_like(q); v=torch.randn_like(q)
    dense=dense_causal(q,k,v); chunked,trace=chunked_causal(q,k,v,3)
    torch.testing.assert_close(chunked,dense,rtol=1e-5,atol=1e-6)
    assert trace==[(0,3,(1,2,3,4)),(3,6,(1,2,6,4)),(6,7,(1,2,7,4))]
    accounting=preemption_accounting(7,4,512*1024,24)
    assert accounting["released_blocks"]==2 and accounting["recompute_tokens"]==7
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE} Q=K=V={tuple(q.shape)} chunk_size=3")
    print(f"chunk_trace={trace}")
    print(f"chunked_vs_full_max_abs_diff={(chunked-dense).abs().max():.3e}")
    print(f"preemption_accounting={accounting}")

if __name__=="__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/chunked_prefill_preemption.py。full/chunked 最大差 2.980e-08。示意每 token KV 512 KiB、7 cached tokens、block size 4：recompute 可释放 2 blocks，恢复重算 7 tokens；假设 host 链路 24 GB/s，纯 payload round-trip 下界约 0.306 ms。

该 swap 数字是公式代入，不是实测 PCIe benchmark：真实传输还包括 page pinning、协议、并发、block 对齐与调度；小示意 payload 也不代表真实模型每 token KV。

调度权衡与失败场景

chunk 太小：launch/调度次数增加，prefill kernel 形态变小，TTFT 可能变差；
chunk 太大：长 prompt 仍会阻塞 decode，TPOT 抖动；
recompute 频繁：重复 prefill 放大 GPU 工作，形成 preemption thrashing；
swap 频繁：PCIe/CPU memory 成为瓶颈，恢复延迟不可预测；
优先级不带 aging：长/低优先级请求可能 starvation；
释放 blocks 与 graph slots、collective 顺序不同步：会产生地址/并发错误。

上下游关系

Prefill/Decode定义 full prefill reference 与 TTFT/TPOT。
KV Cache提供要释放、重算或换出的状态 payload。
PagedAttention以 blocks 为分配/回收单位；最后 block 的碎片影响实际释放量。
Continuous Batching决定 chunk budget、admission、priority 与 preemption 时机。
Prefix Cache可能让 recompute 恢复命中一部分 prefix，但命中 blocks 仍占容量且可能被驱逐。
KV Offload 与 P/D Disaggregation实际测量本机 pinned CPU round-trip，并展开跨 worker handoff 的 bytes、ownership 与一致性协议。

Chunked Prefill：改变调度粒度，不改变 causal 结果

Preemption：释放哪种状态，未来付什么成本

可运行等价性与成本核算

调度权衡与失败场景

上下游关系

参考资料