chore: import upstream snapshot with attribution

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

@@ -0,0 +1,39 @@
+import importlib
+import logging
+import pkgutil
+
+from sglang.srt.dllm.config import DllmConfig
+
+logger = logging.getLogger(__name__)
+
+
+def import_algorithms():
+    mapping = {}
+    package_name = "sglang.srt.dllm.algorithm"
+    package = importlib.import_module(package_name)
+    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
+        if ispkg:
+            continue
+        try:
+            module = importlib.import_module(name)
+        except Exception as e:
+            logger.warning(f"Ignore import error when loading {name}: {e}")
+            continue
+        if not hasattr(module, "Algorithm"):
+            continue
+
+        algo = module.Algorithm
+        mapping[algo.__name__] = algo
+
+    return mapping
+
+
+def get_algorithm(config: DllmConfig):
+    try:
+        name = config.algorithm
+        return algo_name_to_cls[name](config)
+    except:
+        raise RuntimeError(f"Unknown diffusion LLM algorithm: {name}")
+
+
+algo_name_to_cls = import_algorithms()

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

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from typing import Any, List, Optional, Tuple, Union
+
+import torch
+
+from sglang.srt.dllm.algorithm import get_algorithm
+from sglang.srt.dllm.config import DllmConfig
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.server_args import ServerArgs
+
+DllmRunOutput = Tuple[
+    Union[LogitsProcessorOutput, torch.Tensor],
+    List,
+    Optional[List[int]],
+    Optional[List[Any]],
+    bool,
+]
+
+
+class DllmAlgorithm:
+    """dLLM algorithm: subclasses implement ``step``; the base owns the
+    synchronous and FDFO (``--dllm-fdfo``) execution loops in ``run``.
+    """
+
+    def __init__(self, config: DllmConfig):
+        self.block_size = config.block_size
+        self.mask_id = config.mask_id
+        self.fdfo = config.first_done_first_out_mode
+
+    @staticmethod
+    def from_server_args(server_args: ServerArgs):
+        config = DllmConfig.from_server_args(server_args)
+        return get_algorithm(config)
+
+    def init_step_state(self, forward_batch: ForwardBatch) -> List[Any]:
+        return [None] * forward_batch.batch_size
+
+    def max_steps(self, block_size: int) -> int:
+        return block_size + 1
+
+    def step(
+        self,
+        forward_batch: ForwardBatch,
+        full_logits: torch.Tensor,
+        states: List[Any],
+    ) -> List[bool]:
+        """One denoise step, advancing ``forward_batch.input_ids``/``states`` in
+        place. Returns, per block, whether it was already complete *on entry* --
+        i.e. this forward persisted its final KV cache and it can be emitted.
+        """
+        raise NotImplementedError
+
+    def run(
+        self,
+        model_runner: ModelRunner,
+        forward_batch: ForwardBatch,
+        algo_states: Optional[List[Any]] = None,
+    ) -> DllmRunOutput:
+        if self.fdfo:
+            return self._run_fdfo(model_runner, forward_batch, algo_states)
+        return self._run_sync(model_runner, forward_batch)
+
+    def _block_start_list(self, forward_batch: ForwardBatch) -> List[int]:
+        batch_size = forward_batch.batch_size
+        input_ids = forward_batch.input_ids.view(batch_size, self.block_size)
+        return (input_ids != self.mask_id).sum(dim=1).tolist()
+
+    def _run_sync(
+        self, model_runner: ModelRunner, forward_batch: ForwardBatch
+    ) -> DllmRunOutput:
+        batch_size = forward_batch.batch_size
+        start_list = self._block_start_list(forward_batch)
+
+        out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
+        # No mask to denoise: return empty so process_batch_result_dllm skips the
+        # stream branch (matches the pre-refactor behavior).
+        if all(start == self.block_size for start in start_list):
+            return out.logits_output, [], None, None, out.can_run_graph
+
+        states = self.init_step_state(forward_batch)
+        for _ in range(self.max_steps(self.block_size)):
+            done = self.step(forward_batch, out.logits_output.full_logits, states)
+            if all(done):
+                break
+            out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
+
+        next_token_ids = forward_batch.input_ids.view(batch_size, self.block_size)
+        next_token_ids_list = [
+            next_token_ids[i, start_list[i] :] for i in range(batch_size)
+        ]
+        return out.logits_output, next_token_ids_list, None, None, out.can_run_graph
+
+    def _run_fdfo(
+        self,
+        model_runner: ModelRunner,
+        forward_batch: ForwardBatch,
+        algo_states: Optional[List[Any]],
+    ) -> DllmRunOutput:
+        batch_size = forward_batch.batch_size
+
+        if algo_states is None:
+            algo_states = [None] * batch_size
+        fresh: Optional[List[Any]] = None
+        states: List[Any] = []
+        for i, carried in enumerate(algo_states):
+            if carried is None:
+                if fresh is None:
+                    fresh = self.init_step_state(forward_batch)
+                states.append(fresh[i])
+            else:
+                states.append(carried)
+
+        out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
+        done = self.step(forward_batch, out.logits_output.full_logits, states)
+
+        accept_length_per_req_cpu = [self.block_size if d else 0 for d in done]
+        next_token_ids_list = forward_batch.input_ids.view(
+            batch_size, self.block_size
+        ).tolist()
+        states_out = [None if done[i] else states[i] for i in range(batch_size)]
+
+        return (
+            out.logits_output,
+            next_token_ids_list,
+            accept_length_per_req_cpu,
+            states_out,
+            out.can_run_graph,
+        )

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

