[DLLM] Add JointThreshold algorithm for joint M2T and T2T decoding (#18171)

Signed-off-by: Junlin Zhou <zhoujunlin.zjl@antgroup.com>
Co-authored-by: Tiwei Bie <tiwei.btw@antgroup.com>

docs/supported_models/text_generation/diffusion_language_models.md

@@ -6,6 +6,8 @@ Diffusion language models have shown promise for non-autoregressive text generat
 
 ## Example Launch Command
 
+SGLang supports different DLLM algorithms such as `LowConfidence` and `JointThreshold`.
+
 ```shell
 python3 -m sglang.launch_server \
   --model-path inclusionAI/LLaDA2.0-mini \ # example HF/local path
@@ -17,6 +19,10 @@ python3 -m sglang.launch_server \
 
 ## Example Configuration File
 
+Depending on the algorithm selected, the configuration parameters vary.
+
+LowConfidence Config:
+
 ```yaml
 # Confidence threshold for accepting predicted tokens
 # - Higher values: More conservative, better quality but slower
@@ -28,6 +34,25 @@ threshold: 0.95
 block_size: 32
 ```
 
+JointThreshold Config:
+
+```yaml
+# Decoding threshold for Mask-to-Token (M2T) phase
+# - Higher values: More conservative, better quality but slower
+# - Lower values: More aggressive, faster but potentially lower quality
+# Range: 0.0 - 1.0
+threshold: 0.5
+# Decoding threshold for Token-to-Token (T2T) phase
+# Range: 0.0 - 1.0
+# Setting to 0.0 allows full editing (recommended for most cases).
+edit_threshold: 0.0
+# Max extra T2T steps after all masks are removed. Prevents infinite loops.
+max_post_edit_steps: 16
+# 2-gram repetition penalty (default 0).
+# An empirical value of 3 is often sufficient to mitigate most repetitions.
+penalty_lambda: 0
+```
+
 ## Example Client Code Snippet
 
 Just like other supported models, diffusion language models can be used via the REST API or Python client.

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

@@ -0,0 +1,139 @@
+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.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+class JointThreshold(DllmAlgorithm):
+
+    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 run(
+        self,
+        model_runner: ModelRunner,
+        forward_batch: ForwardBatch,
+    ) -> tuple[LogitsProcessorOutput | torch.Tensor, torch.Tensor | None, bool]:
+        batch_size = forward_batch.batch_size
+        device = forward_batch.input_ids.device
+
+        mask_index = forward_batch.input_ids == self.mask_id
+        if not mask_index.any():
+            out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
+            return out.logits_output, [], out.can_run_graph
+
+        start_list = []
+        prompt_masks = []
+        for i in range(batch_size):
+            block_start = i * self.block_size
+            block_end = block_start + self.block_size
+            block_input_ids = forward_batch.input_ids[block_start:block_end]
+
+            prompt_mask = block_input_ids != self.mask_id
+            prompt_masks.append(prompt_mask)
+            start_list.append(prompt_mask.sum().item())
+
+        post_edit_steps = torch.zeros(batch_size, dtype=torch.int32, device=device)
+
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        # Controls whether to perform an additional forward pass for KV cache persistence.
+        # For certain decoding rounds where the terminal step yields no state change,
+        # this can be set to False to bypass the overhead of an idle forward pass.
+        any_changed_in_last_step = False
+
+        max_iterations = self.block_size + self.max_post_edit_steps
+        for _ in range(max_iterations):
+            if finished.all():
+                break
+
+            out = model_runner.forward(forward_batch, pp_proxy_tensors=None)
+            logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph
+
+            any_changed_in_last_step = False
+
+            for i in range(batch_size):
+                if finished[i]:
+                    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 = logits_output.full_logits[block_start:block_end]
+                curr_prompt_mask = prompt_masks[i]
+
+                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)
+                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:
+                    post_edit_steps[i] += 1
+                    if post_edit_steps[i] > self.max_post_edit_steps:
+                        finished[i] = True
+                        continue
+
+                # 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 not transfer_index.any():
+                    finished[i] = True
+                    continue
+
+                curr_input_ids[transfer_index] = x[transfer_index]
+                any_changed_in_last_step = True
+
+        if any_changed_in_last_step:
+            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 = 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)
+        ]
+
+        return logits_output, next_token_ids_list, can_run_cuda_graph
+
+
+Algorithm = JointThreshold