From 977f04aafecac1e474751d009ceed61826ef2266 Mon Sep 17 00:00:00 2001 From: Ilia Yastrebov Date: Wed, 29 Jul 2026 15:02:13 +0200 Subject: [PATCH] [PD] NIXL connector: shard by destination (#32025) --- python/sglang/srt/disaggregation/nixl/conn.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 1ab1562fb..831d2f267 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -2254,13 +2254,16 @@ class NixlKVManager(CommonKVManager): if self.enable_staging: self._prefetch_staging_reqs(bootstrap_room) - # Transfer is async: just enqueue the chunk; the per-queue worker - # (transfer_worker) does the actual gather + RDMA. Routing by - # ``room % N`` keeps every chunk of a given room on the same - # worker -- and therefore on the same private staging buffer -- - # which is required for the staging ring's offset/watermark - # state machine to advance correctly. - shard_idx = bootstrap_room % len(self.transfer_queues) + if bootstrap_room not in self.transfer_infos: + # Dummy rank or already cleared; nothing to enqueue. + return None + + # Shard by destination (mirror mooncake): same dst endpoint(s) -> same + # worker, keeping a room's chunks on one private staging buffer. + session_port_sum = sum( + info.dst_port for info in self.transfer_infos[bootstrap_room].values() + ) + shard_idx = session_port_sum % len(self.transfer_queues) self.transfer_queues[shard_idx].put( TransferKVChunk( room=bootstrap_room,