@@ -0,0 +1,119 @@
+from typing import Any, List
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.dllm.algorithm.base import DllmAlgorithm
+from sglang.srt.dllm.config import DllmConfig
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class JointThreshold(DllmAlgorithm):
+    """Joint-threshold denoising: mask-to-token (M2T) unmasking plus token-to-token
+    (T2T) edits, finishing on no-change or an exhausted edit budget. Stateful (edit
+    budget + prompt mask), carried across FDFO rounds via ``dllm_algo_state``.
+    """
+
+    def __init__(self, config: DllmConfig):
+        super().__init__(config)
+        self.threshold = config.algorithm_config.get("threshold", 0.5)
+        self.edit_threshold = config.algorithm_config.get("edit_threshold", 0)
+        self.max_post_edit_steps = config.algorithm_config.get(
+            "max_post_edit_steps", 16
+        )
+        self.penalty_lambda = config.algorithm_config.get("penalty_lambda", 0)
+
+    def max_steps(self, block_size: int) -> int:
+        return block_size + self.max_post_edit_steps + 1
+
+    def init_step_state(self, forward_batch: ForwardBatch) -> List[Any]:
+        batch_size = forward_batch.batch_size
+        input_ids = forward_batch.input_ids.view(batch_size, self.block_size)
+        # Built once as a GPU tensor and reused across steps (no per-step
+        # host/device transfer); the FDFO carry keeps it in-process.
+        prompt_mask = input_ids != self.mask_id
+        return [
+            {
+                "post_edit_steps": 0,
+                "finished": False,
+                "prompt_mask": prompt_mask[i],
+            }
+            for i in range(batch_size)
+        ]
+
+    def step(
+        self,
+        forward_batch: ForwardBatch,
+        full_logits: torch.Tensor,
+        states: List[Any],
+    ) -> List[bool]:
+        batch_size = forward_batch.batch_size
+        done: List[bool] = []
+
+        for i in range(batch_size):
+            state = states[i]
+            if state["finished"]:
+                done.append(True)
+                continue
+
+            block_start = i * self.block_size
+            block_end = block_start + self.block_size
+            curr_input_ids = forward_batch.input_ids[block_start:block_end]
+            curr_logits = full_logits[block_start:block_end]
+            curr_prompt_mask = state["prompt_mask"]
+
+            if self.penalty_lambda > 0:
+                prev_ids = curr_input_ids[:-1]
+                curr_logits[1:, :].scatter_(
+                    1, prev_ids.unsqueeze(-1), -self.penalty_lambda, reduce="add"
+                )
+
+            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,
+            )
+
+            mask_index = curr_input_ids == self.mask_id
+            has_mask = mask_index.any()
+
+            # Mask to token (M2T)
+            mask_transfer_index = torch.zeros_like(mask_index)
+            budget_exhausted = False
+            if has_mask:
+                confidence = torch.where(mask_index, p, -np.inf)
+                mask_transfer_index = confidence > self.threshold
+                if not mask_transfer_index.any():
+                    _, select_index = torch.topk(confidence, k=1)
+                    mask_transfer_index[select_index] = True
+            else:
+                state["post_edit_steps"] += 1
+                if state["post_edit_steps"] > self.max_post_edit_steps:
+                    state["finished"] = True
+                    budget_exhausted = True
+
+            if not budget_exhausted:
+                # Token to token (T2T)
+                edit_mask = ~mask_index & ~curr_prompt_mask
+                edit_transfer_index = (
+                    (p > self.edit_threshold) & (curr_input_ids != x) & edit_mask
+                )
+                transfer_index = mask_transfer_index | edit_transfer_index
+                if transfer_index.any():
+                    curr_input_ids[transfer_index] = x[transfer_index]
+                else:
+                    state["finished"] = True
+
+            # A terminating step changes nothing, so this forward already holds the
+            # block's final KV: emit it now rather than after an extra forward.
+            done.append(state["finished"])
+
+        return done
+
+
+Algorithm = JointThreshold

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

@@ -0,0 +1,55 @@
+from typing import Any, List
+
+import torch
+
+from sglang.srt.dllm.algorithm.base import DllmAlgorithm
+from sglang.srt.dllm.config import DllmConfig
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class LowConfidence(DllmAlgorithm):
+    """Each step unmasks positions whose predicted-token confidence exceeds a
+    threshold (falling back to the highest-confidence masked position).
+    """
+
+    def __init__(self, config: DllmConfig):
+        super().__init__(config)
+        self.threshold = config.algorithm_config.get("threshold", 0.95)
+
+    def step(
+        self,
+        forward_batch: ForwardBatch,
+        full_logits: torch.Tensor,
+        states: List[Any],
+    ) -> List[bool]:
+        batch_size = forward_batch.batch_size
+        vocab_size = full_logits.shape[-1]
+        logits = full_logits.view(batch_size, self.block_size, vocab_size)
+        input_ids = forward_batch.input_ids.view(batch_size, self.block_size)
+        block_mask_index = input_ids == self.mask_id
+        done = block_mask_index.sum(dim=1) == 0
+
+        x = torch.argmax(logits, dim=-1)
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        confidence = torch.gather(probs, dim=-1, index=x.unsqueeze(-1)).squeeze(-1)
+        confidence = torch.where(block_mask_index, confidence, -float("inf"))
+
+        transfer_index = confidence > self.threshold
+        has_transfer = transfer_index.sum(dim=1) > 0
+        top1_indices = torch.argmax(confidence, dim=1)
+        batch_indices = torch.arange(batch_size, device=top1_indices.device)
+        top1_mask = torch.zeros_like(transfer_index, dtype=torch.bool)
+        top1_mask[batch_indices, top1_indices] = True
+        transfer_index = torch.where(
+            has_transfer.unsqueeze(-1), transfer_index, top1_mask
+        )
+
+        x = torch.where(block_mask_index, x, input_ids)
+        new_input_ids = torch.where(transfer_index, x, input_ids)
+        # In-place to preserve the input_ids tensor identity (CUDA graph safe).
+        forward_batch.input_ids.copy_(new_input_ids.view(-1))
+
+        return done.tolist()
+
+
+Algorithm = LowConfidence

python/sglang/srt/dllm/config.py

