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