预填充与解码分离
主线位置:「一个请求的一生」从调度到前向这几段被拆到两个实例上:预填充实例算 prompt、写 KV,解码实例收 KV、逐 token 生成。引擎内部的分离逻辑两边都是 Python;KV 的搬运交给 NIXL、Mooncake 等传输库;把一个请求同时或先后发给两个实例的路由层是 Rust(SGLang model gateway、vllm-project/router)。
版本:vLLM v0.32.0、SGLang v0.5.22;路由层对照 SGLang 仓库里的
sgl-model-gateway与 vllm-project/router v0.1.16
先说结论:
- 为什么分。 预填充算力受限,解码访存受限;放在同一个实例里,长 prompt 的预填充步会拉长或打断别人的解码步。分开后两边各自批、各自调参,代价是每个请求要把整段 prompt 的 KV 从一个实例搬到另一个:Llama-3.1-8B 一个 2048 token 的 prompt 约 256 MiB(估算)。
- vLLM 用 KV connector。 一个抽象类分成调度器一侧和 worker 一侧,按名字注册了十几种实现,用
--kv-transfer-config的 JSON 选。默认不开。PD 常用NixlConnector(拉取式)。流程是串行的:代理先把请求发给预填充实例(max_tokens=1),拿回响应里的kv_transfer_params,再带着它发给解码实例,解码实例从预填充实例读 KV。 - SGLang 用 bootstrap room。 用
--disaggregation-mode prefill|decode起两类服务;预填充实例在主进程里起一个 bootstrap server(默认端口 8998)。路由器给请求生成一个随机的bootstrap_room,同时发给两边;解码端先分配好 KV、把块号连同 room 发给预填充端,预填充端凭 room 认出是同一个请求,算完把 KV 写过去。传输后端默认mooncake。 - 路由层双发。 SGLang model gateway 的 PD router 往请求里注入
bootstrap_host/bootstrap_port/bootstrap_room,同时发出两个 HTTP 请求。vllm-project/router 的VllmPDRouter对 NIXL 拉取模式先发预填充、再发解码,只有 NIXL 推送模式才并发。路由层的细节在路由层一页。 - 效果(TTFT、TPOT、吞吐、传输耗时)全部待实测(C3-7)。
为什么分
为什么服务大模型难 里算过:预填充的算术强度约等于 prompt 长度,解码的算术强度约等于批大小。两种负载在同一个实例里,有两个后果:
- 互相干扰。 一步里混了预填充块,这一步就变长,同批的解码请求这一个 token 等得也长。分块预填充 把这一步的长度封了顶,但没有消除:vLLM 每步仍是「解码 + 至多 2048 个预填充 token」,SGLang 默认配置下预填充块之间解码整步停。
- 只能用一套参数。 预填充想要大 token 预算,解码想要大并发和多 KV。同一个实例的
max_num_batched_tokens、显存比例只能取一个折中。
分开之后,预填充实例只跑预填充,解码实例只跑解码;两边可以用不同的卡数、不同的并行方式,按负载各自加副本。
代价是 KV 要搬。按 Llama-3.1-8B 每 token 128 KiB 算(估算,脚本只用 Python 标准库):
KV per token = 128 KiB
prompt 512: KV to transfer = 64 MiB
prompt 2048: KV to transfer = 256 MiB
prompt 8192: KV to transfer = 1024 MiB
这些数据要在预填充算完到解码第一步之间送到。走 RDMA 还是 TCP、两边 TP 是否相同,都会影响这段时间。另外两个实例各自的 KV 池要单独规划,解码实例的前缀缓存也不像单实例那样自然(SGLang 解码端默认关掉基数树,见下文)。
flowchart LR
subgraph one["单实例"]
s1["步 n:解码 B 个 + 预填充块"] --> s2["步 n+1:解码 B 个 + 预填充块"]
end
subgraph two["分离"]
p["预填充实例<br/>只跑预填充"] -->|"KV(每 token 128 KiB,估算)"| d["解码实例<br/>只跑解码"]
end
vLLM:KV connector
接口
KVConnectorBase_V1 是一个抽象类,同一个类的实例在两个地方各建一个,角色不同:
vllm/distributed/kv_transfer/kv_connector/v1/base.py#L1
"""KVConnectorBase_V1 Class for Distributed KV Cache & Hidden State
communication in vLLM v1
The class provides the following primitives:
Scheduler-side: runs in the scheduler, binds metadata, which
is used by the worker-side to load/save KV cache.
get_num_new_matched_tokens() - get number of new tokens
that exist in the remote KV cache. ...
update_state_after_alloc() - update KVConnector state after
temporary buffer alloc by the CacheManager.
# ...
request_finished() - called once when a request is finished,
with the computed kv cache blocks for the request.
Returns whether KV cache should be freed now or if the
connector now assumes responsibility for freeing the
the blocks asynchronously. ...
Worker-side: runs in each worker, loads/saves KV cache to/from
the Connector based on the metadata.
# ...
start_load_kv() - starts loading all KVs (maybe async)
wait_for_layer_load() - blocks until layer i load is done
save_kv_layer() - starts saving KV for layer i (maybe async)
wait_for_save() - submits saves and performs required synchronization
- 调度器一侧在 EngineCore 的
Scheduler里建(scheduler.py#L176)。它回答「这个请求有多少 token 的 KV 在远端」(get_num_new_matched_tokens,base.py#L501),以及请求结束时块要不要先别释放(request_finished,#L616)。 - worker 一侧在每个 worker 里建(kv_transfer_state.py#L90),前向前后被 model runner 调用(kv_connector_model_runner_mixin.py#L85),按调度器给的元数据真正搬 KV,可以按层边算边传。
接口对上的是「远端有一段 KV」这件事,不只用于 PD:把 KV 卸到 CPU、存进外部 KV 存储,也是同一个接口。远端命中的请求,调度器先给它分好本地块,状态设为 WAITING_FOR_REMOTE_KVS,等 worker 报告收完再参加调度(scheduler.py#L1323)。
配置
vllm/config/kv_transfer.py#L57
class KVTransferConfig:
"""Configuration for distributed KV cache transfer."""
kv_connector: str | None = None
# ...
kv_role: KVRole | None = None
"""Whether this vLLM instance produces, consumes KV cache, or both. Choices
are 'kv_producer', 'kv_consumer', and 'kv_both'."""
# ...
kv_connector_extra_config: dict[str, Any] = field(default_factory=dict)
# ...
kv_connector_module_path: str | None = None
命令行是 --kv-transfer-config '<JSON>'(arg_utils.py#L1843),EngineArgs.kv_transfer_config 默认 None,也就是默认不建任何 connector。上游的 NIXL 集成测试给预填充实例配 {"kv_connector":"NixlConnector","kv_role":"kv_producer"},解码实例配 kv_consumer(run_accuracy_test.sh#L59)。kv_connector_module_path 可以指向自己写的模块,优先于内置注册表(factory.py#L107)。
有哪些实现
KVConnectorFactory 里按名字注册(factory.py#L153 起):
| 名字 | 用途 |
|---|---|
NixlConnector | PD:经 NIXL 传 KV;是 NixlPullConnector 的别名(nixl/connector.py#L414),解码端 READ |
NixlPushConnector | PD:预填充端 WRITE 到解码端(#L373) |
MooncakeConnector、MoRIIOConnector | PD:经 Mooncake 传输引擎、AMD 的 MoRI-IO |
LMCacheConnectorV1、LMCacheMPConnector、MooncakeStoreConnector、FlexKVConnectorV1、HF3FSKVConnector | 外部 KV 存储 / 缓存层 |
OffloadingConnector、SimpleCPUOffloadConnector | 把 KV 卸到 CPU |
MultiConnector | 组合几个:从第一个有命中的读,往所有的写(multi_connector.py#L135) |
ExampleConnector、ExampleHiddenStatesConnector、DecodeBenchConnector、HiSparseConnector | 示例(存到磁盘)、隐状态、解码压测用的假 KV、稀疏注意力 |
NIXL 下的一次请求
代理(或路由器)先把请求改成 max_tokens=1、非流式,带上 kv_transfer_params: {"do_remote_decode": true} 发给预填充实例。预填充实例照常跑完这个请求;结束时,connector 的 request_finished 不让调度器立刻释放块,而是返回一组参数,随 HTTP 响应带回去:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py#L278
return delay_free_blocks, dict(
do_remote_prefill=is_p_node,
do_remote_decode=is_d_node,
remote_block_ids=block_ids,
remote_engine_id=self.engine_id,
remote_request_id=request.request_id,
remote_host=self.side_channel_host,
remote_port=self.side_channel_port,
# ...
代理把这组参数原样塞进原请求,发给解码实例。解码实例的调度器看到 do_remote_prefill,在 get_num_new_matched_tokens 里报告「整段 prompt 都在远端、异步加载」(pull_scheduler.py#L59),分好本地块后进入 WAITING_FOR_REMOTE_KVS;worker 经 NIXL 从预填充实例的显存里读这些块。NIXL 的握手走一个旁路端口,VLLM_NIXL_SIDE_CHANNEL_PORT 默认 5600(envs.py#L1739)。
预填充端的块不会一直留着:解码端读完会通知它释放;没人来读时,按 kv_lease_duration 到期释放,默认 30 秒,可在 kv_connector_extra_config 里改(base_scheduler.py#L82)。
sequenceDiagram participant C as 客户端 participant R as 代理 / vllm-router participant P as 预填充实例 participant D as 解码实例 C->>R: POST /v1/chat/completions R->>P: 原请求,max_tokens=1,stream=false<br/>kv_transfer_params.do_remote_decode=true P->>P: 预填充,块延迟释放 P-->>R: 响应带 kv_transfer_params<br/>(remote_block_ids、remote_host、remote_port) R->>D: 原请求 + P 返回的 kv_transfer_params D->>D: get_num_new_matched_tokens:全部在远端<br/>分块,WAITING_FOR_REMOTE_KVS D->>P: NIXL 握手(旁路端口)、READ KV 块 D->>P: 读完,通知释放 D->>D: 进入解码 D-->>C: 流式返回
这条路是串行的:解码实例在预填充完成、响应回到代理之后才收到请求,分配 KV 和建立连接都在这之后。上游的多轮示例代理就是这么写的(disagg_proxy_multiturn.py#L432)。
SGLang:bootstrap server 与 bootstrap room
配置
python/sglang/srt/arg_groups/fields/disagg.py#L55
disaggregation_mode: A[
Literal["null", "prefill", "decode"],
# ...
] = "null"
disaggregation_decode_allocation_policy: A[
Literal["early", "prefill_complete"],
# ...
] = "early"
disaggregation_transfer_backend: A[
str,
# ...
] = "mooncake"
# ...
disaggregation_bootstrap_port: A[
int, "Bootstrap server port on the prefill server. Default is 8998."
] = 8998
--disaggregation-mode:默认null,不分离;prefill或decode起对应的服务。--disaggregation-transfer-backend:默认mooncake,可选nixl、mori、ascend、fake、mooncake_tcp(choices.py#L122)。mooncake_tcp不是独立后端:启动时设MC_FORCE_TCP=1,改写成mooncake,不走 RDMA(pd_disaggregation_hook.py#L30)。fake不搬 KV,只给解码端压测用,预填充端不接受。- 每个后端提供一组四个类:Manager、Sender、Receiver、BootstrapServer(disaggregation/utils.py#L676)。
- 解码端默认关掉基数树(
disable_radix_cache=True,pd_disaggregation_hook.py#L132);要开得加--disaggregation-decode-enable-radix-cache,标为实验。
bootstrap server
只在预填充实例上起(开了角色互换时每个实例都起),由 TokenizerManager 初始化时调用,所以在主进程里,是一个跑 aiohttp 的后台线程:
python/sglang/srt/managers/disagg_service.py#L21
def start_disagg_service():
# Start kv bootstrap server on prefill
# ...
if start_bootstrap:
kv_bootstrap_server_class = get_kv_class(
transfer_backend, KVClassType.BOOTSTRAP_SERVER
)
bootstrap_server = kv_bootstrap_server_class(
host=get_serving().host,
port=get_disagg().disaggregation_bootstrap_port,
)
它只是一个注册表(common/conn.py#L2614):预填充实例的每个 rank 启动时 PUT /route 登记自己的 IP、端口和并行布局;解码端 GET /route 查「该连预填充的哪个 rank」。注册数没到 DP × CP × TP × PP 之前,查询返回 503(_is_ready,#L2652)。KV 不经过它。
打开 SGLANG_RUST_SERVER=1 时,这个注册表由嵌在调度器进程里的 Rust 服务提供,直接挂在 API 端口上,启动时把 disaggregation_bootstrap_port 改成 API 端口(pd_disaggregation_hook.py#L251)。这是 PD 路径上引擎内唯一的一段 Rust,且是可选的。
bootstrap room 怎么配对
请求里多三个字段:bootstrap_host、bootstrap_port、bootstrap_room(io_struct.py#L291)。前两个指向预填充实例的 bootstrap server,room 是路由器生成的随机 63 位整数,同一个请求在两边带的是同一个 room。两边各自把请求放进队列:
- 预填充端(prefill.py#L1):为请求建一个
KVSender,记下 room(prefill.py#L373),进 bootstrap 队列。状态从Bootstrapping变成WaitingForInput(base/conn.py#L118)之后,请求才进等待队列参加调度。也就是说,默认配置下预填充要等解码端准备好才开算(optimistic_prefill_attempts默认 0)。算完一块就发一块的 KV,最后进 inflight 队列等传完。 - 解码端(decode.py#L1):为请求建一个
KVReceiver(decode.py#L875),向 bootstrap server 查到预填充端对应 rank 的地址,有空余 KV 时预先分配好整段 prompt 的位置,把块号连同 room 发过去(send_metadata,以 Mooncake 为例在 mooncake/conn.py#L3361)。收完后不再跑预填充前向,只补元数据,并入运行批开始解码。
预填充端有一个线程专门收这种消息,按 room 归档:
python/sglang/srt/disaggregation/mooncake/conn.py#L2943
required_dst_info_num = int(waiting_req_bytes[7].decode("ascii"))
room = int(room)
if room not in self.transfer_infos:
self.transfer_infos[room] = {}
self.transfer_infos[room][mooncake_session_id] = (
TransferInfo.from_zmq(waiting_req_bytes)
)
# NOTE: after bootstrapping we can mark the req as waiting for input
if len(self.transfer_infos[room]) == required_dst_info_num:
两边的 TP 可以不同,解码端可能有几个 rank 要同一份 KV 的不同部分,所以要等齐 required_dst_info_num 份才算准备好。
room 还决定 DP 组。预填充实例 DP > 1 时,load_balance_method 在 auto 下解析成 follow_bootstrap_room(见 张量并行),按 room % DP 组数 选组(data_parallel_controller.py#L831);解码端用同一个公式算出该连哪个组,不必再查(decode.py#L871)。直接把请求发给预填充实例而不带 room,会在这里断言失败,报错信息叫你走路由器。
两边等对方都有上限:bootstrap 和等待传输各 300 秒(SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT、SGLANG_DISAGGREGATION_WAITING_TIMEOUT,environ.py#L706)。
sequenceDiagram
participant C as 客户端
participant G as SGLang model gateway
participant PB as 预填充 bootstrap server
participant P as 预填充 Scheduler
participant D as 解码 Scheduler
Note over P,PB: 启动时每个 rank PUT /route 登记
C->>G: POST /v1/chat/completions
G->>G: 选一对 P / D,生成 bootstrap_room
par 同时发出
G->>P: 请求 + bootstrap_host / port / room
and
G->>D: 同一请求 + 同一 room
end
P->>P: 建 KVSender,状态 Bootstrapping
D->>PB: GET /route(首次,之后缓存)
PB-->>D: 预填充各 rank 的地址
D->>D: 建 KVReceiver,预分配 KV
D->>P: room + 本地块号(ZMQ)
P->>P: 按 room 配上,WaitingForInput,进等待队列
P->>P: 预填充,按块写 KV
P->>D: KV 写入解码端显存(Mooncake / NIXL)
P-->>G: 预填充响应
D->>D: 传输完成,并入运行批解码
D-->>G: 流式输出
G-->>C: 流式返回
和 vLLM 的串行流程比:SGLang 两边同时收到请求,解码端的分配、握手和预填充端排队可以重叠;预填充端默认等解码端分好 KV 再算,解码端的 KV 容量会反过来卡住预填充。vLLM 的预填充端不等谁,算完把块留 30 秒等人来读。
路由层怎么双发
两条路由都是 Rust 写的独立进程,这里只看 PD 一处,其余(负载策略、缓存感知、服务发现)在路由层一页。
SGLang model gateway(sgl-model-gateway)。PD 模式用 --pd-disaggregation 启动,--prefill 后面跟 URL 和该实例的 bootstrap 端口(main.rs#L99)。每个请求选一对实例后,往 JSON 里注入三个字段:
sgl-model-gateway/src/routers/http/pd_router.rs#L242
fn inject_bootstrap_into_value(
mut original: Value,
prefill_worker: &dyn Worker,
batch_size: Option<usize>,
) -> Result<Value, String> {
// ...
obj.insert(
Self::BOOTSTRAP_HOST_KEY.to_string(),
Value::from(prefill_worker.bootstrap_host()),
);
// ...
obj.insert(
Self::BOOTSTRAP_ROOM_KEY.to_string(),
Value::from(super::pd_types::generate_room_id()),
);
room 取 rand::random::<u64>() & i64::MAX,和 Python 端 random.randint(0, 2**63 - 1) 的范围一致(pd_types.rs#L60)。批量请求每条一个 room。然后两个 HTTP 请求一起发,同一个任务里 select! 两个 future,等预填充先回来;预填充失败就丢掉解码请求,免得解码端干等 300 秒(pd_router.rs#L710)。客户端拿到的是解码端的响应。
vllm-project/router。它是 SGLang 路由器的分支,但 PD 模式(--vllm-pd-disaggregation)由工厂建 VllmPDRouter,按 vLLM 的 kv_transfer_params 协议走(factory.rs#L22)。--kv-connector 默认 nixl,可选 mooncake、moriio(types.rs#L536)。预填充请求被改成 max_tokens=1、stream=false(vllm_pd_router.rs#L578)。NIXL 默认的拉取模式下是两段串行:等预填充响应,取出里面的 kv_transfer_params 转给解码;只有预填充实例用 NixlPushConnector、路由器已缓存它的身份信息时,才构造两份请求并发出去(vllm_pd_router.rs#L1680)。Mooncake 模式下它去查预填充实例的 bootstrap 服务;vLLM 的 MooncakeConnector 也有一个 bootstrap 端口,VLLM_MOONCAKE_BOOTSTRAP_PORT 默认 8998,和 SGLang 的默认值相同(envs.py#L1783)。
flowchart TB
req["客户端请求"] --> which{"哪个路由器"}
which -->|"SGLang model gateway"| inj["注入 bootstrap_host / port / room"]
inj --> both["同时发 P 和 D"]
which -->|"vllm-router,NIXL 拉取(默认)"| p1["先发 P:max_tokens=1"]
p1 --> p2["取 P 响应里的 kv_transfer_params"]
p2 --> d1["再发 D"]
which -->|"vllm-router,NIXL 推送"| both2["构造两份请求,同时发"]
待实测
C3-7 在多卡机器上做:
- 同一个 8B 模型,1P1D(各一张卡)对两张卡的单实例 DP,在长短 prompt 混合负载下的 TTFT、TPOT 和总吞吐;两个引擎各一组。
- KV 传输耗时:prompt 512 / 2048 / 8192,走 RDMA 和走 TCP(SGLang
mooncake_tcp)各一次;对照上面 64 MiB / 256 MiB / 1 GiB 的估算。 - vLLM 串行流程比 SGLang 并发双发多出来的 TTFT。
小结
- 分离的理由是预填充和解码的瓶颈不同、放一起互相拖;代价是每个请求整段 prompt 的 KV 要跨实例搬,Llama-3.1-8B 2048 token 约 256 MiB(估算)。
- vLLM:
KVConnectorBase_V1分调度器一侧和 worker 一侧,用--kv-transfer-config选实现,默认不开;NixlConnector是拉取式的别名。代理先发预填充(max_tokens=1),把响应里的kv_transfer_params转给解码,解码端从预填充端读 KV,预填充端的块默认留 30 秒。 - SGLang:
--disaggregation-mode prefill|decode,传输后端默认mooncake;预填充实例主进程里起 bootstrap server(默认 8998)只做注册表;同一个bootstrap_room让两边认出同一个请求,解码端先分 KV、发块号,预填充端默认等它再算,按块写过去。解码端默认关基数树。 - 路由层:SGLang model gateway 注入三个 bootstrap 字段后同时发两边;vllm-project/router 对 NIXL 拉取模式两段串行,推送模式才并发。两者都是 Rust,引擎内的分离逻辑都是 Python。
- 效果全部待实测(C3-7)。