feat: add direct routing strategy to DP worker (#6884)

This commit is contained in:
ishandhanani
2025-06-09 11:44:05 -07:00
committed by GitHub
parent 3465d7ae78
commit f1569876d5
12 changed files with 78 additions and 8 deletions

View File

@@ -407,9 +407,10 @@ class NixlKVReceiver(CommonKVReceiver):
mgr: NixlKVManager,
bootstrap_addr: str,
bootstrap_room: Optional[int] = None,
data_parallel_rank: Optional[int] = None,
):
self.started_transfer = False
super().__init__(mgr, bootstrap_addr, bootstrap_room)
super().__init__(mgr, bootstrap_addr, bootstrap_room, data_parallel_rank)
def init(self, kv_indices: npt.NDArray[np.int64], aux_index: Optional[int] = None):
for bootstrap_info in self.bootstrap_infos: