[BugFix] Report prefiller cached tokens in PD proxy (vllm-project#11630)

在这个PR的基础上做了一些优化，vllm-project#10598
优化点：
SSE模式下，只有最后一次的chunk的cached token数是正确的，数值=prefill节点的cached
token数，而无需每轮chunk都修改。

测试流程：
单机分别部署PD节点：
P节点：
vllm serve /xxx/Qwen3-8B \
--host 0.0.0.0 \
--port 1000 \
--enable-prefix-caching \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name Qwen3-8B \
--max-model-len 40000 \
--max-num-batched-tokens 40000 \
--enable-prompt-tokens-details \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--kv-transfer-config \
'{"kv_connector": "MooncakeConnectorV1",
"kv_role": "kv_producer",
"kv_port": "30000",
"kv_connector_extra_config": {
"prefill": {
"dp_size": 1,
"tp_size": 1
},
"decode": {
"dp_size": 1,
"tp_size": 1
}
}'

D节点：
vllm serve /xxx/Qwen3-8B \
--host 0.0.0.0 \
--port 1001 \
--no-enable-prefix-caching \
--tensor-parallel-size 1 \
--seed 1024 \
--served-model-name Qwen3-8B \
--max-model-len 40000 \
--enable-chunked-prefill \
--max-num-batched-tokens 40000 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--enable-force-include-usage \
--kv-transfer-config \
'{"kv_connector": "MooncakeConnectorV1",
"kv_role": "kv_consumer",
"kv_port": "30100",
"kv_connector_extra_config": {
"prefill": {
"dp_size": 1,
"tp_size": 1
},
"decode": {
"dp_size": 1,
"tp_size": 1
}
}'
启动：
python3
/xxx/vllm-ascend/examples/disaggregated_prefill_v1/load_balance_proxy_server_example.py
\
--host 0.0.0.0 \
--port 8898 \
--prefiller-hosts 0.0.0.0 \
--prefiller-ports 1000 \
--decoder-hosts 0.0.0.0 \
--decoder-ports 1001

结果如下：
root@test-xxx# curl -X POST http://localhost:8898/v1/chat/completions \
>     -H "Content-Type: application/json" \
>     -d '{
>           "model": "Qwen3-8B",
> "messages": [{"role": "user", "content": "The Aegean Sea is an
elongated embayment of the Mediterranean Sea, lying between Europes
Balkan peninsula and Asia Minor. It is home to thousands of islands, the
most famous being Crete, Rhodes, Mykonos, and Santorini. The Aegean is
widely regarded as the cradle of ancient Greek civilization. It was here
that the Minoans and Mycenaeans built some of Europes earliest advanced
societies, and where legends of Troy and Atlantis were born.Today, the
sea is known for its crystal-clear waters, white-washed villages perched
on cliffs, and a rich maritime heritage that continues to shape the
culture of Greece and Turkey."}],
>           "max_tokens": 10,
>            "stream": true,
>            "stream_options": {
>               "include_usage": true
>            }
>        }'
data: {"id": "chatcmpl-xxxx", "object": "chat.completion.chunk",
"created": 1783499708, "model": "Qwen3-8B", "choi
ces": [{"index": 0, "delta": {"role": "assistant", "content": ""},
"logprobs": null, "finish_reason": null}], "usage": {"prompt_tokens":
145, "total_tokens": 145, "completion_tokens": 0,
"prompt_tokens_details": {"cached_tokens": 145}}, "prompt_token_ids":
null, "prompt_text": null}

data: {"id": "chatcmpl-xxxx", "object": "chat.completion.chunk",
"created": 1783499708, "model": "Qwen3-8B", "choi
ces": [{"index": 0, "delta": {"content": "<think>"}, "logprobs": null,
"finish_reason": null, "token_ids": null}], "usage": {"prompt_tokens":
145, "total_tokens": 146, "completion_tokens": 1,
"prompt_tokens_details": {"cached_tokens": 145}}}

data: {"id": "chatcmpl-xxxx", "object": "chat.completion.chunk",
"created": 1783499708, "model": "Qwen3-8B", "choi
ces": [{"index": 0, "delta": {"content": "\n"}, "logprobs": null,
"finish_reason": null, "token_ids": null}], "usage": {"prompt_tokens":
145, "total_tokens": 147, "completion_tokens": 2,
"prompt_tokens_details": {"cached_tokens": 145}}}

......

data: {"id": "chatcmpl-xxxx", "object": "chat.completion.chunk",
"created": 1783499708, "model": "Qwen3-8B", "choi
ces": [{"index": 0, "delta": {"content": " about"}, "logprobs": null,
"finish_reason": "length", "stop_reason": null, "token_ids": null}],
"usage": {"prompt_tokens": 145, "total_tokens": 155,
"completion_tokens": 10, "prompt_tokens_details": {"cached_tokens":
145}}}