@@ -0,0 +1,78 @@
+from typing import Any
+
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.server_args import ServerArgs
+
+
+class DllmConfig:
+    def __init__(
+        self,
+        algorithm: str,
+        algorithm_config: dict[str, Any],
+        block_size: int,
+        mask_id: int,
+        max_running_requests: int,
+        first_done_first_out_mode: bool = False,
+    ):
+        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
+        self.first_done_first_out_mode = first_done_first_out_mode
+
+    @staticmethod
+    def from_server_args(
+        server_args: ServerArgs,
+    ):
+        if server_args.dllm_algorithm is None:
+            return None
+
+        model_config = ModelConfig.from_server_args(
+            server_args,
+            model_path=server_args.model_path,
+            model_revision=server_args.revision,
+        )
+        DLLM_PARAMS = {
+            "LLaDA2MoeModelLM": {"block_size": 32, "mask_id": 156895},
+            "SDARForCausalLM": {"block_size": 4, "mask_id": 151669},
+            "SDARMoeForCausalLM": {"block_size": 4, "mask_id": 151669},
+        }
+
+        arch = model_config.hf_config.architectures[0]
+        if arch in DLLM_PARAMS:
+            params = DLLM_PARAMS[arch]
+            block_size = params["block_size"]
+            mask_id = params["mask_id"]
+        else:
+            raise RuntimeError(f"Unknown diffusion LLM: {arch}")
+
+        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:
+                import yaml
+            except ImportError:
+                raise ImportError(
+                    "Please install PyYAML to use YAML config files. "
+                    "`pip install pyyaml`"
+                )
+            with open(server_args.dllm_algorithm_config, "r") as f:
+                algorithm_config = yaml.safe_load(f)
+
+            # Parse common algorithm configurations
+            block_size = algorithm_config.get("block_size", block_size)
+
+        return DllmConfig(
+            algorithm=server_args.dllm_algorithm,
+            algorithm_config=algorithm_config,
+            block_size=block_size,
+            mask_id=mask_id,
+            max_running_requests=max_running_requests,
+            first_done_first_out_mode=server_args.dllm_fdfo,
+        )

python/sglang/srt/dllm/mixin/req.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import enum
+from array import array
+from typing import TYPE_CHECKING, Optional
+
+from sglang.srt.dllm.config import DllmConfig
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
+
+class DllmReqPhase(str, enum.Enum):
+    STAGING_PREFILL = "staging_prefill"
+    STAGING_DECODE = "staging_decode"
+    INCOMING_PREFILL = "incoming_prefill"
+    INCOMING_DECODE = "incoming_decode"
+
+
+class ReqDllmMixin:
+    def init_diffusion_llm(self: Req, dllm_config: DllmConfig):
+        self.dllm_phase: Optional[DllmReqPhase] = None
+        self.dllm_incomplete_ids = array("q")
+        self.dllm_algo_state = None
+        self.dllm_block_offset = 0
+        self.dllm_config = dllm_config
+
+        if self.dllm_config is not None:
+            if len(self.origin_input_ids) < self.dllm_config.block_size:
+                self.dllm_phase = DllmReqPhase.INCOMING_DECODE
+            else:
+                self.dllm_phase = DllmReqPhase.INCOMING_PREFILL
+
+    def is_dllm(self: Req) -> bool:
+        return self.dllm_config is not None
+
+    def is_dllm_prefill(self: Req) -> bool:
+        return self.dllm_phase in [
+            DllmReqPhase.STAGING_PREFILL,
+            DllmReqPhase.INCOMING_PREFILL,
+        ]
+
+    def determine_dllm_phase(self: Req):
+        if self.dllm_incomplete_ids:
+            self.dllm_phase = DllmReqPhase.STAGING_DECODE
+            return
+
+        prefix_length = len(self.prefix_indices)
+        min_required_length = prefix_length + self.dllm_config.block_size
+
+        if len(self.full_untruncated_fill_ids) < min_required_length:
+            # still incoming stage
+            return
+
+        input_block = self.full_untruncated_fill_ids[prefix_length:min_required_length]
+        is_prefill_phase = self.dllm_config.mask_id not in input_block
+
+        if is_prefill_phase:
+            self.dllm_phase = DllmReqPhase.STAGING_PREFILL
+        else:
+            self.dllm_phase = DllmReqPhase.STAGING_DECODE
+
+    def _init_fill_ids_for_dllm(self: Req):
+        if self.dllm_incomplete_ids:
+            prefix_len = len(self.prefix_indices)
+            assert len(self.dllm_incomplete_ids) == self.dllm_config.block_size
+            self.full_untruncated_fill_ids = (
+                self.full_untruncated_fill_ids[:prefix_len] + self.dllm_incomplete_ids
+            )
+            # extend_range is (re)computed by the staging adder
+            # (add_dllm_staging_req) before this req is scheduled, mirroring the
+            # non-incomplete path which also defers it to the adder.
+            return
+
+        self.dllm_block_offset = (
+            0
+            if not self.dllm_initialized
+            else self.dllm_block_offset + self.dllm_config.block_size
+        )
+        self.full_untruncated_fill_ids = (
+            self.origin_input_ids
+            + self.output_ids
+            + array("q", [self.dllm_config.mask_id] * self.dllm_config.block_size)
+        )
+        self.dllm_initialized = True
+
+    def _update_block_offset_for_dllm(self):
+        prefix_len = len(self.prefix_indices)
+        assert (
+            prefix_len % self.dllm_config.block_size == 0
+        ), f"Unexpected prefix len: {prefix_len}"
+        if prefix_len > self.dllm_block_offset:
+            self.dllm_block_offset = prefix_len

