[feature] Initial block diffusion language model support (#12588)

Co-authored-by: Tiwei Bie <tiwei.btw@antgroup.com>

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,18 @@
+from sglang.srt.dllm.algorithm import get_algorithm
+from sglang.srt.dllm.config import DllmConfig
+from sglang.srt.server_args import ServerArgs
+
+
+class DllmAlgorithm:
+
+    def __init__(
+        self,
+        config: DllmConfig,
+    ):
+        self.block_size = config.block_size
+        self.mask_id = config.mask_id
+
+    @staticmethod
+    def from_server_args(server_args: ServerArgs):
+        config = DllmConfig.from_server_args(server_args)
+        return get_algorithm(config)

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

@@ -0,0 +1,59 @@
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.dllm.algorithm.base import DllmAlgorithm
+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 LowConfidence(DllmAlgorithm):
+
+    def run(
+        self,
+        model_runner: ModelRunner,
+        forward_batch: ForwardBatch,
+    ) -> Tuple[
+        Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor], bool
+    ]:
+        mask_index = forward_batch.input_ids == self.mask_id
+        start = len(forward_batch.input_ids) - torch.sum(mask_index).item()
+
+        for _ in range(self.block_size):
+            mask_index = forward_batch.input_ids == self.mask_id
+            if torch.sum(mask_index).item() == 0:
+                break
+
+            logits_output, can_run_cuda_graph = model_runner.forward(
+                forward_batch, pp_proxy_tensors=None
+            )
+
+            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)
+            transfer_index = torch.zeros_like(x, dtype=torch.bool, device=x.device)
+            _, select_index = torch.topk(confidence, k=1)
+            transfer_index[select_index] = True
+
+            forward_batch.input_ids[transfer_index] = x[transfer_index]
+
+        logits_output, can_run_cuda_graph = model_runner.forward(
+            forward_batch, pp_proxy_tensors=None
+        )
+
+        next_token_ids = forward_batch.input_ids[start:]
+        return logits_output, next_token_ids, can_run_cuda_graph
+
+
+Algorithm = LowConfidence

python/sglang/srt/dllm/config.py

@@ -0,0 +1,40 @@
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.server_args import ServerArgs
+
+
+class DllmConfig:
+    def __init__(
+        self,
+        mask_id: int,
+        block_size: int,
+        algorithm: str,
+    ):
+        self.algorithm = algorithm
+        self.block_size = block_size
+        self.mask_id = mask_id
+
+    @staticmethod
+    def from_server_args(
+        server_args: ServerArgs,
+    ):
+        if server_args.dllm_algorithm is None:
+            return None
+
+        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":
+            mask_id = 156895
+        else:
+            raise RuntimeError(
+                f"Unknown diffusion LLM: {config.hf_config.architectures[0]}"
+            )
+
+        return DllmConfig(
+            algorithm=server_args.dllm_algorithm,
+            block_size=server_args.dllm_block_size,
+            mask_id=mask_id,
+        )