[Feature] PD-Multiplexing Context and Scheduler, lazy import spatial. (#12275)

Author: ykcombat

python/sglang/srt/layers/logits_processor.py

@@ -134,10 +134,7 @@ class LogitsMetadata:
     @classmethod
     def from_forward_batch(cls, forward_batch: ForwardBatch):
         if (
-            (
-                forward_batch.forward_mode.is_extend()
-                or forward_batch.forward_mode.is_split_prefill()
-            )
+            forward_batch.forward_mode.is_extend()
             and forward_batch.return_logprob
             and not forward_batch.forward_mode.is_target_verify()
         ):
@@ -384,8 +381,8 @@ def forward(
             input_logprob_indices = None
         elif (
             logits_metadata.forward_mode.is_extend()
-            or logits_metadata.forward_mode.is_split_prefill()
-        ) and not logits_metadata.extend_return_logprob:
+            and not logits_metadata.extend_return_logprob
+        ):
             # Prefill without input logprobs.
             if logits_metadata.padded_static_len < 0:
                 last_index = torch.cumsum(logits_metadata.extend_seq_lens, dim=0) - 1

python/sglang/srt/managers/schedule_batch.py

@@ -72,7 +72,11 @@
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
 from sglang.srt.metrics.collector import SchedulerMetricsCollector, TimeStats
-from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import ServerArgs, get_global_server_args

python/sglang/srt/managers/scheduler.py

@@ -152,6 +152,7 @@
 from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
+from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin
 from sglang.srt.parser.reasoning_parser import ReasoningParser
 from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -213,6 +214,7 @@ class Scheduler(
     SchedulerMetricsMixin,
     SchedulerDisaggregationDecodeMixin,
     SchedulerDisaggregationPrefillMixin,
+    SchedulerMultiplexMixin,
     SchedulerRuntimeCheckerMixin,
     SchedulerPPMixin,
 ):
@@ -252,6 +254,7 @@ def __init__(
         self.enable_lora = server_args.enable_lora
         self.max_loras_per_batch = server_args.max_loras_per_batch
         self.enable_overlap = not server_args.disable_overlap_schedule
+        self.enable_pdmux = server_args.enable_pdmux
         self.skip_tokenizer_init = server_args.skip_tokenizer_init
         self.enable_metrics = server_args.enable_metrics
         self.enable_metrics_for_all_schedulers = (
@@ -285,6 +288,10 @@ def __init__(
         # Init inter-process communication
         self.init_sockets(server_args, port_args)
 
+        # Init pdmux context
+        if self.enable_pdmux:
+            self.init_pdmux()
+
         # Init tokenizer
         self.init_tokenizer()
 
@@ -424,6 +431,8 @@ def __init__(
         self.running_batch: ScheduleBatch = ScheduleBatch(reqs=[], batch_is_full=False)
         # The current forward batch
         self.cur_batch: Optional[ScheduleBatch] = None
+        # The current split prefill batch
+        self.split_prefill_batch: Optional[ScheduleBatch] = None
         # The last forward batch
         self.last_batch: Optional[ScheduleBatch] = None
         self.forward_ct = 0
@@ -1952,7 +1961,6 @@ def run_batch(
 
         # Run forward
         if self.is_generation:
-
             batch_or_worker_batch = batch
 
             if self.enable_overlap or self.spec_algorithm.is_none():
@@ -2009,6 +2017,9 @@ def run_batch(
                     # The future value, usually for next batch preparation
                     # Current implementation strictly synchronizes the seq_lens
                     batch.seq_lens = batch_result.next_draft_input.new_seq_lens
+            elif self.enable_pdmux and batch.forward_mode.is_split_prefill():
+                batch_result = self.tp_worker.forward_batch_split_prefill(batch)
+                future_indices_or_next_token_ids = batch_result.next_token_ids
             else:
                 batch_result = self.model_worker.forward_batch_generation(
                     batch_or_worker_batch
@@ -2791,7 +2802,9 @@ def run_scheduler_process(
 
         disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode
         if disaggregation_mode == DisaggregationMode.NULL:
-            if server_args.pp_size > 1:
+            if scheduler.enable_pdmux:
+                scheduler.event_loop_pdmux()
+            elif server_args.pp_size > 1:
                 scheduler.event_loop_pp()
             elif scheduler.enable_overlap:
                 scheduler.event_loop_overlap()

python/sglang/srt/managers/tp_worker.py

@@ -35,7 +35,7 @@
     UpdateWeightsFromIPCReqInput,
     UpdateWeightsFromTensorReqInput,
 )
-from sglang.srt.managers.schedule_batch import ModelWorkerBatch
+from sglang.srt.managers.schedule_batch import ModelWorkerBatch, ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@@ -425,3 +425,26 @@ def sample_batch_func():
                 pp_hidden_states_proxy_tensors=pp_proxy_tensors,
                 can_run_cuda_graph=can_run_cuda_graph,
             )
+
+    def forward_batch_split_prefill(self, batch: ScheduleBatch):
+        if batch.split_index == 0:
+            model_worker_batch = batch.get_model_worker_batch()
+            forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
+            batch.split_forward_batch = forward_batch
+            batch.seq_lens_cpu_cache = model_worker_batch.seq_lens_cpu
+        else:
+            model_worker_batch = batch.get_model_worker_batch(batch.seq_lens_cpu_cache)
+
+        logits_output, can_run_cuda_graph = self.model_runner.forward(
+            batch.split_forward_batch, split_forward_count=batch.split_forward_count
+        )
+        if logits_output:
+            next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
+        else:
+            next_token_ids = None
+        batch_result = GenerationBatchResult(
+            logits_output=logits_output,
+            can_run_cuda_graph=can_run_cuda_graph,
+        )
+        batch_result.next_token_ids = next_token_ids
+        return batch_result

python/sglang/srt/mem_cache/memory_pool.py

@@ -509,6 +509,7 @@ def __init__(
         enable_memory_saver: bool,
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
+        enable_alt_stream: bool = True,
         enable_kv_cache_copy: bool = False,
     ):
         super().__init__(
@@ -527,7 +528,9 @@ def __init__(
         self._create_buffers()
 
         self.device_module = torch.get_device_module(self.device)
-        self.alt_stream = self.device_module.Stream() if _is_cuda else None
+        self.alt_stream = (
+            self.device_module.Stream() if _is_cuda and enable_alt_stream else None
+        )
 
         if enable_kv_cache_copy:
             self._init_kv_copy_and_warmup()

python/sglang/srt/model_executor/forward_batch_info.py

@@ -96,6 +96,7 @@ def is_extend(self, include_draft_extend_v2: bool = False):
                 else False
             )
             or self == ForwardMode.TARGET_VERIFY
+            or self == ForwardMode.SPLIT_PREFILL
         )
 
     def is_decode(self):

python/sglang/srt/model_executor/model_runner.py

@@ -1765,6 +1765,7 @@ def init_memory_pool(
                     enable_memory_saver=self.server_args.enable_memory_saver,
                     start_layer=self.start_layer,
                     end_layer=self.end_layer,
+                    enable_alt_stream=not self.server_args.enable_pdmux,
                     enable_kv_cache_copy=(
                         self.server_args.speculative_algorithm is not None
                     ),
@@ -1833,12 +1834,18 @@ def init_cublas(self):
 
     def init_attention_backend(self):
         """Init attention kernel backend."""
-        if self.server_args.enable_two_batch_overlap and not self.is_draft_worker:
+        if self.server_args.enable_pdmux:
+            self.attn_backend = self._get_attention_backend(init_new_workspace=True)
+            self.decode_attn_backend_group = []
+            for _ in range(self.server_args.sm_group_num):
+                self.decode_attn_backend_group.append(self._get_attention_backend())
+            self.decode_attn_backend = self.decode_attn_backend_group[0]
+        elif self.server_args.enable_two_batch_overlap and not self.is_draft_worker:
             self.attn_backend = TboAttnBackend.init_new(self._get_attention_backend)
         else:
             self.attn_backend = self._get_attention_backend()
 
-    def _get_attention_backend(self):
+    def _get_attention_backend(self, init_new_workspace: bool = False):
         """Init attention kernel backend."""
         self.prefill_attention_backend_str, self.decode_attention_backend_str = (
             self.server_args.get_attention_backends()
@@ -1852,10 +1859,12 @@ def _get_attention_backend(self):
             attn_backend = HybridAttnBackend(
                 self,
                 decode_backend=self._get_attention_backend_from_str(
-                    self.decode_attention_backend_str
+                    self.decode_attention_backend_str,
+                    init_new_workspace=init_new_workspace,
                 ),
                 prefill_backend=self._get_attention_backend_from_str(
-                    self.prefill_attention_backend_str
+                    self.prefill_attention_backend_str,
+                    init_new_workspace=init_new_workspace,
                 ),
             )
             logger.info(
@@ -1869,7 +1878,8 @@ def _get_attention_backend(self):
             )
         else:
             attn_backend = self._get_attention_backend_from_str(
-                self.server_args.attention_backend
+                self.server_args.attention_backend,
+                init_new_workspace=init_new_workspace,
             )
 
         (
@@ -1878,9 +1888,12 @@ def _get_attention_backend(self):
         ) = (self.prefill_attention_backend_str, self.decode_attention_backend_str)
         return attn_backend
 
-    def _get_attention_backend_from_str(self, backend_str: str):
+    def _get_attention_backend_from_str(
+        self, backend_str: str, init_new_workspace: bool = False
+    ):
         if backend_str not in ATTENTION_BACKENDS:
             raise ValueError(f"Invalid attention backend: {backend_str}")
+        self.init_new_workspace = init_new_workspace
         full_attention_backend = ATTENTION_BACKENDS[backend_str](self)
         return attn_backend_wrapper(self, full_attention_backend)
 
@@ -1978,14 +1991,21 @@ def apply_torch_tp(self):
         device_mesh = torch.distributed.init_device_mesh(self.device, (self.tp_size,))
         tensor_parallel(self.model, device_mesh)
 
+    def update_decode_attn_backend(self, stream_idx: int):
+        self.decode_attn_backend = self.decode_attn_backend_group[stream_idx]
+
     def forward_decode(
         self,
         forward_batch: ForwardBatch,
         skip_attn_backend_init: bool = False,
         pp_proxy_tensors=None,
     ) -> LogitsProcessorOutput:
         if not skip_attn_backend_init:
-            self.attn_backend.init_forward_metadata(forward_batch)
+            if self.server_args.enable_pdmux:
+                self.decode_attn_backend.init_forward_metadata(forward_batch)
+                forward_batch.attn_backend = self.decode_attn_backend
+            else:
+                self.attn_backend.init_forward_metadata(forward_batch)
         # FIXME: add pp_proxy_tensors arg to all models
         kwargs = {}
         if self.support_pp:
@@ -2123,18 +2143,18 @@ def _forward_raw(
                 skip_attn_backend_init=skip_attn_backend_init,
                 pp_proxy_tensors=pp_proxy_tensors,
             )
-        elif forward_batch.forward_mode.is_extend():
-            ret = self.forward_extend(
-                forward_batch,
-                skip_attn_backend_init=skip_attn_backend_init,
-                pp_proxy_tensors=pp_proxy_tensors,
-            )
         elif forward_batch.forward_mode.is_split_prefill():
             ret = self.forward_split_prefill(
                 forward_batch,
                 reinit_attn_backend=reinit_attn_backend,
                 forward_count=split_forward_count,
             )
+        elif forward_batch.forward_mode.is_extend():
+            ret = self.forward_extend(
+                forward_batch,
+                skip_attn_backend_init=skip_attn_backend_init,
+                pp_proxy_tensors=pp_proxy_tensors,
+            )
         elif forward_batch.forward_mode.is_idle():
             ret = self.forward_idle(forward_batch, pp_proxy_tensors=pp_proxy_tensors)
         else: