[DLLM] Implement initial dynamic batching for diffusion LLM (#14883)

python/sglang/srt/dllm/algorithm/low_confidence.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, Union
+from typing import List, Tuple, Union
 
 import numpy as np
 import torch
@@ -24,11 +24,29 @@ class LowConfidence(DllmAlgorithm):
         self,
         model_runner: ModelRunner,
         forward_batch: ForwardBatch,
-    ) -> Tuple[
-        Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor], bool
-    ]:
+    ) -> Tuple[Union[LogitsProcessorOutput, torch.Tensor], List[torch.Tensor], bool]:
+        batch_size = forward_batch.batch_size
+        # Here, the forward_batch full logits contains all the blocks
+        # such as [dllm_block_size * batch_size, hidden_size]
+        start_list = []
         mask_index = forward_batch.input_ids == self.mask_id
-        start = len(forward_batch.input_ids) - torch.sum(mask_index).item()
+
+        # Fast path: if there is no mask token, forward and save kv cache
+        if torch.sum(mask_index).item() == 0:
+            out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
+            logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
+
+            next_token_ids = []
+            return logits_output, next_token_ids, can_run_cuda_graph
+
+        # Calculate start positions for each block
+        for block_id in range(batch_size):
+            block_start = block_id * self.block_size
+            block_end = block_start + self.block_size
+            block_input_ids = forward_batch.input_ids[block_start:block_end]
+            block_mask_index = block_input_ids == self.mask_id
+            start = self.block_size - torch.sum(block_mask_index).item()
+            start_list.append(start)
 
         for _ in range(self.block_size):
             mask_index = forward_batch.input_ids == self.mask_id
@@ -37,31 +55,50 @@ class LowConfidence(DllmAlgorithm):
 
             out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
             logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
+            assert batch_size == forward_batch.input_ids.shape[0] // self.block_size
+            for batch_id in range(batch_size):
+                curr_block_start = batch_id * self.block_size
+                curr_block_end = curr_block_start + self.block_size
+                block_input_ids = forward_batch.input_ids[
+                    curr_block_start:curr_block_end,
+                ]
+                block_mask_index = block_input_ids == self.mask_id
+                if torch.sum(block_mask_index).item() == 0:
+                    continue
+                curr_logits = logits_output.full_logits[
+                    curr_block_start:curr_block_end,
+                ]
 
-            x = torch.argmax(logits_output.full_logits, dim=-1)
-            p = torch.squeeze(
-                torch.gather(
-                    F.softmax(logits_output.full_logits, dim=-1),
-                    dim=-1,
-                    index=torch.unsqueeze(x, -1),
-                ),
-                -1,
-            )
-            x = torch.where(mask_index, x, forward_batch.input_ids)
-            confidence = torch.where(mask_index, p, -np.inf)
+                x = torch.argmax(curr_logits, dim=-1)
+                p = torch.squeeze(
+                    torch.gather(
+                        F.softmax(curr_logits, dim=-1),
+                        dim=-1,
+                        index=torch.unsqueeze(x, -1),
+                    ),
+                    -1,
+                )
+                x = torch.where(block_mask_index, x, block_input_ids)
+                confidence = torch.where(block_mask_index, p, -np.inf)
 
-            transfer_index = confidence > self.threshold
-            if transfer_index.sum().item() == 0:
-                _, select_index = torch.topk(confidence, k=1)
-                transfer_index[select_index] = True
+                transfer_index = confidence > self.threshold
 
-            forward_batch.input_ids[transfer_index] = x[transfer_index]
+                if transfer_index.sum().item() == 0:
+                    _, select_index = torch.topk(confidence, k=1)
+                    transfer_index[select_index] = True
+
+                block_input_ids[transfer_index] = x[transfer_index]
 
         out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
         logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
+        # Here next token ids is tricky to implement the dynamic lengths,
+        # so we return a list of tensors
+        next_token_ids = torch.reshape(forward_batch.input_ids, (batch_size, -1))
+        next_token_ids_list = [
+            next_token_ids[i, start_list[i] :] for i in range(batch_size)
+        ]
 
-        next_token_ids = forward_batch.input_ids[start:]
-        return logits_output, next_token_ids, can_run_cuda_graph
+        return logits_output, next_token_ids_list, can_run_cuda_graph
 
 
 Algorithm = LowConfidence

python/sglang/srt/dllm/config.py

@@ -11,11 +11,13 @@ class DllmConfig:
         algorithm_config: dict[str, Any],
         block_size: int,
         mask_id: int,
+        max_running_requests: int,
     ):
         self.algorithm = algorithm
         self.algorithm_config = algorithm_config
         self.block_size = block_size
         self.mask_id = mask_id
+        self.max_running_requests = max_running_requests
 
     @staticmethod
     def from_server_args(
@@ -38,6 +40,12 @@ class DllmConfig:
                 f"Unknown diffusion LLM: {model_config.hf_config.architectures[0]}"
             )
 
+        max_running_requests = (
+            1
+            if server_args.max_running_requests is None
+            else server_args.max_running_requests
+        )
+
         algorithm_config = {}
         if server_args.dllm_algorithm_config is not None:
             try:
@@ -58,4 +66,5 @@ class DllmConfig:
             algorithm_config=algorithm_config,
             block_size=block_size,
             mask_id=mask_id,
+            max_running_requests=max_running_requests,
         )

python/sglang/srt/managers/schedule_batch.py

@@ -1159,6 +1159,62 @@ class Req:
         )
 
 
+class DllmStagingReqs:
+    def __init__(self, dllm_config: Optional[DllmConfig] = None):
+        self.dllm_config = dllm_config
+        self.max_running_reqs = (
+            dllm_config.max_running_requests if dllm_config is not None else 1
+        )
+        self.reqs: List[Req] = []
+
+    def add_reqs(self, req: Union[Req, List[Req], "DllmStagingReqs"]):
+        assert self.dllm_config is not None, "Diffusion LLM config is not set."
+
+        if isinstance(req, DllmStagingReqs):
+            reqs_to_add = req.reqs
+        elif isinstance(req, list):
+            reqs_to_add = req
+        else:
+            reqs_to_add = [req]
+
+        num_to_add = len(reqs_to_add)
+
+        # Sanity check:
+        if self.check_redundant_reqs(reqs_to_add):
+            raise RuntimeError("Redundant requests detected in dLLM requests.")
+
+        if len(self.reqs) + num_to_add > self.max_running_reqs:
+            raise RuntimeError(
+                f"Exceeding maximum number of concurrent diffusion LLM requests: {self.max_running_reqs}"
+            )
+
+        self.reqs.extend(reqs_to_add)
+
+    def check_redundant_reqs(self, reqs: List[Req]) -> bool:
+        existing_rids: Set[str] = {r.rid for r in self.reqs}
+        return any(req.rid in existing_rids for req in reqs)
+
+    def init_next_round(self):
+        for req in self.reqs:
+            req.init_next_round_input()
+
+    def non_empty(self) -> bool:
+        return self.dllm_config is not None and len(self.reqs) > 0
+
+    def empty(self) -> bool:
+        return self.dllm_config is None or len(self.reqs) == 0
+
+    def update_chunked_status(self):
+        for req in self.reqs:
+            req.is_chunked += 1
+
+    def filter_finished_reqs(self):
+        self.reqs = [req for req in self.reqs if not req.finished()]
+
+    def __iter__(self):
+        return iter(self.reqs)
+
+
 @dataclasses.dataclass
 class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     """Store all information of a batch on the scheduler."""
@@ -1279,6 +1335,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     hicache_consumer_index: int = -1
 
     # Diffusion LLM
+    dllm_staging_reqs: Optional[DllmStagingReqs] = None
     dllm_config: Optional[DllmConfig] = None
 
     # Metrics
@@ -1295,6 +1352,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         enable_overlap: bool,
         spec_algorithm: SpeculativeAlgorithm,
         chunked_req: Optional[Req] = None,
+        dllm_staging_reqs: Optional[DllmStagingReqs] = None,
         dllm_config: Optional[DllmConfig] = None,
     ):
         return_logprob = any(req.return_logprob for req in reqs)
@@ -1320,6 +1378,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             return_routed_experts=any(req.return_routed_experts for req in reqs),
             is_prefill_only=all(req.is_prefill_only for req in reqs),
             chunked_req=chunked_req,
+            dllm_staging_reqs=dllm_staging_reqs,
             dllm_config=dllm_config,
         )
 

python/sglang/srt/managers/schedule_policy.py

@@ -32,8 +32,9 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
 
 import torch
 
+from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
-from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+from sglang.srt.managers.schedule_batch import DllmStagingReqs, Req, ScheduleBatch
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
 from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
@@ -373,6 +374,7 @@ class PrefillAdder:
         priority_scheduling_preemption_threshold: int = 0,
         prefill_max_requests: Optional[int] = None,
         prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
