[P/D] layerwise connector support recompute scheduler

[liziyu179]

What this PR does / why we need it?

layerwise connector support recompute scheduler.

NOTE：
Triggering recompute will invoke the tokenizer again, which may lead to precision fluctuations.

[RFC]: CDCP Scheduling for Disaggregated Prefilling with KV Cache Layerwise Push Support #4842

Does this PR introduce any user-facing change?

How was this patch tested?

• vLLM version: v0.13.0
• vLLM main: vllm-project/vllm@bde38c1

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Code Review

This pull request adds support for a recompute scheduler in the layerwise connector example. The implementation introduces complex logic to handle retries when a 'recomputed' signal is received from the backend.

My review has identified several issues:

• There are two critical bugs that could lead to unhandled exceptions: one is an IndexError if the messages list in a chat request is empty, and the other is a TypeError in the completion_tokens calculation for non-streaming responses.
• I also found a couple of high-severity issues related to code quality and robustness. The released_kv flag seems to be part of an incomplete copy-paste and might hide a resource leak. Additionally, the logic for recalculating max_tokens during a retry is brittle and should be refactored for better clarity and reliability.

I have provided specific comments with suggestions for fixing these issues.

[gemini-code-assist]

The code at line 439 (messages[0].get("content", "")) assumes that the messages list is not empty. If an API request is sent with an empty messages list (e.g., "messages": []), this will raise an IndexError, causing an unhandled exception. While the OpenAI API spec requires at least one message, it's best to code defensively.

A similar issue exists on lines 507-508. You should add checks to ensure messages is not empty before accessing its elements.

Here's a suggested way to fix this:

# At lines 438-439
            messages = req_data["messages"]
            origin_prompt = messages[0].get("content", "") if messages else ""

# And at lines 506-508
                            if chat_flag and messages:
                                messages[0][
                                    "content"] = origin_prompt + generated_token

Suggested change

	messages = req_data["messages"]
	origin_prompt = messages[0].get("content", "")
	messages = req_data["messages"]
	origin_prompt = messages[0].get("content", "") if messages else ""

[gemini-code-assist]

In the non-streaming case, usage.get("completion_tokens") can return None if the key is not present in the usage dictionary. Adding None to completion_tokens will raise a TypeError. You should provide a default value of 0 to prevent this.

Suggested change

	completion_tokens = (completion_tokens + 1) if stream_flag else \
	(completion_tokens + usage.get("completion_tokens"))
	completion_tokens = (completion_tokens + 1) if stream_flag else \
	(completion_tokens + usage.get("completion_tokens", 0))

[gemini-code-assist]

The released_kv variable is declared nonlocal and then reassigned, but it's never read within the generate_stream function. This indicates dead code. It seems that logic for releasing the KV cache, which is present in load_balance_proxy_server_example.py, is missing here. This could lead to a resource leak if the KV cache is not freed. Please either add the KV cache release logic or remove the unused released_kv variable and its related declarations.

[gemini-code-assist]

The calculation for max_tokens on retry, origin_max_tokens - completion_tokens + retry_count, is brittle. It seems to compensate for completion_tokens being incremented for the "recomputed" signal chunk, which is not a real token. This makes the logic hard to follow and prone to errors if the backend behavior changes.

A more robust approach would be to reorder the logic to check for stop_reason == "recomputed" before processing the chunk as a regular token and incrementing completion_tokens. This way, completion_tokens would accurately reflect the number of actual generated tokens, and the max_tokens calculation would simplify to origin_max_tokens - completion_tokens.

[LCAIZJ]

Could you please provide a more detailed explanation of the background and design of this PR?

[winson-00178005]

如果抛出异常的时候，是不是也要调用这个来释放？

[liziyu179]

Running to this line of the function signifies the end of stream return, releasing the key-value cache records of node D.

[winson-00178005]

可以把request_id/retry_count加入日志中，方便问题定位

[liziyu179]

Thanks for reviewing the code. The number of retries will be printed in stream_service_response_with_retry, and the request_id will be printed on the next line.