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

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

python/sglang/srt/managers/schedule_batch.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import enum
 
+from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 # Copyright 2023-2024 SGLang Team
@@ -442,6 +443,7 @@ class Req:
         sampling_params: SamplingParams,
         return_logprob: bool = False,
         top_logprobs_num: int = 0,
+        dllm_config: Optional[DllmConfig] = None,
         token_ids_logprob: List[int] = None,
         stream: bool = False,
         origin_input_ids_unpadded: Optional[Tuple[int]] = None,
@@ -683,6 +685,11 @@ class Req:
         # For Matryoshka embeddings
         self.dimensions = dimensions
 
+        # For diffusion LLM
+        self.dllm_ids = []
+        self.dllm_block_offset = 0
+        self.dllm_config = dllm_config
+
     @property
     def seqlen(self):
         return len(self.origin_input_ids) + len(self.output_ids)
@@ -751,8 +758,28 @@ class Req:
         # Whether request reached finished condition
         return self.finished_reason is not None
 
+    def is_dllm(self):
+        return self.dllm_config is not None
+
     def init_next_round_input(self, tree_cache: Optional[BasePrefixCache] = None):
-        self.fill_ids = self.origin_input_ids + self.output_ids
+        if self.is_dllm():
+            if not self.fill_ids:
+                self.dllm_ids = (
+                    self.origin_input_ids
+                    + [
+                        self.dllm_config.mask_id,
+                    ]
+                    * self.dllm_config.block_size
+                )
+            else:
+                self.dllm_block_offset += self.dllm_config.block_size
+                self.dllm_ids += [
+                    self.dllm_config.mask_id
+                ] * self.dllm_config.block_size
+            self.fill_ids = self.dllm_ids
+        else:
+            self.fill_ids = self.origin_input_ids + self.output_ids
+
         input_len = len(self.fill_ids)
         # NOTE: the matched length is at most 1 less than the input length to enable logprob computation
         max_prefix_len = input_len - 1
@@ -1127,6 +1154,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     # hicache pointer for synchronizing data loading from CPU to GPU
     hicache_consumer_index: int = -1
 
+    # Diffusion LLM
+    dllm_config: Optional[DllmConfig] = None
+
     @classmethod
     def init_new(
         cls,
@@ -1138,6 +1168,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         enable_overlap: bool,
         spec_algorithm: SpeculativeAlgorithm,
         chunked_req: Optional[Req] = None,
+        dllm_config: Optional[DllmConfig] = None,
     ):
         return_logprob = any(req.return_logprob for req in reqs)
 
@@ -1166,6 +1197,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             return_hidden_states=any(req.return_hidden_states for req in reqs),
             is_prefill_only=all(req.is_prefill_only for req in reqs),
             chunked_req=chunked_req,
+            dllm_config=dllm_config,
         )
 
     def batch_size(self):
@@ -1174,6 +1206,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     def is_empty(self):
         return len(self.reqs) == 0
 
+    def is_dllm(self):
+        return self.dllm_config is not None
+
     def prepare_encoder_info_extend(self, input_ids: List[int], seq_lens: List[int]):
         self.encoder_lens_cpu = []
         self.encoder_cached = []
@@ -1886,6 +1921,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             extend_input_logprob_token_ids=self.extend_input_logprob_token_ids,
             is_prefill_only=self.is_prefill_only,
             dimensions=self.dimensions,
+            dllm_block_offsets=[req.dllm_block_offset for req in self.reqs],
+            dllm_config=self.dllm_config,
         )
 
     def copy(self):
@@ -1999,3 +2036,7 @@ class ModelWorkerBatch:
 
     # Whether this batch is prefill-only (no token generation needed)
     is_prefill_only: bool = False
+
+    # Diffusion LLM
+    dllm_block_offsets: Optional[List[int]] = None
+    dllm_config: Optional[DllmConfig] = None

python/sglang/srt/managers/scheduler.py

@@ -60,6 +60,7 @@ from sglang.srt.disaggregation.utils import (
     prepare_abort,
 )
 from sglang.srt.distributed import get_pp_group, get_world_group