+        dllm_config: Optional[DllmConfig] = None,
     ):
         self.page_size = page_size
         self.tree_cache = tree_cache
@@ -381,6 +383,11 @@ class PrefillAdder:
         self.new_token_ratio = new_token_ratio
         self.rem_input_tokens = rem_input_tokens - mixed_with_decode_tokens
         self.rem_chunk_tokens = rem_chunk_tokens
+        self.dllm_config = dllm_config
+
+        if self.dllm_config is not None:
+            self._init_dllm_meta(dllm_config)
+
         if self.rem_chunk_tokens is not None:
             self.rem_chunk_tokens -= mixed_with_decode_tokens
         self.rem_total_token_offset = mixed_with_decode_tokens
@@ -414,6 +421,13 @@ class PrefillAdder:
         self.prefill_max_requests = prefill_max_requests
         self.prefill_delayer_single_pass = prefill_delayer_single_pass
 
+    def _init_dllm_meta(self, dllm_config: DllmConfig):
+        self.dllm_block_size = dllm_config.block_size
+        max_running_reqs = dllm_config.max_running_requests
+
+        self.rem_dllm_tokens = max_running_reqs * self.dllm_block_size
+        self.dllm_staging_reqs = DllmStagingReqs(dllm_config=dllm_config)
+
     def _get_running_request_total_token_offset(self, req: Req) -> int:
         return (
             min(
@@ -473,11 +487,16 @@ class PrefillAdder:
         if self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0:
             return AddReqResult.NO_TOKEN
 
-        if self.rem_input_tokens <= 0 or (
-            self.rem_chunk_tokens is not None and self.rem_chunk_tokens <= 0
-        ):
+        if self.rem_input_tokens <= 0:
             return AddReqResult.OTHER
 
+        if self.dllm_config is not None:
+            if self.rem_dllm_tokens <= 0:
+                return AddReqResult.OTHER
+        else:
+            if self.rem_chunk_tokens is not None and self.rem_chunk_tokens <= 0:
+                return AddReqResult.OTHER
+
         return AddReqResult.CONTINUE
 
     def _update_prefill_budget(
@@ -489,18 +508,61 @@ class PrefillAdder:
         self.rem_total_token_offset += extend_input_len + max_new_tokens
         self.cur_rem_token_offset += extend_input_len
         self.rem_input_tokens -= extend_input_len
-        if self.rem_chunk_tokens is not None:
+
+        if self.dllm_config is not None:
+            self.rem_dllm_tokens -= extend_input_len
+        elif self.rem_chunk_tokens is not None:
             self.rem_chunk_tokens -= extend_input_len
 
         self.log_hit_tokens += prefix_len
         self.log_input_tokens += extend_input_len
 
-    def add_chunked_req(self, req: Req):
-        _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens))
-        # The chunked_req must be added to the list; otherwise, it will cause a memory leak.
-        # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens.
+    def _get_dllm_remain_tokens(self) -> int:
+        _rem_tokens = min(
+            self.rem_dllm_tokens,
+            self.dllm_block_size,
+            int(self.rem_total_tokens),
+        )
         if _rem_tokens <= 0:
-            _rem_tokens = self.rem_chunk_tokens
+            _rem_tokens = self.rem_dllm_tokens
+
+        return _rem_tokens
+
+    def _add_dllm_req(self, req: Req, prefix_len: int):
+        # FIXME: consider the case when rem_dllm_tokens < dllm_block_size,
+        # the diffusion unmask process may have some problems
+        # Make sure at least one page is available
+        trunc_len = (
+            min(self.rem_dllm_tokens, self.dllm_block_size)
+            // self.page_size
+            * self.page_size
+        )
+
+        req.extend_input_len = trunc_len
+        req.fill_ids = req.fill_ids[: prefix_len + trunc_len]
+
+        self.can_run_list.append(req)
+        self.dllm_staging_reqs.add_reqs(req)
+
+        self._update_prefill_budget(prefix_len, trunc_len, 0)
+
+    def _req_inc_lock_ref(self, req: Req):
+        if self.is_hybrid_swa:
+            swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node)
+            req.swa_uuid_for_lock = swa_uuid_for_lock
+        else:
+            self.tree_cache.inc_lock_ref(req.last_node)
+
+    def add_chunked_req(self, req: Req):
+        if self.dllm_config is not None:
+            _rem_tokens = self._get_dllm_remain_tokens()
+        else:
+            _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens))
+            # The chunked_req must be added to the list; otherwise, it will cause a memory leak.
+            # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens.
+            if _rem_tokens <= 0:
+                _rem_tokens = self.rem_chunk_tokens
+
         truncated = req.extend_input_len > _rem_tokens
         req.set_extend_input_len(min(req.extend_input_len, _rem_tokens))
         req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
