Minor style fixes to the scheduler.py (#15218)

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(

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

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

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(

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()

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,

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)

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.

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:

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:

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