Support DP MLA (#1970)

python/sglang/srt/managers/data_parallel_controller.py

@@ -81,20 +81,34 @@ class DataParallelController:
         # Start data parallel workers
         base_gpu_id = 0
         self.workers = []
+        scheduler_pipe_readers = []
         for dp_rank in range(server_args.dp_size):
             tmp_port_args = PortArgs.init_new(server_args)
             tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
             tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
 
-            send_to = self.launch_tensor_parallel_group(
-                server_args,
-                tmp_port_args,
-                base_gpu_id,
-                dp_rank,
-            )
-
+            if server_args.enable_dp_attention:
+                # Share workers for DP and TP
+                send_to, reader = self.launch_tensor_parallel_process(
+                    server_args,
+                    tmp_port_args,
+                    base_gpu_id,
+                    dp_rank,
+                )
+                base_gpu_id += 1
+                scheduler_pipe_readers.append(reader)
+            else:
+                send_to = self.launch_tensor_parallel_group(
+                    server_args,
+                    tmp_port_args,
+                    base_gpu_id,
+                    dp_rank,
+                )
+                base_gpu_id += server_args.tp_size
             self.workers.append(send_to)
-            base_gpu_id += server_args.tp_size
+
+        for reader in scheduler_pipe_readers:
+            reader.recv()
 
     def launch_tensor_parallel_group(
         self,
@@ -132,6 +146,27 @@ class DataParallelController:
 
         return send_to
 
+    def launch_tensor_parallel_process(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        base_gpu_id: int,
+        dp_rank: int,
+    ):
+        reader, writer = mp.Pipe(duplex=False)
+        gpu_id = base_gpu_id
+        tp_rank = dp_rank
+        proc = mp.Process(
+            target=run_scheduler_process,
+            args=(server_args, port_args, gpu_id, tp_rank, dp_rank, writer),
+        )
+        proc.start()
+        send_to = get_zmq_socket(
+            self.context, zmq.PUSH, port_args.scheduler_input_ipc_name
+        )
+
+        return send_to, reader
+
     def round_robin_scheduler(self, req):
         self.workers[self.round_robin_counter].send_pyobj(req)
         self.round_robin_counter = (self.round_robin_counter + 1) % len(self.workers)