data: {"id": "chatcmpl-xxxx", "object": "chat.completion.chunk",
"created": 1783499708, "model": "Qwen3-8B", "choi
ces": [], "usage": {"prompt_tokens": 145, "total_tokens": 155,
"completion_tokens": 10, "prompt_tokens_details": {"cached_tokens":
128}}, "system_fingerprint": "vllm-0.23.0-70fa37d5"}

data: [DONE]

仅最后一次总的token统计数{"cached_tokens": 128}，（本实验中block size设置的128token）符合预期。
- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@1f486d9

Signed-off-by: LiuY <LiuYuan_jjj@163.com>

Author: LYMAXPUP

examples/disaggregated_prefill_v1/load_balance_proxy_server_example.py

@@ -161,10 +161,38 @@ class InstanceInfo:
     decoder_score: float
     decoder_host: str
     decoder_port: int
+    prefiller_cached_tokens: int | None = None
 
 
 TAINT_PRIORITY = 1e15
 
+
+def extract_cached_tokens(response_json: dict) -> int | None:
+    usage = response_json.get("usage") or {}
+    prompt_tokens_details = usage.get("prompt_tokens_details") or {}
+    cached_tokens = prompt_tokens_details.get("cached_tokens")
+    return cached_tokens if isinstance(cached_tokens, int) else None
+
+
+def update_cached_tokens_in_chunk(chunk_json: dict, cached_tokens: int | None) -> bool:
+    if cached_tokens is None:
+        return False
+    usage = chunk_json.get("usage")
+    if not isinstance(usage, dict):
+        return False
+    prompt_tokens_details = usage.get("prompt_tokens_details")
+    if not isinstance(prompt_tokens_details, dict):
+        prompt_tokens_details = {}
+    usage["prompt_tokens_details"] = prompt_tokens_details
+    prompt_tokens_details["cached_tokens"] = cached_tokens
+    return True
+
+
+def encode_response_chunk(chunk_json: dict, is_sse: bool) -> bytes:
+    chunk = json.dumps(chunk_json, ensure_ascii=False).encode("utf-8")
+    return b"data: " + chunk + b"\n\n" if is_sse else chunk
+
+
 global_args: argparse.Namespace | None = None
 shared_scheduler: "SharedProxyScheduler | None" = None
 runtime: "WorkerRuntime | None" = None
@@ -922,9 +950,11 @@ async def assign_instances(
         await _abort_prefill_selection(runtime, prefiller_key, prefiller_score, is_initial_request=is_initial_request)
         raise
 
-    kv_transfer_params = response.json().get("kv_transfer_params", {})
+    response_json = response.json()
+    kv_transfer_params = response_json.get("kv_transfer_params", {})
     if kv_transfer_params:
         req_data["kv_transfer_params"] = kv_transfer_params
+    prefiller_cached_tokens = extract_cached_tokens(response_json)
 
     try:
         decoder = await runtime.schedule("pick_decoder", decoder_score)
@@ -943,6 +973,7 @@ async def assign_instances(
         decoder_score=decoder_score,
         decoder_host=decoder["host"],
         decoder_port=decoder["port"],
+        prefiller_cached_tokens=prefiller_cached_tokens,
     )
 
 
@@ -987,6 +1018,7 @@ async def generate_stream():
             retry_count = 0
             retry = True
             completion_tokens = 0
+            reported_prefiller_cached_tokens = instance_info.prefiller_cached_tokens
 
             async def release_prefill_kv_once() -> None:
                 nonlocal released_kv
@@ -1018,7 +1050,8 @@ async def release_prefill_kv_once() -> None:
                             continue
                         if not chunk_str:
                             continue
-                        if chunk_str.startswith("data: "):
+                        is_sse = chunk_str.startswith("data: ")
+                        if is_sse:
                             chunk_str = chunk_str[len("data: ") :]
                         try:
                             chunk_json = json.loads(chunk_str)
@@ -1028,6 +1061,8 @@ async def release_prefill_kv_once() -> None:
                             continue
                         choices = chunk_json.get("choices", [])
                         if not choices:
+                            if update_cached_tokens_in_chunk(chunk_json, reported_prefiller_cached_tokens):
+                                chunk = encode_response_chunk(chunk_json, is_sse)
                             yield chunk
                             continue
 
@@ -1061,7 +1096,7 @@ async def release_prefill_kv_once() -> None:
                                 choice["message"]["content"] = generated_token
                             else:
                                 choice["text"] = generated_token
-                            chunk = json.dumps(chunk_json).encode("utf-8")
+                            chunk = encode_response_chunk(chunk_json, is_sse)
                         yield chunk
             except asyncio.CancelledError:
                 logger.warning(