Revert "[PD Disaggregation] replace transfer with batch transfer for better performance (#7236)" #7968
Conversation
…better performance (sgl-project#7236) " This reverts commit 2ed68d7.
There was a problem hiding this comment.
Summary of Changes
Hello @fzyzcjy, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request reverts a previous change (#7236) that aimed to improve performance by replacing individual data transfers with a batch transfer mechanism for KV cache synchronization in the disaggregation module. The revert suggests that the intended performance benefits were not realized or that the batch transfer approach introduced unforeseen issues, restoring the previous, more granular transfer method.
Highlights
- Revert of Performance Optimization: This pull request undoes the changes introduced in PR #7236, which had attempted to optimize KV cache synchronization by switching from individual
transfer_synccalls to a singlebatch_transfer_syncoperation. - Restoration of Individual Transfers: Specifically, within the
process_layerfunction ofsend_kvcacheinpython/sglang/srt/disaggregation/mooncake/conn.py, the code has been reverted to perform individualtransfer_synccalls in a loop, rather than collecting all transfer details for a singlebatch_transfer_syncoperation. - Reintroduction of Per-Transfer Error Checking: The reverted code also reintroduces immediate error checking (
if status != 0: return status) after each individualtransfer_synccall, which was implicitly handled differently by the batch transfer.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request reverts the use of batch_transfer_sync back to transfer_sync in the send_kvcache function, likely to fix a correctness or stability issue. The review identifies that batch_transfer_sync is still used in two other functions within the same file (send_kvcache_slice and send_aux). It is recommended to revert these usages as well to maintain consistency and prevent potential issues, assuming the batched transfer method is the root cause of the problem.
| def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int: | ||
| src_addr_list = [] | ||
| dst_addr_list = [] | ||
| length_list = [] | ||
| for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks): | ||
| src_addr = src_ptr + int(prefill_index[0]) * item_len | ||
| dst_addr = dst_ptr + int(decode_index[0]) * item_len | ||
| length = item_len * len(prefill_index) | ||
| src_addr_list.append(src_addr) | ||
| dst_addr_list.append(dst_addr) | ||
| length_list.append(length) | ||
| return self.engine.batch_transfer_sync( | ||
| mooncake_session_id, src_addr_list, dst_addr_list, length_list | ||
| ) | ||
|
|
||
| status = self.engine.transfer_sync( | ||
| mooncake_session_id, src_addr, dst_addr, length | ||
| ) | ||
| if status != 0: | ||
| return status | ||
| return 0 |
There was a problem hiding this comment.
This revert from batch_transfer_sync to transfer_sync seems to address a potential issue with the batched transfer. However, batch_transfer_sync is still being used in other parts of this file, specifically in send_kvcache_slice (line 436) and send_aux (line 476).
If batch_transfer_sync is indeed problematic, its continued use in other functions poses a risk to correctness and stability. To ensure consistency, these other usages should also be reverted to their non-batched counterparts.
For example, send_aux could be changed to:
def send_aux(
self,
mooncake_session_id: str,
prefill_aux_index: int,
dst_aux_ptrs: list[int],
dst_aux_index: int,
):
prefill_aux_ptrs = self.kv_args.aux_data_ptrs
prefill_aux_item_lens = self.kv_args.aux_item_lens
for i, dst_aux_ptr in enumerate(dst_aux_ptrs):
length = prefill_aux_item_lens[i]
src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
dst_addr = dst_aux_ptrs[i] + length * dst_aux_index
status = self.engine.transfer_sync(
mooncake_session_id, src_addr, dst_addr, length
)
if status != 0:
return status
return 0Please consider reverting all usages of batch_transfer_sync in this file.
|
@fzyzcjy @ShangmingCai do we need to use the older version of transfer engine |
|
I test using latest mooncake and it works |
|
@zhyncs No need, we only need to fix batch_transfer_sync with nvlink transport, then I will release a new package and bring this change back. The older transfer_sync is functioning well. |
…better performance (sgl-project#7236)" (sgl-project#7968)
…better performance (sgl-project#7236)" (sgl-project#7968)
…better performance (sgl-project#7236)" (sgl-project#7968)
…better performance (sgl-project#7236)" (sgl-project#7968)
…fer for better performance (sgl-project#7236)" (sgl-project#7968)" This reverts commit 880221b.
cc @ShangmingCai