diff --git a/python/sglang/srt/dllm/mixin/req.py b/python/sglang/srt/dllm/mixin/req.py index dc6214b1d..720b9d1db 100644 --- a/python/sglang/srt/dllm/mixin/req.py +++ b/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_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.dllm_block_offset = ( + 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 + ) - 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 diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index e5404b0e3..2b852507c 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/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( diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c294893b8..12afdfd74 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index cafa87b09..6d64ffd9a 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 53d660419..9893ee243 100644 --- a/python/sglang/srt/server_args.py +++ b/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: - logger.warning( - "Radix cache is disabled because of using diffusion LLM inference" - ) - self.disable_radix_cache = True + from sglang.srt.dllm.config import DllmConfig + + config = DllmConfig.from_server_args(self) + if self.page_size % config.block_size != 0: + logger.warning( + f"Setting page size to {config.block_size} for diffusion LLM inference" + ) + 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" diff --git a/test/registered/dllm/test_llada2_mini.py b/test/registered/dllm/test_llada2_mini.py index 0f17af759..a853e0a9b 100644 --- a/test/registered/dllm/test_llada2_mini.py +++ b/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)