python/sglang/srt/dllm/mixin/scheduler.py

@@ -0,0 +1,420 @@
+from __future__ import annotations
+
+import logging
+from array import array
+from typing import TYPE_CHECKING, List, Optional, Set, Union
+
+from sglang.srt.dllm.config import DllmConfig
+from sglang.srt.dllm.mixin.req import DllmReqPhase
+from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder
+from sglang.srt.mem_cache.common import release_kv_cache
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.observability.req_time_stats import set_time_batch
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import GenerationBatchResult, Scheduler
+
+
+class SchedulerDllmMixin:
+    def init_diffusion_llm(self: Scheduler):
+        self.dllm_config = (
+            DllmConfig.from_server_args(self.server_args)
+            if self.server_args.dllm_algorithm is not None
+            else None
+        )
+        self.dllm_manager = DllmManager(dllm_config=self.dllm_config)
+
+    def get_new_batch_dllm(
+        self: Scheduler, running_batch: ScheduleBatch
+    ) -> Optional[ScheduleBatch]:
+        """Generate a new batch for DLLM (Diffusion LLM) scheduling."""
+        self.running_batch = running_batch
+        if self.enable_priority_preemption:
+            self.running_batch.batch_is_full = False
+
+        # Early exit if batch is full or no requests available
+        if self._should_skip_prefill():
+            return None
+
+        running_bs = len(self.running_batch.reqs)
+        self.policy.calc_priority(self.waiting_queue)
+
+        # Create prefill adder with resource constraints
+        adder = self._create_dllm_prefill_adder(running_bs)
+
+        # Initialize DLLM manager and transfer requests
+        self.dllm_manager.init_next_round()
+        self._fetch_waiting_reqs()
+
+        # Process batches
+        forward_mode = self._process_dllm_batches(adder)
+
+        can_run_list = adder.can_run_list
+        if not can_run_list:
+            return None
+
+        # Record metrics and update state
+        set_time_batch(can_run_list, "set_forward_entry_time")
+        self._update_state_for_batch(can_run_list, adder, running_bs)
+
+        # Create and prepare batch
+        new_batch = self._create_dllm_batch(can_run_list, forward_mode)
+        return new_batch
+
+    def process_batch_result_dllm(
+        self: Scheduler,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+    ):
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+
+        fdfo_mode = self.dllm_config.first_done_first_out_mode
+        assert (
+            not fdfo_mode or result.accept_length_per_req_cpu is not None
+        ), "FDFO dLLM result is missing accept lengths."
+
+        # Sync mode emits tokens only once a block fully resolves; FDFO always
+        # commits (resolved blocks decode, unresolved blocks stash + free KV).
+        if fdfo_mode or result.next_token_ids:
+            block_size = self.dllm_config.block_size
+            algo_states = result.dllm_algo_state
+
+            self.token_to_kv_pool_allocator.free_group_begin()
+            for idx in range(batch.batch_size()):
+                req = batch.reqs[idx]
+
+                if not fdfo_mode:
+                    next_token_ids = result.next_token_ids[idx].tolist()
+                    new_tokens = len(next_token_ids)
+                    if new_tokens == 0:
+                        continue
+
+                    req.full_untruncated_fill_ids[
+                        req.extend_range.end - new_tokens : req.extend_range.end
+                    ] = array("q", next_token_ids)
+                    self.metrics_reporter.num_generated_tokens += new_tokens
+
+                    req.output_ids.extend(next_token_ids)
+                    req.update_finish_state(new_accepted_len=new_tokens)
+
+                    if req.finished():
+                        release_kv_cache(req, self.tree_cache)
+                        req.time_stats.set_completion_time()
+                    continue
+
+                next_token_ids = result.next_token_ids[idx]
+                assert len(next_token_ids) == block_size
+
+                if result.accept_length_per_req_cpu[idx] == 0:
+                    # Block unresolved: stash partial state and free the KV slots
+                    # of the still-masked block so the next FDFO round can
+                    # re-denoise it without leaking the previous allocation.
+                    req.dllm_incomplete_ids = array("q", next_token_ids)
+                    req.dllm_algo_state = (
+                        algo_states[idx] if algo_states is not None else None
+                    )
+                    old_prefix_len = len(req.prefix_indices)
+                    new_fill_len = req.extend_range.end
+                    if new_fill_len > old_prefix_len:
+                        kv_indices_to_free = self.req_to_token_pool.req_to_token[
+                            req.req_pool_idx, old_prefix_len:new_fill_len
+                        ]
+                        self.token_to_kv_pool_allocator.free(kv_indices_to_free)
+                    continue
+
+                req.dllm_incomplete_ids = array("q")
+                req.dllm_algo_state = None
+
+                # Mirror the resolved block into the committed fill ids so the
+                # prefix cache keys on the real tokens, not the mask block, next
+                # round. Index relative to extend_range.end (the truncated/
+                # committed length), which can be shorter than
+                # full_untruncated_fill_ids when the staging adder truncates the
+                # block to the KV budget.
+                req.full_untruncated_fill_ids[
+                    req.extend_range.end - block_size : req.extend_range.end
+                ] = array("q", next_token_ids)
+
+                len_input = len(req.origin_input_ids)
+                len_fill = req.extend_range.end
+                if len_fill <= len_input:
+                    continue
+
+                if len_fill - len(next_token_ids) < len_input:
+                    next_token_ids = next_token_ids[len_input - len_fill :]
+
+                self.metrics_reporter.num_generated_tokens += len(next_token_ids)
+                req.output_ids.extend(next_token_ids)
+                req.update_finish_state(new_accepted_len=len(next_token_ids))
+
+                if req.finished():
+                    release_kv_cache(req, self.tree_cache)
+                    req.time_stats.set_completion_time()
+
+            self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
+            self.token_to_kv_pool_allocator.free_group_end()
+
+        self.metrics_reporter.report_prefill_stats(
+            batch=batch,
+            prefill_stats=batch.prefill_stats,
+            can_run_cuda_graph=result.can_run_cuda_graph,
+            dp_cooperation_info=batch.dp_cooperation_info,
+        )
+
+    def _fetch_waiting_reqs(self: Scheduler):
+        # Calculate how many requests can be added to DLLM manager
+        max_dllm_capacity = self.dllm_config.max_running_requests - len(
+            self.dllm_manager.waiting_queue
+        )
+        num_requests_to_add = min(max_dllm_capacity, len(self.waiting_queue))
+
+        if num_requests_to_add > 0:
+            requests_to_add = self.waiting_queue[:num_requests_to_add]
+            self.dllm_manager.add_waiting_reqs(requests_to_add)
+            self.waiting_queue = self.waiting_queue[num_requests_to_add:]
+
+    def _should_skip_prefill(self: Scheduler) -> bool:
+        """Check if DLLM prefill should be skipped."""
+        if (
+            self.running_batch.batch_is_full or not self.waiting_queue
+        ) and self.dllm_manager.is_empty():
+            return True
+
+        running_bs = len(self.running_batch.reqs)
+        if (
+            self.get_num_allocatable_reqs(running_bs) <= 0
+            and self.dllm_manager.is_empty()
+            and not self.enable_priority_preemption
+        ):
+            self.running_batch.batch_is_full = True
+            return True
+
+        return False
+
+    def _create_dllm_prefill_adder(self: Scheduler, running_bs: int) -> PrefillAdder:
+        """Create a prefill adder configured for DLLM scheduling."""
+        return PrefillAdder(
+            self.page_size,
+            self.tree_cache,
+            self.token_to_kv_pool_allocator,
+            self.running_batch,
+            self.new_token_ratio_tracker.current,
+            self.max_prefill_tokens,
+            self.chunked_prefill_size,
+            running_bs if self.is_mixed_chunk else 0,
+            self.priority_scheduling_preemption_threshold,
+            prefill_max_requests=self.server_args.prefill_max_requests,
+            dllm_config=self.dllm_config,
+        )
+
+    def _process_dllm_batches(self: Scheduler, adder: PrefillAdder) -> ForwardMode:
+        """Process prefill or decode batches for DLLM."""
+        forward_mode = ForwardMode.DLLM_EXTEND
+
+        # Try prefill batch first
+        prefill_reqs = self.dllm_manager.get_prefill_requests()
+        if prefill_reqs:
+            self._process_batch_by_phase(
+                adder,
+                prefill_reqs,
+                DllmReqPhase.STAGING_PREFILL,
+                DllmReqPhase.INCOMING_PREFILL,
+            )
+        else:
+            # Fall back to decode batch
+            decode_reqs = self.dllm_manager.get_decode_requests()
+            self._process_batch_by_phase(
+                adder,
+                decode_reqs,
+                DllmReqPhase.STAGING_DECODE,
+                DllmReqPhase.INCOMING_DECODE,
+            )
+
+        return forward_mode
+
+    def _process_batch_by_phase(
+        self,
+        adder: PrefillAdder,
+        batch: List[Req],
+        staging_phase: DllmReqPhase,
+        incoming_phase: DllmReqPhase,
+    ) -> None:
+        """Process a batch, separating staging and incoming requests."""
+        staging_reqs = [req for req in batch if req.dllm_phase == staging_phase]
+        if staging_reqs:
+            staging_result = self.process_dllm_staging_reqs(adder, staging_reqs)
+            if staging_result != AddReqResult.CONTINUE:
+                return
+
+        incoming_reqs = [req for req in batch if req.dllm_phase == incoming_phase]
+        if incoming_reqs:
+            self.process_dllm_incoming_reqs(adder, incoming_reqs)
+
+    def _update_state_for_batch(
+        self: Scheduler, can_run_list: List[Req], adder: PrefillAdder, running_bs: int
+    ) -> None:
+        """Update state for the batch."""
+
+        if adder.preempt_list:
+            for req in adder.preempt_list:
+                self._add_request_to_queue(req)
+
+        if can_run_list:
+            self.dllm_manager.add_staging_reqs(can_run_list)
+            self.dllm_manager.increment_inflight_middle_chunks()
+
+        self.adder = adder
+        self.can_run_list = can_run_list
+        self.running_bs = len(self.running_batch.reqs)
+
+    def _create_dllm_batch(
+        self: Scheduler, can_run_list: List[Req], forward_mode: ForwardMode
+    ) -> ScheduleBatch:
+        """Create and prepare a new DLLM batch."""
+        new_batch = ScheduleBatch.init_new(
+            can_run_list,
+            self.req_to_token_pool,
+            self.token_to_kv_pool_allocator,
+            self.tree_cache,
+            self.model_config,
+            self.enable_overlap,
+            self.spec_algorithm,
+            dllm_config=self.dllm_config,
+        )
+        new_batch.prepare_for_extend()
+        new_batch.forward_mode = forward_mode
+        new_batch.decoding_reqs = None
+
+        # Record prefill stats for logging after forward
+        from sglang.srt.managers.scheduler_components.metrics_reporter import (
+            PrefillStats,
+        )
+
+        new_batch.prefill_stats = PrefillStats.from_adder(
+            self.adder, self.running_batch.reqs, self.enable_priority_scheduling
+        )
+
+        return new_batch
+
+    def process_dllm_incoming_reqs(
+        self: Scheduler, adder: PrefillAdder, reqs: List[Req]
+    ) -> AddReqResult:
+        """Process incoming DLLM requests with resource allocation and preemption."""
+        res = AddReqResult.CONTINUE
+        for req in reqs:
+            # Check if batch is full
+            running_bs = len(self.running_batch.reqs)
+            if len(adder.can_run_list) >= self.get_num_allocatable_reqs(running_bs):
+                self.running_batch.batch_is_full = True
+
+            # Try preemption if batch is full
+            if self.running_batch.batch_is_full:
+                if (
+                    not self.enable_priority_preemption
+                    or not adder.preempt_to_schedule(req, self.server_args)
+                ):
+                    break
+
+            # Prepare and add request
+            req.init_next_round_input(self.tree_cache)
+            res = adder.add_one_req(
+                req,
+                has_chunked_req=True,
+                truncation_align_size=self.truncation_align_size,
+            )
+
+            if res != AddReqResult.CONTINUE:
+                if res == AddReqResult.NO_TOKEN:
+                    self.running_batch.batch_is_full = True
+                break
+
+        return res
+
+    def process_dllm_staging_reqs(
+        self: Scheduler, adder: PrefillAdder, reqs: List[Req]
+    ) -> AddReqResult:
+        """Process staging DLLM requests with resource allocation."""
+        for req in reqs:
+            res = adder.add_dllm_staging_req(req)
+            if res == AddReqResult.NO_TOKEN:
+                return res
+
+        return AddReqResult.CONTINUE
+
+
+class DllmManager:
+    """
+    Manager for Diffusion LLM request scheduling.
+
+    Maintains two queues:
+    - waiting_queue: The requests waiting to be scheduled with max running requests limit
+    - staging_queue: Requests allocated resources by PrefillAdder
+    """
+
+    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.waiting_queue: List[Req] = []
+        self.staging_queue: List[Req] = []
+
+    def get_prefill_requests(self) -> List[Req]:
+        """Get all prefill requests from waiting queue."""
+        return [req for req in self.waiting_queue if req.is_dllm_prefill()]
+
+    def get_decode_requests(self) -> List[Req]:
+        """Get all decode requests from waiting queue."""
+        return [req for req in self.waiting_queue if not req.is_dllm_prefill()]
+
+    def add_waiting_reqs(self, reqs: Union[Req, List[Req]]) -> None:
+        """Add requests to waiting queue with redundancy check."""
+        assert self.dllm_config is not None, "Diffusion LLM config is not set."
+
+        reqs_to_add = reqs if isinstance(reqs, list) else [reqs]
+
+        # Check for duplicate request IDs
+        if self._has_duplicate_reqs(reqs_to_add):
+            raise RuntimeError("Redundant requests detected in dLLM requests.")
+
+        self.waiting_queue.extend(reqs_to_add)
+
+    def add_staging_reqs(self, reqs: Union[Req, List[Req]]) -> None:
+        """Add requests to staging queue (allocated by PrefillAdder)."""
+        reqs_to_add = reqs if isinstance(reqs, list) else [reqs]
+        self.staging_queue.extend(reqs_to_add)
+
+    def _has_duplicate_reqs(self, reqs: List[Req]) -> bool:
+        """Check if any request ID already exists in waiting queue."""
+        existing_rids: Set[str] = {r.rid for r in self.waiting_queue}
+        return any(req.rid in existing_rids for req in reqs)
+
+    def any_staging_reqs(self) -> bool:
+        """Check if there are requests in staging queue."""
+        return self.dllm_config is not None and len(self.staging_queue) > 0
+
+    def is_empty(self) -> bool:
+        """Check if both queues are empty or DLLM is not configured."""
+        if self.dllm_config is None:
+            return True
+        return len(self.waiting_queue) == 0
+
+    def increment_inflight_middle_chunks(self) -> None:
+        """Increment chunked count for all staging requests."""
+        for req in self.staging_queue:
+            req.inflight_middle_chunks += 1
+
+    def filter_finished_reqs(self) -> None:
+        """Remove finished requests from both queues."""
+        self.waiting_queue = [req for req in self.waiting_queue if not req.finished()]
+        self.staging_queue = [req for req in self.staging_queue if not req.finished()]
+
+    def init_next_round(self) -> None:
+        """Initialize staging requests for next round and clear staging queue."""
+        for req in self.staging_queue:
+            req.init_next_round_input()
+        self.staging_queue = []