diff --git a/python/sglang/srt/dllm/algorithm/low_confidence.py b/python/sglang/srt/dllm/algorithm/low_confidence.py index c8912b47c..6acdcc7f8 100644 --- a/python/sglang/srt/dllm/algorithm/low_confidence.py +++ b/python/sglang/srt/dllm/algorithm/low_confidence.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union +from typing import List, Tuple, Union import numpy as np import torch @@ -24,11 +24,29 @@ class LowConfidence(DllmAlgorithm): self, model_runner: ModelRunner, forward_batch: ForwardBatch, - ) -> Tuple[ - Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor], bool - ]: + ) -> Tuple[Union[LogitsProcessorOutput, torch.Tensor], List[torch.Tensor], bool]: + batch_size = forward_batch.batch_size + # Here, the forward_batch full logits contains all the blocks + # such as [dllm_block_size * batch_size, hidden_size] + start_list = [] mask_index = forward_batch.input_ids == self.mask_id - start = len(forward_batch.input_ids) - torch.sum(mask_index).item() + + # Fast path: if there is no mask token, forward and save kv cache + if torch.sum(mask_index).item() == 0: + out = model_runner.forward(forward_batch, pp_proxy_tensors=None) + logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph + + next_token_ids = [] + return logits_output, next_token_ids, can_run_cuda_graph + + # Calculate start positions for each block + for block_id in range(batch_size): + block_start = block_id * self.block_size + block_end = block_start + self.block_size + block_input_ids = forward_batch.input_ids[block_start:block_end] + block_mask_index = block_input_ids == self.mask_id + start = self.block_size - torch.sum(block_mask_index).item() + start_list.append(start) for _ in range(self.block_size): mask_index = forward_batch.input_ids == self.mask_id @@ -37,31 +55,50 @@ class LowConfidence(DllmAlgorithm): out = model_runner.forward(forward_batch, pp_proxy_tensors=None) logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph + assert batch_size == forward_batch.input_ids.shape[0] // self.block_size + for batch_id in range(batch_size): + curr_block_start = batch_id * self.block_size + curr_block_end = curr_block_start + self.block_size + block_input_ids = forward_batch.input_ids[ + curr_block_start:curr_block_end, + ] + block_mask_index = block_input_ids == self.mask_id + if torch.sum(block_mask_index).item() == 0: + continue + curr_logits = logits_output.full_logits[ + curr_block_start:curr_block_end, + ] - 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) + x = torch.argmax(curr_logits, dim=-1) + p = torch.squeeze( + torch.gather( + F.softmax(curr_logits, dim=-1), + dim=-1, + index=torch.unsqueeze(x, -1), + ), + -1, + ) + x = torch.where(block_mask_index, x, block_input_ids) + confidence = torch.where(block_mask_index, p, -np.inf) - transfer_index = confidence > self.threshold - if transfer_index.sum().item() == 0: - _, select_index = torch.topk(confidence, k=1) - transfer_index[select_index] = True + transfer_index = confidence > self.threshold - forward_batch.input_ids[transfer_index] = x[transfer_index] + if transfer_index.sum().item() == 0: + _, select_index = torch.topk(confidence, k=1) + transfer_index[select_index] = True + + block_input_ids[transfer_index] = x[transfer_index] out = model_runner.forward(forward_batch, pp_proxy_tensors=None) logits_output, can_run_cuda_graph = out.logits_output, out.can_run_graph + # Here next token ids is tricky to implement the dynamic lengths, + # so we return a list of tensors + next_token_ids = torch.reshape(forward_batch.input_ids, (batch_size, -1)) + next_token_ids_list = [ + next_token_ids[i, start_list[i] :] for i in range(batch_size) + ] - next_token_ids = forward_batch.input_ids[start:] - return logits_output, next_token_ids, can_run_cuda_graph + return logits_output, next_token_ids_list, can_run_cuda_graph Algorithm = LowConfidence diff --git a/python/sglang/srt/dllm/config.py b/python/sglang/srt/dllm/config.py index 3809db1bf..ce601cbdc 100644 --- a/python/sglang/srt/dllm/config.py +++ b/python/sglang/srt/dllm/config.py @@ -11,11 +11,13 @@ class DllmConfig: algorithm_config: dict[str, Any], block_size: int, mask_id: int, + max_running_requests: int, ): self.algorithm = algorithm self.algorithm_config = algorithm_config self.block_size = block_size self.mask_id = mask_id + self.max_running_requests = max_running_requests @staticmethod def from_server_args( @@ -38,6 +40,12 @@ class DllmConfig: f"Unknown diffusion LLM: {model_config.hf_config.architectures[0]}" ) + max_running_requests = ( + 1 + if server_args.max_running_requests is None + else server_args.max_running_requests + ) + algorithm_config = {} if server_args.dllm_algorithm_config is not None: try: @@ -58,4 +66,5 @@ class DllmConfig: algorithm_config=algorithm_config, block_size=block_size, mask_id=mask_id, + max_running_requests=max_running_requests, ) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5fae27b64..068bd1849 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1159,6 +1159,62 @@ class Req: ) +class DllmStagingReqs: + def __init__(self, dllm_config: Optional[DllmConfig] = None): + self.dllm_config = dllm_config + self.max_running_reqs = ( + dllm_config.max_running_requests if dllm_config is not None else 1 + ) + self.reqs: List[Req] = [] + + def add_reqs(self, req: Union[Req, List[Req], "DllmStagingReqs"]): + assert self.dllm_config is not None, "Diffusion LLM config is not set." + + if isinstance(req, DllmStagingReqs): + reqs_to_add = req.reqs + elif isinstance(req, list): + reqs_to_add = req + else: + reqs_to_add = [req] + + num_to_add = len(reqs_to_add) + + # Sanity check: + if self.check_redundant_reqs(reqs_to_add): + raise RuntimeError("Redundant requests detected in dLLM requests.") + + if len(self.reqs) + num_to_add > self.max_running_reqs: + raise RuntimeError( + f"Exceeding maximum number of concurrent diffusion LLM requests: {self.max_running_reqs}" + ) + + self.reqs.extend(reqs_to_add) + + def check_redundant_reqs(self, reqs: List[Req]) -> bool: + existing_rids: Set[str] = {r.rid for r in self.reqs} + return any(req.rid in existing_rids for req in reqs) + + def init_next_round(self): + for req in self.reqs: + req.init_next_round_input() + + def non_empty(self) -> bool: + return self.dllm_config is not None and len(self.reqs) > 0 + + def empty(self) -> bool: + return self.dllm_config is None or len(self.reqs) == 0 + + def update_chunked_status(self): + for req in self.reqs: + req.is_chunked += 1 + + def filter_finished_reqs(self): + self.reqs = [req for req in self.reqs if not req.finished()] + + def __iter__(self): + return iter(self.reqs) + + @dataclasses.dataclass class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): """Store all information of a batch on the scheduler.""" @@ -1279,6 +1335,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): hicache_consumer_index: int = -1 # Diffusion LLM + dllm_staging_reqs: Optional[DllmStagingReqs] = None dllm_config: Optional[DllmConfig] = None # Metrics @@ -1295,6 +1352,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): enable_overlap: bool, spec_algorithm: SpeculativeAlgorithm, chunked_req: Optional[Req] = None, + dllm_staging_reqs: Optional[DllmStagingReqs] = None, dllm_config: Optional[DllmConfig] = None, ): return_logprob = any(req.return_logprob for req in reqs) @@ -1320,6 +1378,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): return_routed_experts=any(req.return_routed_experts for req in reqs), is_prefill_only=all(req.is_prefill_only for req in reqs), chunked_req=chunked_req, + dllm_staging_reqs=dllm_staging_reqs, dllm_config=dllm_config, ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index fa989650e..49d71a27a 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -32,8 +32,9 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union import torch +from sglang.srt.dllm.config import DllmConfig from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split -from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.managers.schedule_batch import DllmStagingReqs, Req, ScheduleBatch from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator @@ -373,6 +374,7 @@ class PrefillAdder: priority_scheduling_preemption_threshold: int = 0, prefill_max_requests: Optional[int] = None, prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None, + dllm_config: Optional[DllmConfig] = None, ): self.page_size = page_size self.tree_cache = tree_cache @@ -381,6 +383,11 @@ class PrefillAdder: self.new_token_ratio = new_token_ratio self.rem_input_tokens = rem_input_tokens - mixed_with_decode_tokens self.rem_chunk_tokens = rem_chunk_tokens + self.dllm_config = dllm_config + + if self.dllm_config is not None: + self._init_dllm_meta(dllm_config) + if self.rem_chunk_tokens is not None: self.rem_chunk_tokens -= mixed_with_decode_tokens self.rem_total_token_offset = mixed_with_decode_tokens @@ -414,6 +421,13 @@ class PrefillAdder: self.prefill_max_requests = prefill_max_requests self.prefill_delayer_single_pass = prefill_delayer_single_pass + def _init_dllm_meta(self, dllm_config: DllmConfig): + self.dllm_block_size = dllm_config.block_size + max_running_reqs = dllm_config.max_running_requests + + self.rem_dllm_tokens = max_running_reqs * self.dllm_block_size + self.dllm_staging_reqs = DllmStagingReqs(dllm_config=dllm_config) + def _get_running_request_total_token_offset(self, req: Req) -> int: return ( min( @@ -473,11 +487,16 @@ class PrefillAdder: if self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0: return AddReqResult.NO_TOKEN - if self.rem_input_tokens <= 0 or ( - self.rem_chunk_tokens is not None and self.rem_chunk_tokens <= 0 - ): + if self.rem_input_tokens <= 0: return AddReqResult.OTHER + if self.dllm_config is not None: + if self.rem_dllm_tokens <= 0: + return AddReqResult.OTHER + else: + if self.rem_chunk_tokens is not None and self.rem_chunk_tokens <= 0: + return AddReqResult.OTHER + return AddReqResult.CONTINUE def _update_prefill_budget( @@ -489,18 +508,61 @@ class PrefillAdder: self.rem_total_token_offset += extend_input_len + max_new_tokens self.cur_rem_token_offset += extend_input_len self.rem_input_tokens -= extend_input_len - if self.rem_chunk_tokens is not None: + + if self.dllm_config is not None: + self.rem_dllm_tokens -= extend_input_len + elif self.rem_chunk_tokens is not None: self.rem_chunk_tokens -= extend_input_len self.log_hit_tokens += prefix_len self.log_input_tokens += extend_input_len - def add_chunked_req(self, req: Req): - _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - # The chunked_req must be added to the list; otherwise, it will cause a memory leak. - # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens. + def _get_dllm_remain_tokens(self) -> int: + _rem_tokens = min( + self.rem_dllm_tokens, + self.dllm_block_size, + int(self.rem_total_tokens), + ) if _rem_tokens <= 0: - _rem_tokens = self.rem_chunk_tokens + _rem_tokens = self.rem_dllm_tokens + + return _rem_tokens + + def _add_dllm_req(self, req: Req, prefix_len: int): + # FIXME: consider the case when rem_dllm_tokens < dllm_block_size, + # the diffusion unmask process may have some problems + # Make sure at least one page is available + trunc_len = ( + min(self.rem_dllm_tokens, self.dllm_block_size) + // self.page_size + * self.page_size + ) + + req.extend_input_len = trunc_len + req.fill_ids = req.fill_ids[: prefix_len + trunc_len] + + self.can_run_list.append(req) + self.dllm_staging_reqs.add_reqs(req) + + self._update_prefill_budget(prefix_len, trunc_len, 0) + + def _req_inc_lock_ref(self, req: Req): + if self.is_hybrid_swa: + swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node) + req.swa_uuid_for_lock = swa_uuid_for_lock + else: + self.tree_cache.inc_lock_ref(req.last_node) + + def add_chunked_req(self, req: Req): + if self.dllm_config is not None: + _rem_tokens = self._get_dllm_remain_tokens() + else: + _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) + # The chunked_req must be added to the list; otherwise, it will cause a memory leak. + # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens. + if _rem_tokens <= 0: + _rem_tokens = self.rem_chunk_tokens + truncated = req.extend_input_len > _rem_tokens req.set_extend_input_len(min(req.extend_input_len, _rem_tokens)) req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len] @@ -588,7 +650,12 @@ class PrefillAdder: return AddReqResult.NO_TOKEN tokens_freed += tokens_occupied - if ( + if self.dllm_config is not None: + if self.rem_dllm_tokens <= 0: + return AddReqResult.OTHER + + self._add_dllm_req(req, 0) + elif ( self.rem_chunk_tokens is None # chunked prefill is disabled or req.extend_input_len <= self.rem_chunk_tokens # it is the last chunk ): @@ -671,14 +738,21 @@ class PrefillAdder: ): return AddReqResult.OTHER - if self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens: + if self.dllm_config is not None: + if self.rem_dllm_tokens <= 0: + return AddReqResult.OTHER + + assert ( + truncation_align_size is None + ), "truncation_align_size is not supported for dllm prefill" + + self._add_dllm_req(req, prefix_len) + self._req_inc_lock_ref(req) + elif self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens: # Non-chunked prefill self.can_run_list.append(req) - if self.is_hybrid_swa: - swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node) - req.swa_uuid_for_lock = swa_uuid_for_lock - else: - self.tree_cache.inc_lock_ref(req.last_node) + + self._req_inc_lock_ref(req) self._update_prefill_budget( prefix_len, input_tokens, @@ -690,6 +764,7 @@ class PrefillAdder: else: # Make sure at least one page is available trunc_len = self.rem_chunk_tokens // self.page_size * self.page_size + if trunc_len <= 0: return AddReqResult.OTHER @@ -710,11 +785,8 @@ class PrefillAdder: self.can_run_list.append(req) self.new_chunked_req = req - if self.is_hybrid_swa: - swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node) - req.swa_uuid_for_lock = swa_uuid_for_lock - else: - self.tree_cache.inc_lock_ref(req.last_node) + + self._req_inc_lock_ref(req) self._update_prefill_budget(prefix_len, trunc_len, 0) return self.budget_state() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 812d72859..5471bfeed 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -136,6 +136,7 @@ from sglang.srt.managers.schedule_batch import ( ) from sglang.srt.managers.schedule_policy import ( AddReqResult, + DllmStagingReqs, PrefillAdder, SchedulePolicy, ) @@ -351,6 +352,9 @@ class Scheduler( # Init chunked prefill self.init_chunked_prefill() + # Init diffusion LLM + self.init_diffusion_llm() + # Init schedule policy and new token estimation self.init_schedule_policy() @@ -726,10 +730,6 @@ class Scheduler( def init_chunked_prefill(self): # Init chunked prefill self.chunked_prefill_size = self.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 @@ -752,6 +752,9 @@ class Scheduler( ) self.enable_dynamic_chunking = False + def init_diffusion_llm(self): + self.dllm_staging_reqs = DllmStagingReqs(dllm_config=self.dllm_config) + def init_schedule_policy(self): # Init schedule policy and new token estimation self.policy = SchedulePolicy( @@ -1297,6 +1300,7 @@ class Scheduler( # If it is a health check generation request and there are running requests, ignore it. if is_health_check_generate_req(recv_req) and ( self.chunked_req is not None + or self.dllm_staging_reqs.non_empty() or not self.running_batch.is_empty() or len(self.offload_tags) > 0 ): @@ -1765,27 +1769,37 @@ class Scheduler( for tokenized_req in recv_req: self.handle_embedding_request(tokenized_req) + def stash_chunked_request(self, req: Req): + self.tree_cache.cache_unfinished_req(req, chunked=True) + # Chunked request keeps its rid but will get a new req_pool_idx + if self.tp_worker.model_runner.mambaish_config is not None: + self.req_to_token_pool.free(req.req_pool_idx, free_mamba_cache=False) + else: + self.req_to_token_pool.free(req.req_pool_idx) + def get_next_batch_to_run(self) -> Optional[ScheduleBatch]: self._abort_on_queued_timeout() if self.dllm_config is not None: - if self.chunked_req is not None and self.chunked_req.finished(): - self.chunked_req = None + self.dllm_staging_reqs.filter_finished_reqs() # Merge the prefill batch into the running batch chunked_req_to_exclude = set() - if self.chunked_req: + + if self.dllm_config is not None: + assert ( + self.chunked_req is None + ), "chunked_req should be None when dllm_config is set" + + if self.dllm_staging_reqs.non_empty(): + chunked_req_to_exclude.update(self.dllm_staging_reqs) + for req in self.dllm_staging_reqs: + self.stash_chunked_request(req) + + if self.chunked_req is not None: # Move the chunked request out of the batch so that we can merge # only finished requests to running_batch. chunked_req_to_exclude.add(self.chunked_req) - self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True) - - # chunked request keeps its rid but will get a new req_pool_idx - if self.tp_worker.model_runner.mambaish_config is not None: - self.req_to_token_pool.free( - self.chunked_req.req_pool_idx, free_mamba_cache=False - ) - else: - self.req_to_token_pool.free(self.chunked_req.req_pool_idx) + self.stash_chunked_request(self.chunked_req) if self.last_batch and self.last_batch.forward_mode.is_extend(): if self.last_batch.chunked_req is not None: @@ -1793,6 +1807,9 @@ class Scheduler( # We need to discard it. chunked_req_to_exclude.add(self.last_batch.chunked_req) + if self.last_batch.dllm_staging_reqs.non_empty(): + chunked_req_to_exclude.update(self.last_batch.dllm_staging_reqs) + # Filter batch last_bs = self.last_batch.batch_size() self.last_batch.filter_batch( @@ -1880,20 +1897,20 @@ class Scheduler( # Reset batch_is_full to try preemption with a prefill adder. self.running_batch.batch_is_full = False - if ( - self.running_batch.batch_is_full or len(self.waiting_queue) == 0 - ) and self.chunked_req is None: + if (self.running_batch.batch_is_full or len(self.waiting_queue) == 0) and ( + not self.dllm_staging_reqs.non_empty() and self.chunked_req is None + ): return None running_bs = len(self.running_batch.reqs) # Ignore the check if self.chunked_req is not None. # In the non-PP case, when self.chunked_req is not None, num_allocatable_reqs should always be greater than 0, - # as the space for the chunked request has just been released. - # In PP case, a chunked req can start in one microbatch and end in another microbatch, so the max_running_requests per microbatch should not be strict. - # Instead, we should always allow chunked request to be added, otherwise, there will be a memory leak. + # as the space for the chunked requests has just been released. + # In PP case, chunked requests (or dllm requests) can start in one microbatch and end in another microbatch, so the max_running_requests per microbatch should not be strict. + # Instead, we should always allow chunked requests to be added, otherwise, there will be a memory leak. if ( self.get_num_allocatable_reqs(running_bs) <= 0 - and not self.chunked_req + and (self.dllm_staging_reqs.empty() or self.chunked_req is not None) and not self.try_preemption ): self.running_batch.batch_is_full = True @@ -1932,8 +1949,20 @@ class Scheduler( self.priority_scheduling_preemption_threshold, prefill_max_requests=self.server_args.prefill_max_requests, prefill_delayer_single_pass=prefill_delayer_single_pass, + dllm_config=self.dllm_config, ) + if self.dllm_config is not None: + assert ( + self.chunked_req is None + ), "chunked_req should be None when dllm_config is set" + + if self.dllm_staging_reqs.non_empty(): + self.dllm_staging_reqs.init_next_round() + for req in self.dllm_staging_reqs: + adder.add_chunked_req(req) + self.dllm_staging_reqs.add_reqs(adder.dllm_staging_reqs) + if self.chunked_req is not None: self.chunked_req.init_next_round_input() self.chunked_req = adder.add_chunked_req(self.chunked_req) @@ -1981,7 +2010,9 @@ class Scheduler( req.init_next_round_input(self.tree_cache) res = adder.add_one_req( req, - has_chunked_req=(self.chunked_req is not None), + has_chunked_req=( + self.dllm_staging_reqs.non_empty() or self.chunked_req is not None + ), truncation_align_size=self.truncation_align_size, ) @@ -2013,14 +2044,25 @@ class Scheduler( for req in adder.preempt_list: self._add_request_to_queue(req) - # Update chunked prefill + if self.dllm_config is not None: + assert ( + self.chunked_req is None + ), "chunked_req should be None when dllm_config is set" + + if adder.dllm_staging_reqs.non_empty(): + self.dllm_staging_reqs.add_reqs(adder.dllm_staging_reqs) + if adder.new_chunked_req is not None: + # Update chunked prefill assert self.chunked_req is None self.chunked_req = adder.new_chunked_req - if self.chunked_req: + if self.chunked_req is not None: self.chunked_req.is_chunked += 1 + if self.dllm_staging_reqs.non_empty(): + self.dllm_staging_reqs.update_chunked_status() + # Print stats if self.current_scheduler_metrics_enabled: self.log_prefill_stats( @@ -2049,6 +2091,7 @@ class Scheduler( self.enable_overlap, self.spec_algorithm, chunked_req=self.chunked_req, + dllm_staging_reqs=self.dllm_staging_reqs, dllm_config=self.dllm_config, ) if self.enable_hierarchical_cache: diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index ed614fea9..984025685 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -332,24 +332,26 @@ class SchedulerOutputProcessorMixin: 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() + for idx in range(batch.batch_size()): + # If no new tokens generated, meaning the prefilling stage + if not result.next_token_ids: break - self.tree_cache.cache_unfinished_req(req) + 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.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() diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index cdbb9dec8..056671d51 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -2,7 +2,7 @@ from __future__ import annotations import dataclasses import logging -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, List, Optional, Union import torch @@ -24,7 +24,7 @@ logger = logging.getLogger(__name__) class GenerationBatchResult: logits_output: Optional[LogitsProcessorOutput] = None pp_hidden_states_proxy_tensors: Optional[PPProxyTensors] = None - next_token_ids: Optional[torch.Tensor] = None + next_token_ids: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None num_accepted_tokens: int = 0 accept_length_per_req_cpu: Optional[List[int]] = None can_run_cuda_graph: bool = False diff --git a/test/registered/dllm/test_dllm_batching.py b/test/registered/dllm/test_dllm_batching.py new file mode 100644 index 000000000..ecf755d52 --- /dev/null +++ b/test/registered/dllm/test_dllm_batching.py @@ -0,0 +1,71 @@ +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=500, suite="stage-b-test-large-1-gpu") + +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +""" +Test dLLM batching capability on CUDA GPUs. + +As current dLLM batching performance is suboptimal to BS=1, this test only verifies correctness. +The test will be removed once dLLM batching performance improves. +""" + + +class TestBatching(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = "inclusionAI/LLaDA2.0-mini" + cls.base_url = DEFAULT_URL_FOR_TEST + + other_args = [ + "--trust-remote-code", + "--mem-fraction-static", + "0.9", + "--max-running-requests", + "4", + "--attention-backend", + "flashinfer", + "--dllm-algorithm", + "LowConfidence", + ] + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"{metrics=}") + + self.assertGreater(metrics["accuracy"], 0.88) + + +if __name__ == "__main__": + unittest.main()