@@ -588,7 +650,12 @@ class PrefillAdder:
                     return AddReqResult.NO_TOKEN
                 tokens_freed += tokens_occupied
 
-        if (
+        if self.dllm_config is not None:
+            if self.rem_dllm_tokens <= 0:
+                return AddReqResult.OTHER
+
+            self._add_dllm_req(req, 0)
+        elif (
             self.rem_chunk_tokens is None  # chunked prefill is disabled
             or req.extend_input_len <= self.rem_chunk_tokens  # it is the last chunk
         ):
@@ -671,14 +738,21 @@ class PrefillAdder:
             ):
                 return AddReqResult.OTHER
 
-            if self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens:
+            if self.dllm_config is not None:
+                if self.rem_dllm_tokens <= 0:
+                    return AddReqResult.OTHER
+
+                assert (
+                    truncation_align_size is None
+                ), "truncation_align_size is not supported for dllm prefill"
+
+                self._add_dllm_req(req, prefix_len)
+                self._req_inc_lock_ref(req)
+            elif self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens:
                 # Non-chunked prefill
                 self.can_run_list.append(req)
-                if self.is_hybrid_swa:
-                    swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node)
-                    req.swa_uuid_for_lock = swa_uuid_for_lock
-                else:
-                    self.tree_cache.inc_lock_ref(req.last_node)
+
+                self._req_inc_lock_ref(req)
                 self._update_prefill_budget(
                     prefix_len,
                     input_tokens,
@@ -690,6 +764,7 @@ class PrefillAdder:
             else:
                 # Make sure at least one page is available
                 trunc_len = self.rem_chunk_tokens // self.page_size * self.page_size
+
                 if trunc_len <= 0:
                     return AddReqResult.OTHER
 
@@ -710,11 +785,8 @@ class PrefillAdder:
 
                 self.can_run_list.append(req)
                 self.new_chunked_req = req
-                if self.is_hybrid_swa:
-                    swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node)
-                    req.swa_uuid_for_lock = swa_uuid_for_lock
-                else:
-                    self.tree_cache.inc_lock_ref(req.last_node)
+
+                self._req_inc_lock_ref(req)
                 self._update_prefill_budget(prefix_len, trunc_len, 0)
 
         return self.budget_state()

python/sglang/srt/managers/scheduler.py

@@ -136,6 +136,7 @@ from sglang.srt.managers.schedule_batch import (
 )
 from sglang.srt.managers.schedule_policy import (
     AddReqResult,
+    DllmStagingReqs,
     PrefillAdder,
     SchedulePolicy,
 )
