06-PD 分离问题反馈

[cr7258]

在 PD 分离场景下，prefill worker 和 decode worker 的 TP size 要一样吗？

不一定，只要能够处理不同 TP size 之间的 kv layout 的转换，比如 Dynamo 的 block_copy.cu kernel 会做这个事情：

For decode and prefill with different KV layouts (i.e., due to different TP), Dynamo applies a high-performance kernel that transposes the KV blocks into their matching layout in the KV receiver after the NIXL reads and before the NIXL writes. https://github.com/ai-dynamo/dynamo/blob/main/docs/design_docs/disagg_serving.md

[cr7258]

PD 分离时，用 NCCL 和 NIXL 的 P2P 性能差异如何？

在 PD 分离场景下，NIXL 相比 NCCL 具有更高的性能，因为它采用了 InfiniBand GPU-Async（IBGDA）机制，使控制流和数据流都可直接从 GPU 传输到网卡，无需经过 CPU 中转。

To transfer from prefill nodes to decoding nodes, low latency high bandwidth communication transfer libraries is needed. NIXL will be using InfiniBand GPU-Async Initialized (IBGDA). Currently in NCCL, the control flow goes through CPU proxy threads while the dataflow goes directly to the NIC without needing to go through the CPU buffering. But with IBGDA, both control flow & dataflow doesn’t need to go through the CPU and goes directly from the GPU to the NIC.
https://semianalysis.com/2025/03/19/nvidia-gtc-2025-built-for-reasoning-vera-rubin-kyber-cpo-dynamo-inference-jensen-math-feynman/#nixl-nvidia-inference-transfer-engine

[cr7258]

在 PD 分离场景下使用 Speculative Decoding 出现模型乱说话的现象，为什么？

在 Speculative Decoding 模式 下运行 PD 分离除了要传输的 KV cache，还必须额外传递 Speculative Decoding 所需的 logits 和 hidden states 元数据。

After the transfer of the last chunk is complete, any additional metadata is also copied over: speculative decoding or MTP require logits and hidden states to be moved to the decoder. https://www.perplexity.ai/hub/blog/disaggregated-prefill-and-decode

[cr7258]

vllm 的 PD 分离里，decode 阶段要重新算 prifill 阶段算出来的首字 token 么?

这个不一定，看具体的实现，比如 vllm 官方的 example 里没有传递 prefill 首字 token，decode 阶段就要重新算。https://github.com/vllm-project/vllm/blob/main/examples/offline_inference/disaggregated_prefill.py

而像 dynamo 则是会显式传递首字 token 给 decode worker。代码路径：

• Prefill Worker 生成首字 token（文件：components/src/dynamo/vllm/handlers.py:242-253）

async for res in gen:
    logger.debug(f"kv transfer params: {res.kv_transfer_params}")
    
    # ← res 是 vLLM 的 RequestOutput
    token_ids = res.outputs[0].token_ids if res.outputs else []
    #           ↑ 这就是首字 token！
    
    output: Dict[str, Any] = {
        "token_ids": list(token_ids),  # ← 包含首字 token
        "disaggregated_params": (
            {"kv_transfer_params": res.kv_transfer_params}
            if res.kv_transfer_params
            else None
        ),
    }
    
    yield output  # ← 返回给 PrefillRouter

• PrefillRouter 提取并传递（文件: lib/llm/src/kv_router/prefill_router.rs:153-187）

async fn call_prefill(
    &self,
    request: SingleIn<PreprocessedRequest>,
) -> Result<serde_json::Value> {
    # 调用 Prefill Worker
    let mut prefill_response = match prefill_router {
        InnerPrefillRouter::KvRouter(router) => router.generate(request).await?,
        InnerPrefillRouter::SimpleRouter(router) => router.generate(request).await?,
    };
    
    # 获取第一个输出（包含首字 token）
    let Some(first_output) = prefill_response.next().await else {
        bail!("Prefill router returned no output");
    };
    
    let Some(output) = &first_output.data else {
        bail!("Prefill router output has no data field");
    };
    
    # 提取 disaggregated_params（包含 token_ids 和 kv_transfer_params）
    let Some(disaggregated_params) = output.disaggregated_params.clone() else {
        bail!("Prefill router output missing disaggregated_params");
    };
    
    Ok(disaggregated_params)  // ← 返回，包含首字 token
}

• PrefillRouter 注入到 Decode 请求（文件 ：lib/llm/src/kv_router/prefill_router.rs:226-243）

match self.call_prefill(prefill_request).await {
    Ok(disaggregated_params) => {
        tracing::debug!("Prefill succeeded, using disaggregated params for decode");
        
        # 更新 Decode 请求
        let mut decode_req = req;
        decode_req.disaggregated_params = Some(disaggregated_params);
        #                                     ↑ 包含首字 token！
        
        # 发送给 Decode Worker
        let decode_request = context.map(|_| decode_req);
        next.generate(decode_request).await
    }
    ...
}

• Decode Worker 接收首字 token（文件: components/src/dynamo/vllm/handlers.py:167-183）

# Decode Worker Handler
prompt = TokensPrompt(prompt_token_ids=request["token_ids"])

# 提取 Prefill 传来的参数
disaggregated_params = request.get("disaggregated_params")
#                                  ↑ 从 PrefillRouter 注入的！

if disaggregated_params:
    # Prefill 已完成 - 使用传来的 disaggregated params
    if sampling_params.extra_args is None:
        sampling_params.extra_args = {}
    sampling_params.extra_args["kv_transfer_params"] = \
        disaggregated_params.get("kv_transfer_params")
    
    logger.debug(
        f"Using disaggregated params from prefill for request {request_id}"
    )

[cr7258]

Dynamo kv cache-aware routing 的计算公式

logit = kv_overlap_score_weight × potential_prefill_blocks + potential_active_blocks

• logit 的值越低越好
• kv_overlap_score_weight 默认是 1，> 1 注重 TTFT 优先找 kv cache 匹配长的 wokrer；< 1 注重 TPOP 优先找请求负载低的 worker
• potential_prefill_blocks 是需要计算的 prefill block 数量（kv cache match 越长需要计算的 prefill block 越小）
• potential_active_blocks 是当前正在处理请求的 block 数量，体现了负载情况