+from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
@@ -287,6 +288,9 @@ class Scheduler(
         # Init model config
         self.model_config = ModelConfig.from_server_args(server_args)
 
+        # Init diffusion LLM config
+        self.dllm_config = DllmConfig.from_server_args(server_args)
+
         # Init inter-process communication
         self.init_sockets(server_args, port_args)
 
@@ -449,6 +453,10 @@ class Scheduler(
 
         # Init chunked prefill
         self.chunked_prefill_size = server_args.chunked_prefill_size
+        if self.dllm_config is not None:
+            # We currently leverage chunked prefill to implement block diffusion
+            # for diffusion LLM.
+            self.chunked_prefill_size = self.dllm_config.block_size
         if self.chunked_prefill_size <= 0:  # -1 means disable
             self.chunked_prefill_size = None
         self.chunked_req = None
@@ -1284,6 +1292,7 @@ class Scheduler(
                     self.metrics_collector if self.enable_metrics else None
                 ),
                 http_worker_ipc=recv_req.http_worker_ipc,
+                dllm_config=self.dllm_config,
             )
             req.tokenizer = self.tokenizer
 
@@ -1600,6 +1609,10 @@ class Scheduler(
             self.handle_embedding_request(tokenized_req)
 
     def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
+        if self.dllm_config is not None:
+            if self.chunked_req is not None and self.chunked_req.finished():
+                self.chunked_req = None
+
         # Merge the prefill batch into the running batch
         chunked_req_to_exclude = set()
         if self.chunked_req:
@@ -1832,6 +1845,7 @@ class Scheduler(
             self.enable_overlap,
             self.spec_algorithm,
             chunked_req=self.chunked_req,
+            dllm_config=self.dllm_config,
         )
         if self.enable_hierarchical_cache:
             # todo (zhiqiang): disable cuda graph execution if hicache loading triggered
@@ -2064,7 +2078,10 @@ class Scheduler(
             self.process_batch_result_decode(batch, result)
             trace_slice_batch(RequestStage.DECODE_LOOP, batch.reqs)
         elif batch.forward_mode.is_extend():
-            self.process_batch_result_prefill(batch, result)
+            if batch.is_dllm():
+                self.process_batch_result_dllm(batch, result)
+            else:
+                self.process_batch_result_prefill(batch, result)
         elif batch.forward_mode.is_prebuilt():
             self.process_batch_result_prebuilt(batch)
         elif batch.forward_mode.is_idle():

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -281,6 +281,36 @@ class SchedulerOutputProcessorMixin:
 
         return predict_tokens
 
+    def process_batch_result_dllm(
+        self: Scheduler,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+    ):
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+
+        next_token_ids = result.next_token_ids.tolist()
+        self.num_generated_tokens += len(next_token_ids)
+
+        self.token_to_kv_pool_allocator.free_group_begin()
+
+        assert len(batch.reqs) == 1, "batch size is currently expected to be 1"
+        req = batch.reqs[0]
+
+        for next_token_id in next_token_ids:
+            req.output_ids.append(next_token_id)
+            req.check_finished()
+
+            if req.finished():
+                release_kv_cache(req, self.tree_cache)
+                req.time_stats.completion_time = time.perf_counter()
+                break
+
+            self.tree_cache.cache_unfinished_req(req)
+
+        self.stream_output(batch.reqs, batch.return_logprob)
+        self.token_to_kv_pool_allocator.free_group_end()
+
     def process_batch_result_decode(
         self: Scheduler,
         batch: ScheduleBatch,

python/sglang/srt/managers/tp_worker.py

@@ -22,6 +22,7 @@ import torch
 
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.distributed import get_pp_group, get_world_group
+from sglang.srt.dllm.algorithm.base import DllmAlgorithm
 from sglang.srt.managers.io_struct import (
     DestroyWeightsUpdateGroupReqInput,
     GetWeightsByNameReqInput,
@@ -234,6 +235,9 @@ class TpModelWorker(BaseTpWorker):
             is_draft_model=is_draft_worker,
         )
 
+        if server_args.dllm_algorithm is not None:
+            self.dllm_algorithm = DllmAlgorithm.from_server_args(server_args)
+
         self._model_runner = ModelRunner(
             model_config=self.model_config,
             mem_fraction_static=server_args.mem_fraction_static,
@@ -340,6 +344,9 @@ class TpModelWorker(BaseTpWorker):
             self.model_runner.token_to_kv_pool.size,
         )
 
+    def is_dllm(self):
+        return hasattr(self, "dllm_algorithm")
+
     def forward_batch_generation(
         self,
         model_worker_batch: ModelWorkerBatch,
@@ -368,6 +375,16 @@ class TpModelWorker(BaseTpWorker):
             )
 
         if self.pp_group.is_last_rank:
+            if self.is_dllm():
+                logits_output, next_token_ids, can_run_cuda_graph = (
+                    self.dllm_algorithm.run(self.model_runner, forward_batch)
+                )
+                return GenerationBatchResult(
+                    logits_output=logits_output,
+                    next_token_ids=next_token_ids,
+                    can_run_cuda_graph=can_run_cuda_graph,
+                )
+
             logits_output, can_run_cuda_graph = self.model_runner.forward(
                 forward_batch,
                 pp_proxy_tensors=pp_proxy_tensors,