runnable

Author: ch-wan

python/sglang/srt/managers/data_parallel_controller.py

@@ -23,6 +23,8 @@
 import setproctitle
 import zmq
 
+from python.sglang.srt.disaggregation.utils import DisaggregationMode
+from python.sglang.srt.managers.schedule_batch import Req
 from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 from sglang.srt.managers.io_struct import (
     TokenizedEmbeddingReqInput,
@@ -220,9 +222,12 @@ def launch_tensor_parallel_group(
         self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
         self.max_req_input_len = scheduler_info[0]["max_req_input_len"]
 
-    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)
+    def round_robin_scheduler(self, req: Req):
+        if self.server_args.disaggregation_mode == DisaggregationMode.NULL:
+            self.workers[self.round_robin_counter].send_pyobj(req)
+            self.round_robin_counter = (self.round_robin_counter + 1) % len(self.workers)
+        else:
+            self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req)
 
     def shortest_queue_scheduler(self, input_requests):
         raise NotImplementedError()

python/sglang/srt/managers/scheduler.py

@@ -484,7 +484,7 @@ def init_memory_pool_and_cache(self):
                 self.tree_cache = HiRadixCache(
                     req_to_token_pool=self.req_to_token_pool,
                     token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
-                    tp_cache_group=self.tp_worker.get_tp_cpu_group(),
+                    tp_cache_group=self.tp_cpu_group,
                     page_size=self.page_size,
                     hicache_ratio=server_args.hicache_ratio,
                 )
@@ -553,7 +553,7 @@ def init_disaggregation(self):
 
             # The decode requests polling kv cache
             self.disagg_decode_transfer_queue = DecodeTransferQueue(
-                gloo_group=self.tp_worker.get_attention_tp_cpu_group(),
+                gloo_group=self.attn_tp_cpu_group,
                 req_to_metadata_buffer_idx_allocator=req_to_metadata_buffer_idx_allocator,
                 metadata_buffers=metadata_buffers,
             )
@@ -568,7 +568,7 @@ def init_disaggregation(self):
                 scheduler=self,
                 transfer_queue=self.disagg_decode_transfer_queue,
                 tree_cache=self.tree_cache,
-                gloo_group=self.tp_worker.get_attention_tp_cpu_group(),
+                gloo_group=self.attn_tp_cpu_group,
                 tp_rank=self.tp_rank,
                 tp_size=self.tp_size,
                 bootstrap_port=self.server_args.disaggregation_bootstrap_port,
@@ -597,7 +597,7 @@ def init_disaggregation(self):
                 tp_rank=self.tp_rank,
                 tp_size=self.tp_size,
                 bootstrap_port=self.server_args.disaggregation_bootstrap_port,
-                gloo_group=self.tp_worker.get_attention_tp_cpu_group(),
+                gloo_group=self.attn_tp_cpu_group,
                 transfer_backend=self.transfer_backend,
                 scheduler=self,
             )
@@ -676,6 +676,11 @@ def event_loop_normal_disagg_prefill(self):
             )
             self.process_prefill_chunk()
             batch = self.get_new_batch_prefill()
+            
+            # Handle DP attention
+            if self.server_args.enable_dp_attention or self.server_args.enable_sp_layernorm:
+                batch, _ = self.prepare_dp_attn_batch(batch)
+
             self.cur_batch = batch
 
             if batch:
@@ -704,11 +709,18 @@ def event_loop_normal_disagg_decode(self):
             # polling and allocating kv cache
             self.process_decode_queue()
             batch = self.get_next_disagg_decode_batch_to_run()
+            
+            # Handle DP attention
+            if self.server_args.enable_dp_attention or self.server_args.enable_sp_layernorm:
+                batch, do_extend = self.prepare_dp_attn_batch(batch)
+            elif batch:
+                do_extend = batch.forward_mode.is_extend()
+                
             self.cur_batch = batch
 
             if batch:
                 # Generate fake extend output.
-                if batch.forward_mode.is_extend():
+                if do_extend:
                     # Note: Logprobs should be handled on the prefill engine.
                     self.stream_output(
                         batch.reqs, [False for _ in range(len(batch.reqs))]