[P/D] Performance enhancement of Layerwise connector in TP asymmetric scenarios

[liziyu179]

What this PR does / why we need it?

[P/D] Performance enhancement of Layerwise connector in TP asymmetric scenarios

1. Session fusion: For transmission tasks at each layer, aggregate transmission tasks with the same destination and merge them into a single task for assignment.
2. Alltoall aggregation: For TP asymmetric scenarios, perform all alltoall operations at once according to the block granularity for all requests.

[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@45c1ca1

[gemini-code-assist]

Code Review

This pull request introduces performance enhancements for the Layerwise connector in asymmetric Tensor Parallelism scenarios by refactoring the KV cache transfer logic to batch requests. While the overall direction of refactoring is positive, there are several critical issues in the implementation that need to be addressed. Specifically, the error handling for transfer failures is incorrect and could lead to data corruption. Additionally, there are concerns about resource management and swallowed exceptions related to the new metaserver communication, which could cause the system to hang or fail silently. I have provided detailed comments on these points.

[gemini-code-assist]

The error handling for batch_transfer_sync_write has been changed to no longer raise an exception on failure. Instead, it logs an error and, for the final layer, proceeds to call self.callback_func. This callback signals to the receiving end that the transfer is complete. This is incorrect behavior in case of a failure, as the receiver will assume a successful transfer and may proceed with incomplete or corrupted KV cache data, leading to incorrect model outputs or crashes.

The previous implementation correctly raised a RuntimeError, which would stop the sending thread and prevent signaling a false success. The TODO comment indicates awareness of this issue, but the current implementation is dangerous. The callback should not be called on failure. The exception should be re-raised to ensure the failure is propagated.

                if ret < 0:
                    logger.error(
                        f"Mooncake transfer failed for send requests {transfer_meta.req_ids} kv cache to {session_id}"
                    )
                    raise RuntimeError(f"Mooncake transfer failed for requests {transfer_meta.req_ids}")
                else:
                    if send_task.layer_idx == (self.total_layers - 1):
                        for req_id in transfer_meta.req_ids:
                            req_meta = send_task.send_request[req_id]
                            if req_meta.chunk_finish:
                                self.callback_func(req_id, req_meta)

[gemini-code-assist]

The httpx.Client is initialized with timeout=None. This means that requests to the metaserver can hang indefinitely if the server is unresponsive, which can cause the scheduler to hang. It is recommended to set a reasonable timeout to prevent this. Additionally, max_connections=100000 is a very high limit and could lead to resource exhaustion on the client or server side. Please consider if such a high limit is necessary and if it can be reduced.

        self.metaserver_client = httpx.Client(
            limits=httpx.Limits(max_connections=1000), timeout=30.0)

[gemini-code-assist]

The future.add_done_callback for the _access_metaserver call only logs the exception if one occurs. The main thread continues execution without being aware of the failure. If accessing the metaserver is critical for the request to be processed correctly (which seems to be the case, as it sets up remote prefill), then swallowing the exception can lead to hard-to-debug failures later on. The failure should be propagated or handled in a way that the request is marked as failed. For example, the callback could add the request to a failure queue to be handled by the scheduler, or the main thread could block on future.result() to handle the exception immediately.

[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.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.