[DLLM] Add initial radix cache support (#18724)

python/sglang/srt/dllm/mixin/req.py

@@ -19,7 +19,6 @@ class DllmReqPhase(str, enum.Enum):
 class ReqDllmMixin:
     def init_diffusion_llm(self: Req, dllm_config: DllmConfig):
         self.dllm_phase: Optional[DllmReqPhase] = None
-        self.dllm_ids = []
         self.dllm_block_offset = 0
         self.dllm_config = dllm_config
 
@@ -55,13 +54,21 @@ class ReqDllmMixin:
             self.dllm_phase = DllmReqPhase.STAGING_DECODE
 
     def _init_fill_ids_for_dllm(self: Req):
-        if not self.dllm_ids:
+        self.dllm_block_offset = (
-            self.dllm_ids = (
+            0
+            if not self.fill_ids
+            else self.dllm_block_offset + self.dllm_config.block_size
+        )
+        self.fill_ids = (
             self.origin_input_ids
+            + self.output_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
+    def _update_block_offset_for_dllm(self):
+        prefix_len = len(self.prefix_indices)
+        assert (
+            prefix_len % self.dllm_config.block_size == 0
+        ), f"Unexpected prefix len: {prefix_len}"
+        if prefix_len > self.dllm_block_offset:
+            self.dllm_block_offset = prefix_len

python/sglang/srt/dllm/mixin/scheduler.py

@@ -7,13 +7,14 @@ from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.dllm.mixin.req import DllmReqPhase
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder
+from sglang.srt.mem_cache.common import release_kv_cache
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.observability.req_time_stats import set_time_batch
 
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from sglang.srt.managers.scheduler import Scheduler
+    from sglang.srt.managers.scheduler import GenerationBatchResult, Scheduler
 
 
 class SchedulerDllmMixin:
@@ -59,6 +60,46 @@ class SchedulerDllmMixin:
         new_batch = self._create_dllm_batch(can_run_list, forward_mode)
         return new_batch
 
+    def process_batch_result_dllm(
+        self: Scheduler,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+    ):
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+
+        if result.next_token_ids:
+            self.token_to_kv_pool_allocator.free_group_begin()
+
+            for idx in range(batch.batch_size()):
+                req = batch.reqs[idx]
+
+                next_token_ids = result.next_token_ids[idx].tolist()
+                new_tokens = len(next_token_ids)
+                if new_tokens == 0:
+                    continue
+
+                req.fill_ids[-new_tokens:] = next_token_ids[:]
+                self.num_generated_tokens += new_tokens
+
+                req.output_ids.extend(next_token_ids)
+                req.check_finished(new_accepted_len=new_tokens)
+
+                if req.finished():
+                    release_kv_cache(req, self.tree_cache)
+                    req.time_stats.set_completion_time()
+
+            self.stream_output(batch.reqs, batch.return_logprob)
+            self.token_to_kv_pool_allocator.free_group_end()
+
+        if self.current_scheduler_metrics_enabled:
+            can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False)
+            self.log_prefill_stats(
+                prefill_stats=batch.prefill_stats,
+                can_run_cuda_graph=can_run_cuda_graph,
+                dp_cooperation_info=batch.dp_cooperation_info,
+            )
+
     def _fetch_waiting_reqs(self: Scheduler):
         # Calculate how many requests can be added to DLLM manager
         max_dllm_capacity = self.dllm_config.max_running_requests - len(

python/sglang/srt/managers/schedule_batch.py

@@ -894,6 +894,9 @@ class Req(ReqDllmMixin):
             )
             self.cache_protected_len = len(self.prefix_indices)
 
+            if self.is_dllm():
+                self._update_block_offset_for_dllm()
+
         if (
             self.is_retracted
             and self.multimodal_inputs is not None

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -354,46 +354,6 @@ class SchedulerOutputProcessorMixin:
             batch.reqs, batch.return_logprob, is_idle_batch=True
         )
 
-    def process_batch_result_dllm(
-        self: Scheduler,
-        batch: ScheduleBatch,
-        result: GenerationBatchResult,
-    ):
-        if result.copy_done is not None:
-            result.copy_done.synchronize()
-
-        self.token_to_kv_pool_allocator.free_group_begin()
-
-        for idx in range(batch.batch_size()):
-            # If no new tokens generated, meaning the prefilling stage
-            if not result.next_token_ids:
-                break
-
-            req = batch.reqs[idx]
-            next_token_ids = result.next_token_ids[idx].tolist()
-            self.num_generated_tokens += len(next_token_ids)
-
-            for _token_idx, next_token_id in enumerate(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.set_completion_time()
-                    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()
-
-        if self.current_scheduler_metrics_enabled:
-            can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False)
-            self.log_prefill_stats(
-                prefill_stats=batch.prefill_stats,
-                can_run_cuda_graph=can_run_cuda_graph,
-                dp_cooperation_info=batch.dp_cooperation_info,
-            )
-
     def process_batch_result_decode(
         self: Scheduler,
         batch: ScheduleBatch,

python/sglang/srt/server_args.py

@@ -3063,11 +3063,27 @@ class ServerArgs:
                 "Overlap schedule is disabled because of using diffusion LLM inference"
             )
             self.disable_overlap_schedule = True
+
         if not self.disable_radix_cache:
+            from sglang.srt.dllm.config import DllmConfig
+
+            config = DllmConfig.from_server_args(self)
+            if self.page_size % config.block_size != 0:
                 logger.warning(
-                "Radix cache is disabled because of using diffusion LLM inference"
+                    f"Setting page size to {config.block_size} for diffusion LLM inference"
                 )
-            self.disable_radix_cache = True
+                self.page_size = config.block_size
+            if self.enable_hierarchical_cache:
+                logger.warning(
+                    "Hierarchical cache is disabled because of using diffusion LLM inference"
+                )
+                self.enable_hierarchical_cache = False
+            if self.enable_lmcache:
+                logger.warning(
+                    "LMCache is disabled because of using diffusion LLM inference"
+                )
+                self.enable_lmcache = False
+
         if not self.pp_size > 1:
             logger.warning(
                 "Pipeline parallelism is disabled because of using diffusion LLM inference"

test/registered/dllm/test_llada2_mini.py

@@ -73,7 +73,7 @@ class TestLLaDA2Mini(CustomTestCase):
         if is_in_amd_ci():
             self.assertGreater(metrics["output_throughput"], 80)
         else:
-            self.assertGreater(metrics["output_throughput"], 250)
+            self.assertGreater(metrics["output_throughput"], 350)
 
     def test_bs_1_speed(self):
         args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)