From 9d64a7b24f644fcd96646abcb545e4495dd8c18b Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 16 Dec 2025 17:09:44 -0800 Subject: [PATCH] Minor style fixes to the scheduler.py (#15218) --- python/sglang/srt/entrypoints/engine.py | 5 +- python/sglang/srt/entrypoints/http_server.py | 2 +- python/sglang/srt/managers/overlap_utils.py | 19 +- python/sglang/srt/managers/schedule_batch.py | 10 +- python/sglang/srt/managers/scheduler.py | 313 +++++++++--------- .../srt/managers/scheduler_metrics_mixin.py | 19 +- .../scheduler_output_processor_mixin.py | 8 +- .../scheduler_runtime_checker_mixin.py | 2 +- python/sglang/srt/managers/tp_worker.py | 18 +- python/sglang/srt/server_args.py | 1 + scripts/ci/ci_install_dependency.sh | 14 +- 11 files changed, 207 insertions(+), 204 deletions(-) diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index cb05f05a2..9aff82b22 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -228,10 +228,7 @@ class Engine(EngineBase): ) self.tokenizer_manager = tokenizer_manager self.template_manager = template_manager - - scheduler_info = scheduler_infos[0] - self.scheduler_info = scheduler_info - + self.scheduler_info = scheduler_infos[0] self.port_args = port_args self.remote_instance_transfer_engine_info = ( parse_remote_instance_transfer_engine_info_from_scheduler_infos( diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 243122ee2..1be55d5ba 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -1432,7 +1432,7 @@ def _execute_server_warmup( for _ in range(120): time.sleep(1) try: - res = requests.get(url + "/get_model_info", timeout=5, headers=headers) + res = requests.get(url + "/model_info", timeout=5, headers=headers) assert res.status_code == 200, f"{res=}, {res.text=}" success = True break diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 07745cf58..c47b69afb 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -55,20 +55,22 @@ class FutureMap: self.future_buffer_len = self.future_limit + 2 * max_running_requests self.device = device self.spec_algo = spec_algo - self.buf_initialized = False if self.spec_algo.is_none(): + # For non-speculative decoding, we only need to store the token ids. + self.buf_initialized = True self.token_ids_buf = torch.empty( (self.future_buffer_len,), dtype=torch.int64, device=self.device ) + else: + # For speculative decoding, we lazily initialize the buffers + # This is to make the shape derivation easier. + self.buf_initialized = False def _lazy_init_buf(self, draft_input: EagleDraftInput): - if self.buf_initialized or not self.spec_algo.is_eagle(): - return - self.buf_initialized = True - # get the template for each tensor + # Get a reference for each tensor topk_p0 = draft_input.topk_p[0] topk_index0 = draft_input.topk_index[0] hidden_states0 = draft_input.hidden_states[0] @@ -147,10 +149,13 @@ class FutureMap: self, future_indices: FutureIndices, draft_input: EagleDraftInput ): intv = future_indices.interval - # idle indices do not need store info if self.is_empty_slice(intv): + # idle indices in dp attention do not need store info return - self._lazy_init_buf(draft_input) + + if not self.buf_initialized: + self._lazy_init_buf(draft_input) + self.topk_p_buf[intv] = draft_input.topk_p self.topk_index_buf[intv] = draft_input.topk_index self.hidden_states_buf[intv] = draft_input.hidden_states diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 87eb98758..53eea9b5a 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1819,15 +1819,15 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ) @property - def is_v2_eagle(self): - # FIXME: finally deprecate is_v2_eagle + def is_eagle_v2(self): + # FIXME: finally deprecate is_eagle_v2 return self.enable_overlap and self.spec_algorithm.is_eagle() def prepare_for_decode(self): self.forward_mode = ForwardMode.DECODE bs = len(self.reqs) - if self.is_v2_eagle: + if self.is_eagle_v2: # TODO(spec-v2): all v2 spec should go through this path draft_input: EagleDraftInput = self.spec_info draft_input.prepare_for_decode(self) @@ -1907,7 +1907,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ) def maybe_wait_verify_done(self): - if self.is_v2_eagle: + if self.is_eagle_v2: draft_input: EagleDraftInput = self.spec_info if draft_input.verify_done is not None: draft_input.verify_done.synchronize() @@ -1980,7 +1980,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # NOTE: spec_info filtered before batch filtering only happens in: # - Spec v1's verify phase # - Only for decode batch (running_batch) - has_been_filtered = v1_spec_info_filtered and not self.is_v2_eagle + has_been_filtered = v1_spec_info_filtered and not self.is_eagle_v2 if self.spec_info: self.spec_info.filter_batch( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index e01a53829..0c3fb83f5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -296,6 +296,9 @@ class Scheduler( # Init diffusion LLM config self.dllm_config = DllmConfig.from_server_args(server_args) + # Init metrics stats + self.init_metrics(tp_rank, pp_rank, dp_rank) + # Init inter-process communication self.init_sockets(server_args, port_args) @@ -430,9 +433,6 @@ class Scheduler( f"{'available_cpu_mem' if self.device == 'cpu' else 'available_gpu_mem'}={avail_mem:.2f} GB" ) - # Init metrics stats - self.init_metrics(tp_rank, pp_rank, dp_rank) - # Init cache using the existing memory pool self.init_cache_with_memory_pool() @@ -447,18 +447,11 @@ class Scheduler( # The last forward batch self.last_batch: Optional[ScheduleBatch] = None self.forward_ct = 0 - self.forward_ct_decode = 0 - self.num_generated_tokens = 0 self.last_prefill_tokens = 0 self.return_health_check_ct = 0 self.num_retracted_reqs: int = 0 self.num_paused_reqs: int = 0 self.sessions: Dict[str, Session] = {} - self.default_stream: CudaStream = torch.get_device_module( - self.device - ).current_stream() - if self.device == "cpu": - self.default_stream.synchronize = lambda: None # No-op for CPU self.forward_sleep_time = None self._engine_paused = False @@ -696,7 +689,7 @@ class Scheduler( self.send_to_tokenizer = SenderWrapper(None) self.send_to_detokenizer = SenderWrapper(None) - if self.current_scheduler_metrics_enabled(): + if self.current_scheduler_metrics_enabled: self.send_metrics_from_scheduler = get_zmq_socket( context, zmq.PUSH, port_args.metrics_ipc_name, False ) @@ -974,20 +967,22 @@ class Scheduler( self.disagg_prefill_inflight_queue: List[Req] = [] def init_overlap(self): - self.future_map = None - if not self.enable_overlap and self.pp_size == 1: - return + self.device_module = torch.get_device_module(self.device) + self.default_stream: CudaStream = self.device_module.current_stream() + if self.device == "cpu": + self.default_stream.synchronize = lambda: None # No-op for CPU - self.forward_stream: CudaStream = torch.get_device_module(self.device).Stream() - self.forward_stream_ctx: CudaStreamContext = torch.get_device_module( - self.device - ).stream(self.forward_stream) - self.copy_stream: CudaStream = torch.get_device_module(self.device).Stream() - self.copy_stream_ctx: CudaStreamContext = torch.get_device_module( - self.device - ).stream(self.copy_stream) + self.forward_stream: CudaStream = self.device_module.Stream() + self.forward_stream_ctx: CudaStreamContext = self.device_module.stream( + self.forward_stream + ) + self.copy_stream: CudaStream = self.device_module.Stream() + self.copy_stream_ctx: CudaStreamContext = self.device_module.stream( + self.copy_stream + ) if not self.enable_overlap: + self.future_map = None return self.future_map = FutureMap( @@ -1000,15 +995,6 @@ class Scheduler( self.batch_record_buf = [None] * 2 self.batch_record_ct = 0 - def record_batch_in_overlap(self, model_worker_batch: ModelWorkerBatch): - # FIXME(lsyin): hacky way to keep a reference to avoid GPU tensors being freed by torch GC - # NOTE: More Reliable: record all tensors into the forward stream - # NOTE: - for all future tensors, we shall always read from future map - # - for all non-future tensors (produced only by schedule stream), - # we shall keep its reference not being release during all the forwarding pass - self.batch_record_ct = (self.batch_record_ct + 1) % 2 - self.batch_record_buf[self.batch_record_ct] = model_worker_batch - def init_moe_config(self): if hasattr(self.model_config.hf_config, "num_experts_per_tok"): initialize_moe_config(self.server_args) @@ -1023,15 +1009,17 @@ class Scheduler( def event_loop_normal(self): """A normal scheduler loop.""" while True: + # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) - if self._engine_paused: continue + # Get the next batch to run batch = self.get_next_batch_to_run() self.cur_batch = batch + # Launch the current batch if batch: result = self.run_batch(batch) self.process_batch_result(batch, result) @@ -1039,8 +1027,8 @@ class Scheduler( # When the server is idle, do self-check and re-init some states self.self_check_during_idle() + # Update the last batch self.last_batch = batch - if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.self_check_during_busy() @@ -1048,9 +1036,6 @@ class Scheduler( def event_loop_overlap(self): """A scheduler loop that overlaps the CPU processing and GPU computation.""" self.result_queue: Deque[Tuple[ScheduleBatch, GenerationBatchResult]] = deque() - disable_consecutive_prefill_overlap = ( - envs.SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.get() - ) def pop_and_process(): # Process the results of the last batch @@ -1058,53 +1043,69 @@ class Scheduler( self.process_batch_result(tmp_batch, tmp_result) while True: + # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) - if self._engine_paused: continue + # Get the next batch to run batch = self.get_next_batch_to_run() self.cur_batch = batch + disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) - disable_overlap_for_batch = ( - disable_consecutive_prefill_overlap - and batch - and batch.forward_mode.is_extend() - and self.last_batch - and self.last_batch.forward_mode.is_extend() - ) - - # FIXME(lsyin): remove this grammar sync - need_grammar_sync = ( - batch is not None - and batch.forward_mode.is_decode() - and batch.has_grammar - and batch.is_v2_eagle - and len(self.result_queue) > 0 - ) - - if disable_overlap_for_batch or need_grammar_sync: + # If we do not need to overlap the current batch with the last batch, + # we can process the last batch immediately. + if disable_overlap_for_batch: pop_and_process() + # Launch the current batch batch_result = None if batch: batch_result = self.run_batch(batch) self.result_queue.append((batch.copy(), batch_result)) + # Process the last batch if self.last_batch: - if not disable_overlap_for_batch and not need_grammar_sync: + if not disable_overlap_for_batch: pop_and_process() elif batch is None: # When the server is idle, do self-check and re-init some states self.self_check_during_idle() + # Run sample of the current batch + # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed. self.launch_batch_sample_if_needed(batch_result) - self.last_batch = batch + # Update the last batch + self.last_batch = batch if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.self_check_during_busy() + def is_disable_overlap_for_batch(self, batch: ScheduleBatch) -> bool: + # For two consecutive prefill batches, we disable overlap to improve the TTFT of the first batch. + # This might slightly hurt the throughput, so we use an environment variable to control it. + disable_overlap_for_batch = ( + envs.SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.get() + and batch + and batch.forward_mode.is_extend() + and self.last_batch + and self.last_batch.forward_mode.is_extend() + ) + + # We do not support overlap + spec + grammar yet, + # so we need to turn off overlap for this batch. + # TODO(lsyin): support overlap + spec + grammar + need_grammar_sync = ( + batch + and batch.is_eagle_v2 + and batch.has_grammar + and batch.forward_mode.is_decode() + and len(self.result_queue) > 0 + ) + + return disable_overlap_for_batch or need_grammar_sync + def recv_limit_reached(self, num_recv_reqs: int) -> bool: if self.max_recv_per_poll < 0: return False @@ -1163,32 +1164,7 @@ class Scheduler( if self.server_args.enable_dp_attention: if self.attn_tp_rank == 0: - work_reqs = [ - req - for req in recv_reqs - if isinstance( - req, - ( - TokenizedGenerateReqInput, - TokenizedEmbeddingReqInput, - BatchTokenizedGenerateReqInput, - BatchTokenizedEmbeddingReqInput, - ), - ) - ] - control_reqs = [ - req - for req in recv_reqs - if not isinstance( - req, - ( - TokenizedGenerateReqInput, - TokenizedEmbeddingReqInput, - BatchTokenizedGenerateReqInput, - BatchTokenizedEmbeddingReqInput, - ), - ) - ] + work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs) else: work_reqs = None control_reqs = None @@ -1226,9 +1202,37 @@ class Scheduler( return recv_reqs - def process_input_requests(self, recv_reqs: List): + def _split_work_and_control_reqs(self, recv_reqs: List): + work_reqs = [ + req + for req in recv_reqs + if isinstance( + req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + BatchTokenizedEmbeddingReqInput, + ), + ) + ] + control_reqs = [ + req + for req in recv_reqs + if not isinstance( + req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + BatchTokenizedEmbeddingReqInput, + ), + ) + ] + return work_reqs, control_reqs - # Process MM requests under E disaggregation + def process_input_requests(self, recv_reqs: List): + # Process MM requests under EPD-disaggregation mode if ( self.server_args.language_only and self.server_args.encoder_transfer_backend == "zmq_to_scheduler" @@ -1247,11 +1251,11 @@ class Scheduler( output = self._request_dispatcher(recv_req) if output is not None: - if isinstance(output, RpcReqOutput): + if not isinstance(output, RpcReqOutput): + self.send_to_tokenizer.send_output(output, recv_req) + else: if self.recv_from_rpc is not None: self.recv_from_rpc.send_pyobj(output) - else: - self.send_to_tokenizer.send_output(output, recv_req) def init_req_max_new_tokens(self, req): req.sampling_params.max_new_tokens = min( @@ -1782,6 +1786,7 @@ class Scheduler( ): # Decrease prefill idle as much as possible during high dp load. return None + # Check if the grammar is ready in the grammar queue if self.grammar_queue: self.move_ready_grammar_requests() @@ -1882,9 +1887,9 @@ class Scheduler( self.running_batch.batch_is_full = True if self.running_batch.batch_is_full: - if not self.try_preemption: - break - if not adder.preempt_to_schedule(req, self.server_args): + if not self.try_preemption or not adder.preempt_to_schedule( + req, self.server_args + ): break if self.enable_hicache_storage: @@ -1928,6 +1933,7 @@ class Scheduler( for req in adder.preempt_list: self._add_request_to_queue(req) + # Update chunked prefill if adder.new_chunked_req is not None: assert self.chunked_req is None self.chunked_req = adder.new_chunked_req @@ -1936,12 +1942,12 @@ class Scheduler( self.chunked_req.is_chunked += 1 # Print stats - if self.current_scheduler_metrics_enabled(): + if self.current_scheduler_metrics_enabled: self.log_prefill_stats(adder, can_run_list, running_bs, 0) + # Record metrics for req in can_run_list: if req.time_stats.forward_entry_time == 0: - # Avoid update chunked request many times req.time_stats.forward_entry_time = time.perf_counter() if self.enable_metrics: self.metrics_collector.observe_queue_time( @@ -2034,11 +2040,14 @@ class Scheduler( batch.prepare_for_decode() return batch - # placeholder for override - def update_cache_from_scheduler( - self, schedule_batch: ScheduleBatch, batch_result: GenerationBatchResult - ): - pass + def record_batch_in_overlap(self, model_worker_batch: ModelWorkerBatch): + # FIXME(lsyin): hacky way to keep a reference to avoid GPU tensors being freed by torch GC + # NOTE: More Reliable: record all tensors into the forward stream + # NOTE: - for all future tensors, we shall always read from future map + # - for all non-future tensors (produced only by schedule stream), + # we shall keep its reference not being release during all the forwarding pass + self.batch_record_ct = (self.batch_record_ct + 1) % 2 + self.batch_record_buf[self.batch_record_ct] = model_worker_batch def run_batch( self, @@ -2066,19 +2075,19 @@ class Scheduler( # Run forward if self.is_generation: - batch_or_worker_batch = batch - - if self.enable_overlap or self.spec_algorithm.is_none(): - # FIXME(lsyin): remove this if and finally unify the abstraction - batch_or_worker_batch = batch.get_model_worker_batch() + if self.spec_algorithm.is_none() or self.enable_overlap: + # In most cases, we use the model worker batch to run the forward. + worker_batch_or_batch = batch.get_model_worker_batch() + else: + # In speculative decoding v1 (non-overlap) case, we use the batch directly. + # TODO(lsyin): delete this branch after unifying the abstraction. + worker_batch_or_batch = batch if self.enable_overlap: - # FIXME: remove this assert - assert isinstance(batch_or_worker_batch, ModelWorkerBatch) - model_worker_batch = batch_or_worker_batch + model_worker_batch = worker_batch_or_batch self.record_batch_in_overlap(model_worker_batch) - # Sampling info will be modified during forward + # Sampling info will be modified during forward, so we store a copy. model_worker_batch.sampling_info = ( model_worker_batch.sampling_info.copy_for_forward() ) @@ -2094,9 +2103,7 @@ class Scheduler( # here pp is not compatible with overlap ) # FIXME(lsyin): maybe move this to forward_batch_generation - batch_result.copy_done = torch.get_device_module( - self.device - ).Event() + batch_result.copy_done = self.device_module.Event() if batch_result.delay_sample_func is None: self.future_map.store_to_map(future_indices, batch_result) batch_result.copy_to_cpu(return_logprob=batch.return_logprob) @@ -2106,7 +2113,7 @@ class Scheduler( # FIXME(lsyin): move this assignment elsewhere future_indices_or_next_token_ids = -future_indices.indices - if batch.is_v2_eagle: + if batch.is_eagle_v2: # FIXME(lsyin): tmp code for eagle v2 # We only keep future indices for next draft input @@ -2131,7 +2138,7 @@ class Scheduler( else {} ) batch_result = self.model_worker.forward_batch_generation( - batch_or_worker_batch, **kwargs + worker_batch_or_batch, **kwargs ) future_indices_or_next_token_ids = batch_result.next_token_ids self.update_cache_from_scheduler(batch, batch_result) @@ -2146,21 +2153,16 @@ class Scheduler( # modified by overlap schedule. So we have to copy them here so that # we can use the correct values in output processing. if batch.return_logprob or self.spec_algorithm.is_eagle(): - extend_input_len_per_req = [req.extend_input_len for req in batch.reqs] - else: - extend_input_len_per_req = None - - if batch.return_logprob: - extend_logprob_start_len_per_req = [ + batch_result.extend_input_len_per_req = [ + req.extend_input_len for req in batch.reqs + ] + batch_result.extend_logprob_start_len_per_req = [ req.extend_logprob_start_len for req in batch.reqs ] else: - extend_logprob_start_len_per_req = None + batch_result.extend_input_len_per_req = None + batch_result.extend_logprob_start_len_per_req = None - batch_result.extend_input_len_per_req = extend_input_len_per_req - batch_result.extend_logprob_start_len_per_req = ( - extend_logprob_start_len_per_req - ) ret = batch_result else: # embedding or reward model model_worker_batch = batch.get_model_worker_batch() @@ -2327,31 +2329,27 @@ class Scheduler( self.cur_batch = None self.last_batch = None self.tree_cache.reset() - if self.grammar_backend: - self.grammar_backend.reset() self.req_to_token_pool.clear() self.token_to_kv_pool_allocator.clear() + if self.grammar_backend: + self.grammar_backend.reset() + self.reset_metrics() if self.draft_worker: self.draft_worker.clear_cache_pool() - self.num_generated_tokens = 0 - self.forward_ct_decode = 0 - self.spec_num_accepted_tokens = 0 - self.spec_num_forward_ct = 0 - self.spec_total_num_accepted_tokens = 0 - self.spec_total_num_forward_ct = 0 + # TODO: allow optional empty cache torch.cuda.empty_cache() logger.info("Cache flushed successfully!") - if_success = True + success = True else: logging.warning( f"Cache not flushed because there are pending requests. " f"#queue-req: {len(self.waiting_queue)}, " f"#running-req: {len(self.running_batch.reqs)}" ) - if_success = False - return if_success + success = False + return success def get_internal_state(self, recv_req: GetInternalStateReq): ret = vars(get_global_server_args()) @@ -2362,16 +2360,14 @@ class Scheduler( self.token_to_kv_pool_allocator.get_kvcache().mem_usage, 2 ), "token_capacity": int(self.max_total_num_tokens), + "graph": round(self.tp_worker.model_runner.graph_mem_usage, 2), } - ret["memory_usage"]["graph"] = round( - self.tp_worker.model_runner.graph_mem_usage, 2 - ) - if not self.spec_algorithm.is_none() and self.spec_total_num_forward_ct > 0: ret["avg_spec_accept_length"] = ( self.spec_total_num_accepted_tokens / self.spec_total_num_forward_ct ) + if RECORD_STEP_TIME: ret["step_time_dict"] = self.step_time_dict @@ -2389,6 +2385,7 @@ class Scheduler( "speculative_accept_threshold_acc", ] ) + if_success = True for k, v in server_args_dict.items(): if k not in args_allow_update: @@ -2403,6 +2400,7 @@ class Scheduler( ) if_success = False break + if if_success: if not self.spec_algorithm.is_none() and self.spec_total_num_forward_ct > 0: avg_spec_accept_length = ( @@ -2633,19 +2631,6 @@ class Scheduler( else: del self.sessions[session_id] - def get_print_prefix(self): - prefix = "" - if self.attn_dp_rank is not None: - prefix += f" DP{self.attn_dp_rank}" - if self.server_args.tp_size > 1: - prefix += f" TP{self.tp_rank}" - if self.pp_size > 1: - prefix += f" PP{self.pp_rank}" - return prefix - - def current_scheduler_metrics_enabled(self): - return self.attn_tp_rank == 0 or self.enable_metrics_for_all_schedulers - def maybe_sleep_on_idle(self): if self.idle_sleeper is not None: self.idle_sleeper.maybe_sleep() @@ -2656,6 +2641,12 @@ class Scheduler( self.send_to_detokenizer.send_output(recv_req, recv_req) return None + # placeholder for override + def update_cache_from_scheduler( + self, schedule_batch: ScheduleBatch, batch_result: GenerationBatchResult + ): + pass + def get_remote_instance_transfer_engine_info(self): return self.tp_worker.get_remote_instance_transfer_engine_info() @@ -2791,6 +2782,8 @@ def run_scheduler_process( ) pipe_writer.send(result_dict) + + # Dispatch to the appropriate event loop based on the disaggregation mode disaggregation_mode: DisaggregationMode = scheduler.disaggregation_mode if disaggregation_mode == DisaggregationMode.NULL: if scheduler.enable_pdmux: @@ -2802,22 +2795,20 @@ def run_scheduler_process( else: scheduler.event_loop_normal() elif disaggregation_mode == DisaggregationMode.PREFILL: - if scheduler.enable_overlap: + if server_args.pp_size > 1: + scheduler.event_loop_pp_disagg_prefill() + elif scheduler.enable_overlap: scheduler.event_loop_overlap_disagg_prefill() else: - if server_args.pp_size > 1: - scheduler.event_loop_pp_disagg_prefill() - else: - scheduler.event_loop_normal_disagg_prefill() + scheduler.event_loop_normal_disagg_prefill() elif disaggregation_mode == DisaggregationMode.DECODE: - if scheduler.enable_overlap: + if server_args.pp_size > 1: + scheduler.event_loop_pp_disagg_decode() + elif scheduler.enable_overlap: scheduler.event_loop_overlap_disagg_decode() else: - if server_args.pp_size > 1: - scheduler.event_loop_pp_disagg_decode() - else: - scheduler.event_loop_normal_disagg_decode() + scheduler.event_loop_normal_disagg_decode() except Exception: traceback = get_exception_traceback() diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index 6fe5fdc81..bf67ff5d2 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -39,9 +39,11 @@ class SchedulerMetricsMixin: def init_metrics( self: Scheduler, tp_rank: int, pp_rank: int, dp_rank: Optional[int] ): + # Basic stats + self.forward_ct_decode = 0 + self.num_generated_tokens = 0 self.last_decode_stats_tic = time.perf_counter() self.last_prefill_stats_tic = time.perf_counter() - self.last_gen_throughput: float = 0.0 self.last_input_throughput: float = 0.0 self.step_time_dict = defaultdict(list) # Dict[batch size -> step time] @@ -52,6 +54,8 @@ class SchedulerMetricsMixin: # The total number of accepted tokens and forward ct for the whole server lifetime self.spec_total_num_accepted_tokens = 0 self.spec_total_num_forward_ct = 0 + + # For PD disaggregation self.kv_transfer_speed_gb_s: float = 0.0 self.kv_transfer_latency_ms: float = 0.0 self.kv_transfer_bootstrap_ms: float = 0.0 @@ -59,6 +63,11 @@ class SchedulerMetricsMixin: self.stats = SchedulerStats() + # Metrics + self.current_scheduler_metrics_enabled = ( + self.attn_tp_rank == 0 or self.enable_metrics_for_all_schedulers + ) + if self.enable_metrics: engine_type = "unified" labels = { @@ -82,6 +91,14 @@ class SchedulerMetricsMixin: self.spec_num_forward_ct += bs self.num_generated_tokens += num_accepted_tokens + def reset_metrics(self): + self.forward_ct_decode = 0 + self.num_generated_tokens = 0 + self.spec_num_accepted_tokens = 0 + self.spec_num_forward_ct = 0 + self.spec_total_num_accepted_tokens = 0 + self.spec_total_num_forward_ct = 0 + def log_prefill_stats( self: Scheduler, adder: PrefillAdder, diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index b3144c5cd..f96ed1dcd 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -331,7 +331,7 @@ class SchedulerOutputProcessorMixin: next_token_ids = next_token_ids.tolist() if batch.return_logprob: next_token_logprobs = logits_output.next_token_logprobs.tolist() - elif batch.is_v2_eagle: + elif batch.is_eagle_v2: next_token_ids = self._resolve_spec_overlap_token_ids(result, batch) self.num_generated_tokens += len(batch.reqs) @@ -356,7 +356,7 @@ class SchedulerOutputProcessorMixin: new_accepted_len = 1 if batch.spec_algorithm.is_none(): req.output_ids.append(next_token_id) - elif batch.is_v2_eagle: + elif batch.is_eagle_v2: # Only v2 eagle's output_ids are updated here. req.output_ids.extend(next_token_id) new_accepted_len = len(next_token_id) @@ -406,7 +406,7 @@ class SchedulerOutputProcessorMixin: if batch.spec_algorithm.is_none(): # Normal decode: single token req.grammar.accept_token(next_token_id) - elif batch.is_v2_eagle: + elif batch.is_eagle_v2: # Speculative decode: next_token_id is a list of accepted tokens for token_id in next_token_id: req.grammar.accept_token(token_id) @@ -424,7 +424,7 @@ class SchedulerOutputProcessorMixin: self.forward_ct_decode = (self.forward_ct_decode + 1) % (1 << 30) if ( - self.current_scheduler_metrics_enabled() + self.current_scheduler_metrics_enabled and self.forward_ct_decode % self.server_args.decode_log_interval == 0 ): self.log_decode_stats(can_run_cuda_graph, running_batch=batch) diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 25e3d3429..a3749578d 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -246,7 +246,7 @@ class SchedulerRuntimeCheckerMixin: if ( self.enable_metrics - and self.current_scheduler_metrics_enabled() + and self.current_scheduler_metrics_enabled and time.perf_counter() > self.metrics_collector.last_log_time + 30 ): # During idle time, also collect metrics every 30 seconds. diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index b3d933df4..bed034246 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -383,6 +383,7 @@ class TpModelWorker(BaseTpWorker): # FIXME(lsyin): maybe remove skip_attn_backend_init in forward_batch_generation, # which requires preparing replay to always be in this function + # Get forward batch from model worker batch if model_worker_batch is not None: # update the consumer index of hicache to the running batch self.set_hicache_consumer(model_worker_batch.hicache_consumer_index) @@ -392,10 +393,10 @@ class TpModelWorker(BaseTpWorker): # FIXME(lsyin): unify the interface of forward_batch assert forward_batch is not None - if self.pp_group.is_last_rank: - if self.is_dllm(): - return self._forward_batch_generation_dllm(forward_batch) + if self.is_dllm(): + return self._forward_batch_generation_dllm(forward_batch) + if self.pp_group.is_last_rank: logits_output, can_run_cuda_graph = self.model_runner.forward( forward_batch, pp_proxy_tensors=pp_proxy_tensors, @@ -425,7 +426,12 @@ class TpModelWorker(BaseTpWorker): batch_result.delay_sample_func = sample_batch_func return batch_result - if model_worker_batch.is_prefill_only: + if not model_worker_batch.is_prefill_only: + # For normal requests, sample the next token ids. + batch_result.next_token_ids = self.model_runner.sample( + logits_output, forward_batch + ) + else: # For prefill-only requests, create dummy token IDs on CPU # The size should match the batch size (number of sequences), not total tokens batch_result.next_token_ids = torch.zeros( @@ -441,10 +447,6 @@ class TpModelWorker(BaseTpWorker): self.model_runner.compute_logprobs_only( logits_output, model_worker_batch ) - else: - batch_result.next_token_ids = self.model_runner.sample( - logits_output, forward_batch - ) return batch_result else: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bff7fbf60..6df68aaba 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1966,6 +1966,7 @@ class ServerArgs: self.disable_overlap_schedule = True logger.warning( "Overlap scheduler is disabled because of using eagle3 or standalone speculative decoding." + "You can set env SGLANG_ENABLE_SPEC_V2=True to enable the experimental overlap scheduler." ) if self.enable_mixed_chunk: diff --git a/scripts/ci/ci_install_dependency.sh b/scripts/ci/ci_install_dependency.sh index 4dc9bf99a..2b5627671 100755 --- a/scripts/ci/ci_install_dependency.sh +++ b/scripts/ci/ci_install_dependency.sh @@ -43,6 +43,7 @@ fi # Install protoc for router build (gRPC protobuf compilation) if [ "${INSTALL_PROTOC:-0}" = "1" ]; then + # TODO: move this to a separate script echo "Installing protoc..." if command -v apt-get &> /dev/null; then # Ubuntu/Debian @@ -108,6 +109,7 @@ $PIP_CMD install -e "python[${EXTRAS}]" --extra-index-url https://download.pytor # Install router for pd-disagg test $PIP_CMD install sglang-router $PIP_INSTALL_SUFFIX +# Remove flash_attn folder to avoid conflicts PYTHON_LIB_PATH=$(python3 -c "import site; print(site.getsitepackages()[0])") FLASH_ATTN_PATH="${PYTHON_LIB_PATH}/flash_attn" @@ -215,15 +217,3 @@ python3 -c "import torch; print(torch.version.cuda)" # Prepare the CI runner (cleanup HuggingFace cache, etc.) bash "${SCRIPT_DIR}/prepare_runner.sh" - -# Remove flash_attn folder to avoid conflicts with sgl-kernel -PYTHON_LIB_PATH=$(python3 -c "import site; print(site.getsitepackages()[0])") -FLASH_ATTN_PATH="${PYTHON_LIB_PATH}/flash_attn" - -if [ -d "$FLASH_ATTN_PATH" ]; then - echo "Directory $FLASH_ATTN_PATH exists. Removing..." - rm -rf "$FLASH_ATTN_PATH" - echo "error: this should not happen" -else - echo "Directory $FLASH_ATTN_PATH does not exist." -fi