[DLLM] feat: Add threshold based parallel decoding support (#14412)

Co-authored-by: Jinwei Yao <jinweiy@illinois.edu>
Co-authored-by: 赵晨阳 <zhaochen20@outlook.com>

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

@@ -5,6 +5,7 @@ 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
@@ -12,6 +13,13 @@ from sglang.srt.model_executor.model_runner import ModelRunner
 
 class LowConfidence(DllmAlgorithm):
 
+    def __init__(
+        self,
+        config: DllmConfig,
+    ):
+        super().__init__(config)
+        self.threshold = config.algorithm_config.get("threshold", 0.95)
+
     def run(
         self,
         model_runner: ModelRunner,
@@ -42,9 +50,11 @@ class LowConfidence(DllmAlgorithm):
             )
             x = torch.where(mask_index, x, forward_batch.input_ids)
             confidence = torch.where(mask_index, p, -np.inf)
-            transfer_index = torch.zeros_like(x, dtype=torch.bool, device=x.device)
+
-            _, select_index = torch.topk(confidence, k=1)
+            transfer_index = confidence > self.threshold
-            transfer_index[select_index] = True
+            if transfer_index.sum().item() == 0:
+                _, select_index = torch.topk(confidence, k=1)
+                transfer_index[select_index] = True
 
             forward_batch.input_ids[transfer_index] = x[transfer_index]
 

python/sglang/srt/dllm/config.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.server_args import ServerArgs
 
@@ -5,11 +7,13 @@ from sglang.srt.server_args import ServerArgs
 class DllmConfig:
     def __init__(
         self,
-        mask_id: int,
-        block_size: int,
         algorithm: str,
+        algorithm_config: dict[str, Any],
+        block_size: int,
+        mask_id: int,
     ):
         self.algorithm = algorithm
+        self.algorithm_config = algorithm_config
         self.block_size = block_size
         self.mask_id = mask_id
 
@@ -20,21 +24,38 @@ class DllmConfig:
         if server_args.dllm_algorithm is None:
             return None
 
-        config = ModelConfig.from_server_args(
+        model_config = ModelConfig.from_server_args(
             server_args,
             model_path=server_args.model_path,
             model_revision=server_args.revision,
         )
 
-        if config.hf_config.architectures[0] == "LLaDA2MoeModelLM":
+        if model_config.hf_config.architectures[0] == "LLaDA2MoeModelLM":
+            block_size = 32
             mask_id = 156895
         else:
             raise RuntimeError(
-                f"Unknown diffusion LLM: {config.hf_config.architectures[0]}"
+                f"Unknown diffusion LLM: {model_config.hf_config.architectures[0]}"
             )
 
+        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,
-            block_size=server_args.dllm_block_size,
+            algorithm_config=algorithm_config,
+            block_size=block_size,
             mask_id=mask_id,
         )

python/sglang/srt/server_args.py

@@ -461,7 +461,7 @@ class ServerArgs:
 
     # Diffusion LLM
     dllm_algorithm: Optional[str] = None
-    dllm_block_size: Optional[int] = None
+    dllm_algorithm_config: Optional[str] = None
 
     # Double Sparsity
     enable_double_sparsity: bool = False
@@ -3352,13 +3352,13 @@ class ServerArgs:
             "--dllm-algorithm",
             type=str,
             default=ServerArgs.dllm_algorithm,
-            help="The diffusion LLM algorithm.",
+            help="The diffusion LLM algorithm, such as LowConfidence.",
         )
         parser.add_argument(
-            "--dllm-block-size",
+            "--dllm-algorithm-config",
-            type=int,
+            type=str,
-            default=ServerArgs.dllm_block_size,
+            default=ServerArgs.dllm_algorithm_config,
-            help="The number of tokens processed in each iteration of the block diffusion LLM.",
+            help="The diffusion LLM algorithm configurations. Must be a YAML file.",
         )
 
         # Double Sparsity