Support multi-node DP attention (#2925)

Co-authored-by: dhou-xai <dhou@x.ai>
2025-01-16 11:15:00 -08:00
parent 58f3f2b840
commit 8b6ce52e92
16 changed files with 287 additions and 137 deletions
--- a/python/sglang/srt/managers/tp_worker_overlap_thread.py
+++ b/python/sglang/srt/managers/tp_worker_overlap_thread.py
@@ -92,6 +92,9 @@ class TpModelWorkerClient:
    def get_tp_cpu_group(self):
        return self.worker.get_tp_cpu_group()

+    def get_attention_tp_cpu_group(self):
+        return self.worker.get_attention_tp_cpu_group()
+
    def get_memory_pool(self):
        return (
            self.worker.model_runner.req_to_token_pool,