@@ -351,6 +352,9 @@ class Scheduler(
         # Init chunked prefill
         self.init_chunked_prefill()
 
+        # Init diffusion LLM
+        self.init_diffusion_llm()
+
         # Init schedule policy and new token estimation
         self.init_schedule_policy()
 
@@ -726,10 +730,6 @@ class Scheduler(
     def init_chunked_prefill(self):
         # Init chunked prefill
         self.chunked_prefill_size = self.server_args.chunked_prefill_size
-        if self.dllm_config is not None:
-            # We currently leverage chunked prefill to implement block diffusion
-            # for diffusion LLM.
-            self.chunked_prefill_size = self.dllm_config.block_size
         if self.chunked_prefill_size <= 0:  # -1 means disable
             self.chunked_prefill_size = None
         self.chunked_req = None
@@ -752,6 +752,9 @@ class Scheduler(
                 )
                 self.enable_dynamic_chunking = False
 
+    def init_diffusion_llm(self):
+        self.dllm_staging_reqs = DllmStagingReqs(dllm_config=self.dllm_config)
+
     def init_schedule_policy(self):
         # Init schedule policy and new token estimation
         self.policy = SchedulePolicy(
@@ -1297,6 +1300,7 @@ class Scheduler(
             # If it is a health check generation request and there are running requests, ignore it.
             if is_health_check_generate_req(recv_req) and (
                 self.chunked_req is not None
+                or self.dllm_staging_reqs.non_empty()
                 or not self.running_batch.is_empty()
                 or len(self.offload_tags) > 0
             ):
@@ -1765,27 +1769,37 @@ class Scheduler(
         for tokenized_req in recv_req:
             self.handle_embedding_request(tokenized_req)
 
+    def stash_chunked_request(self, req: Req):
+        self.tree_cache.cache_unfinished_req(req, chunked=True)
+        # Chunked request keeps its rid but will get a new req_pool_idx
+        if self.tp_worker.model_runner.mambaish_config is not None:
+            self.req_to_token_pool.free(req.req_pool_idx, free_mamba_cache=False)
+        else:
+            self.req_to_token_pool.free(req.req_pool_idx)
+
     def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
         self._abort_on_queued_timeout()
         if self.dllm_config is not None:
-            if self.chunked_req is not None and self.chunked_req.finished():
-                self.chunked_req = None
+            self.dllm_staging_reqs.filter_finished_reqs()
 
         # Merge the prefill batch into the running batch
         chunked_req_to_exclude = set()
-        if self.chunked_req:
+
+        if self.dllm_config is not None:
+            assert (
+                self.chunked_req is None
+            ), "chunked_req should be None when dllm_config is set"
+
+            if self.dllm_staging_reqs.non_empty():
+                chunked_req_to_exclude.update(self.dllm_staging_reqs)
+                for req in self.dllm_staging_reqs:
+                    self.stash_chunked_request(req)
+
+        if self.chunked_req is not None:
             # Move the chunked request out of the batch so that we can merge
             # only finished requests to running_batch.
             chunked_req_to_exclude.add(self.chunked_req)
-            self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
-
-            # chunked request keeps its rid but will get a new req_pool_idx
-            if self.tp_worker.model_runner.mambaish_config is not None:
-                self.req_to_token_pool.free(
-                    self.chunked_req.req_pool_idx, free_mamba_cache=False
-                )
-            else:
-                self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
+            self.stash_chunked_request(self.chunked_req)
 
         if self.last_batch and self.last_batch.forward_mode.is_extend():
             if self.last_batch.chunked_req is not None:
@@ -1793,6 +1807,9 @@ class Scheduler(
                 # We need to discard it.
                 chunked_req_to_exclude.add(self.last_batch.chunked_req)
 
+            if self.last_batch.dllm_staging_reqs.non_empty():
+                chunked_req_to_exclude.update(self.last_batch.dllm_staging_reqs)
+
             # Filter batch
             last_bs = self.last_batch.batch_size()
             self.last_batch.filter_batch(
@@ -1880,20 +1897,20 @@ class Scheduler(
             # Reset batch_is_full to try preemption with a prefill adder.
             self.running_batch.batch_is_full = False
 
-        if (
-            self.running_batch.batch_is_full or len(self.waiting_queue) == 0
-        ) and self.chunked_req is None:
+        if (self.running_batch.batch_is_full or len(self.waiting_queue) == 0) and (
+            not self.dllm_staging_reqs.non_empty() and self.chunked_req is None
+        ):
             return None
 
         running_bs = len(self.running_batch.reqs)
         # Ignore the check if self.chunked_req is not None.
         # In the non-PP case, when self.chunked_req is not None, num_allocatable_reqs should always be greater than 0,
-        # as the space for the chunked request has just been released.
-        # In PP case, a chunked req can start in one microbatch and end in another microbatch, so the max_running_requests per microbatch should not be strict.
-        # Instead, we should always allow chunked request to be added, otherwise, there will be a memory leak.
+        # as the space for the chunked requests has just been released.
+        # In PP case, chunked requests (or dllm requests) can start in one microbatch and end in another microbatch, so the max_running_requests per microbatch should not be strict.
+        # Instead, we should always allow chunked requests to be added, otherwise, there will be a memory leak.
         if (
             self.get_num_allocatable_reqs(running_bs) <= 0
-            and not self.chunked_req
+            and (self.dllm_staging_reqs.empty() or self.chunked_req is not None)
             and not self.try_preemption
         ):
             self.running_batch.batch_is_full = True
@@ -1932,8 +1949,20 @@ class Scheduler(
             self.priority_scheduling_preemption_threshold,
             prefill_max_requests=self.server_args.prefill_max_requests,
             prefill_delayer_single_pass=prefill_delayer_single_pass,
+            dllm_config=self.dllm_config,
         )
 
+        if self.dllm_config is not None:
+            assert (
+                self.chunked_req is None
+            ), "chunked_req should be None when dllm_config is set"
+
+            if self.dllm_staging_reqs.non_empty():
+                self.dllm_staging_reqs.init_next_round()
+                for req in self.dllm_staging_reqs:
+                    adder.add_chunked_req(req)
+                self.dllm_staging_reqs.add_reqs(adder.dllm_staging_reqs)
+
         if self.chunked_req is not None:
             self.chunked_req.init_next_round_input()
             self.chunked_req = adder.add_chunked_req(self.chunked_req)
@@ -1981,7 +2010,9 @@ class Scheduler(
             req.init_next_round_input(self.tree_cache)
             res = adder.add_one_req(
                 req,
-                has_chunked_req=(self.chunked_req is not None),
+                has_chunked_req=(
+                    self.dllm_staging_reqs.non_empty() or self.chunked_req is not None
+                ),
                 truncation_align_size=self.truncation_align_size,
             )
 
@@ -2013,14 +2044,25 @@ class Scheduler(
             for req in adder.preempt_list:
                 self._add_request_to_queue(req)
 
-        # Update chunked prefill
+        if self.dllm_config is not None:
+            assert (
+                self.chunked_req is None
+            ), "chunked_req should be None when dllm_config is set"
+
+            if adder.dllm_staging_reqs.non_empty():
+                self.dllm_staging_reqs.add_reqs(adder.dllm_staging_reqs)
+
         if adder.new_chunked_req is not None:
+            # Update chunked prefill
             assert self.chunked_req is None
             self.chunked_req = adder.new_chunked_req
 
-        if self.chunked_req:
+        if self.chunked_req is not None:
             self.chunked_req.is_chunked += 1
 
+        if self.dllm_staging_reqs.non_empty():
+            self.dllm_staging_reqs.update_chunked_status()
+
         # Print stats
         if self.current_scheduler_metrics_enabled:
             self.log_prefill_stats(
@@ -2049,6 +2091,7 @@ class Scheduler(
             self.enable_overlap,
             self.spec_algorithm,
             chunked_req=self.chunked_req,
+            dllm_staging_reqs=self.dllm_staging_reqs,
             dllm_config=self.dllm_config,
         )
         if self.enable_hierarchical_cache:

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -332,24 +332,26 @@ class SchedulerOutputProcessorMixin:
         if result.copy_done is not None:
             result.copy_done.synchronize()
 
-        next_token_ids = result.next_token_ids.tolist()
-        self.num_generated_tokens += len(next_token_ids)
-
         self.token_to_kv_pool_allocator.free_group_begin()
 
-        assert len(batch.reqs) == 1, "batch size is currently expected to be 1"
-        req = batch.reqs[0]
-
-        for next_token_id in next_token_ids:
-            req.output_ids.append(next_token_id)
-            req.check_finished()
-
-            if req.finished():
-                release_kv_cache(req, self.tree_cache)
-                req.time_stats.completion_time = time.perf_counter()
+        for idx in range(batch.batch_size()):
+            # If no new tokens generated, meaning the prefilling stage
+            if not result.next_token_ids:
                 break
 
-            self.tree_cache.cache_unfinished_req(req)
+            req = batch.reqs[idx]
+            next_token_ids = result.next_token_ids[idx].tolist()
+            self.num_generated_tokens += len(next_token_ids)
+
+            for _token_idx, next_token_id in enumerate(next_token_ids):
+                req.output_ids.append(next_token_id)
+                req.check_finished()
+                if req.finished():
+                    release_kv_cache(req, self.tree_cache)
+                    req.time_stats.completion_time = time.perf_counter()
+                    break
+
+                self.tree_cache.cache_unfinished_req(req)
 
         self.stream_output(batch.reqs, batch.return_logprob)
         self.token_to_kv_pool_allocator.free_group_end()

python/sglang/srt/managers/utils.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import dataclasses
 import logging
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, List, Optional, Union
 
 import torch
 
@@ -24,7 +24,7 @@ logger = logging.getLogger(__name__)
 class GenerationBatchResult:
     logits_output: Optional[LogitsProcessorOutput] = None
     pp_hidden_states_proxy_tensors: Optional[PPProxyTensors] = None
-    next_token_ids: Optional[torch.Tensor] = None
+    next_token_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None
     num_accepted_tokens: int = 0
     accept_length_per_req_cpu: Optional[List[int]] = None
     can_run_cuda_graph: bool = False