sglang 0.4.3.post3__py3-none-any.whl → 0.4.3.post4__py3-none-any.whl

sglang/srt/managers/scheduler.py

@@ -159,17 +159,6 @@ class Scheduler:
         )
         self.gpu_id = gpu_id
         self.enable_hierarchical_cache = server_args.enable_hierarchical_cache
-        self.decode_mem_cache_buf_multiplier = (
-            (
-                self.server_args.speculative_num_draft_tokens
-                + (
-                    self.server_args.speculative_eagle_topk
-                    * self.server_args.speculative_num_draft_tokens
-                )
-            )
-            if not self.spec_algorithm.is_none()
-            else 1
-        )
 
         # Distributed rank info
         self.dp_size = server_args.dp_size
@@ -208,42 +197,12 @@ class Scheduler:
             self.send_to_detokenizer = SimpleNamespace(send_pyobj=lambda x: None)
 
         # Init tokenizer
-        self.model_config = ModelConfig(
-            server_args.model_path,
-            trust_remote_code=server_args.trust_remote_code,
-            revision=server_args.revision,
-            context_length=server_args.context_length,
-            model_override_args=server_args.json_model_override_args,
-            is_embedding=server_args.is_embedding,
-            dtype=server_args.dtype,
-            quantization=server_args.quantization,
-        )
-        self.is_generation = self.model_config.is_generation
-
-        if server_args.skip_tokenizer_init:
-            self.tokenizer = self.processor = None
-        else:
-            if self.model_config.is_multimodal:
-                self.processor = get_processor(
-                    server_args.tokenizer_path,
-                    tokenizer_mode=server_args.tokenizer_mode,
-                    trust_remote_code=server_args.trust_remote_code,
-                    revision=server_args.revision,
-                )
-                self.tokenizer = self.processor.tokenizer
-            else:
-                self.tokenizer = get_tokenizer(
-                    server_args.tokenizer_path,
-                    tokenizer_mode=server_args.tokenizer_mode,
-                    trust_remote_code=server_args.trust_remote_code,
-                    revision=server_args.revision,
-                )
+        self.init_tokenizer()
 
         # Check whether overlap can be enabled
         if not self.is_generation:
             self.enable_overlap = False
             logger.info("Overlap scheduler is disabled for embedding models.")
-
         if self.model_config.is_multimodal:
             self.enable_overlap = False
             logger.info("Overlap scheduler is disabled for multimodal models.")
@@ -274,10 +233,8 @@ class Scheduler:
                 target_worker=self.tp_worker,
                 dp_rank=dp_rank,
             )
-            self.prefill_only_one_req = True
         else:
             self.draft_worker = None
-            self.prefill_only_one_req = False
 
         # Get token and memory info from the model worker
         (
@@ -309,32 +266,7 @@ class Scheduler:
         )
 
         # Init memory pool and cache
-        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
-            self.tp_worker.get_memory_pool()
-        )
-
-        if (
-            server_args.chunked_prefill_size is not None
-            and server_args.disable_radix_cache
-        ):
-            self.tree_cache = ChunkCache(
-                req_to_token_pool=self.req_to_token_pool,
-                token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
-            )
-        else:
-            if self.enable_hierarchical_cache:
-                self.tree_cache = HiRadixCache(
-                    req_to_token_pool=self.req_to_token_pool,
-                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
-                )
-            else:
-                self.tree_cache = RadixCache(
-                    req_to_token_pool=self.req_to_token_pool,
-                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
-                    disable=server_args.disable_radix_cache,
-                )
-
-        self.policy = SchedulePolicy(self.schedule_policy, self.tree_cache)
+        self.init_memory_pool_and_cache()
 
         # Init running status
         self.waiting_queue: List[Req] = []
@@ -348,25 +280,13 @@ class Scheduler:
         self.forward_ct = 0
         self.forward_ct_decode = 0
         self.num_generated_tokens = 0
-        self.spec_num_total_accepted_tokens = 0
-        self.spec_num_total_forward_ct = 0
-        self.cum_spec_accept_length = 0
-        self.cum_spec_accept_count = 0
         self.last_decode_stats_tic = time.time()
         self.return_health_check_ct = 0
         self.current_stream = torch.get_device_module(self.device).current_stream()
         if self.device == "cpu":
             self.current_stream.synchronize = lambda: None  # No-op for CPU
 
-        # For metrics only.
-        # The largest prefill length of a single request
-        self._largest_prefill_len: int = 0
-        # The largest context length (prefill + generation) of a single request
-        self._largest_prefill_decode_len: int = 0
-        self.last_gen_throughput: float = 0.0
-        self.step_time_dict = defaultdict(list)  # Dict[batch size -> step time]
-
-        # Session info
+        # Init session info
         self.sessions: Dict[str, Session] = {}
 
         # Init chunked prefill
@@ -387,11 +307,11 @@ class Scheduler:
         else:
             self.grammar_backend = None
 
-        # Init new token estimation
+        # Init schedule policy and new token estimation
+        self.policy = SchedulePolicy(self.schedule_policy, self.tree_cache)
         assert (
             server_args.schedule_conservativeness >= 0
         ), "Invalid schedule_conservativeness"
-
         self.init_new_token_ratio = min(
             global_config.default_init_new_token_ratio
             * server_args.schedule_conservativeness,
@@ -430,14 +350,7 @@ class Scheduler:
         self.profiler_target_forward_ct: Optional[int] = None
 
         # Init metrics stats
-        self.stats = SchedulerStats()
-        if self.enable_metrics:
-            self.metrics_collector = SchedulerMetricsCollector(
-                labels={
-                    "model_name": self.server_args.served_model_name,
-                    # TODO: Add lora name/path in the future,
-                },
-            )
+        self.init_metrics()
 
         # Init request dispatcher
         self._request_dispatcher = TypeBasedDispatcher(
@@ -460,39 +373,104 @@ class Scheduler:
                 (ResumeMemoryOccupationReqInput, self.resume_memory_occupation),
                 (ProfileReq, self.profile),
                 (GetInternalStateReq, self.get_internal_state),
+                (SetInternalStateReq, self.set_internal_state),
             ]
         )
 
-    def watchdog_thread(self):
-        """A watch dog thread that will try to kill the server itself if one forward batch takes too long."""
-        self.watchdog_last_forward_ct = 0
-        self.watchdog_last_time = time.time()
+    def init_tokenizer(self):
+        server_args = self.server_args
 
-        while True:
-            current = time.time()
-            if self.cur_batch is not None:
-                if self.watchdog_last_forward_ct == self.forward_ct:
-                    if current > self.watchdog_last_time + self.watchdog_timeout:
-                        logger.error(f"Watchdog timeout ({self.watchdog_timeout=})")
-                        break
-                else:
-                    self.watchdog_last_forward_ct = self.forward_ct
-                    self.watchdog_last_time = current
-            time.sleep(self.watchdog_timeout // 2)
+        self.model_config = ModelConfig(
+            server_args.model_path,
+            trust_remote_code=server_args.trust_remote_code,
+            revision=server_args.revision,
+            context_length=server_args.context_length,
+            model_override_args=server_args.json_model_override_args,
+            is_embedding=server_args.is_embedding,
+            dtype=server_args.dtype,
+            quantization=server_args.quantization,
+        )
+        self.is_generation = self.model_config.is_generation
 
-        # Print batch size and memory pool info to check whether there are de-sync issues.
-        logger.error(
-            f"{self.cur_batch.batch_size()=}, "
-            f"{self.cur_batch.reqs=}, "
-            f"{self.token_to_kv_pool.available_size()=}, "
-            f"{self.tree_cache.evictable_size()=}, "
+        if server_args.skip_tokenizer_init:
+            self.tokenizer = self.processor = None
+        else:
+            if self.model_config.is_multimodal:
+                self.processor = get_processor(
+                    server_args.tokenizer_path,
+                    tokenizer_mode=server_args.tokenizer_mode,
+                    trust_remote_code=server_args.trust_remote_code,
+                    revision=server_args.revision,
+                )
+                self.tokenizer = self.processor.tokenizer
+            else:
+                self.tokenizer = get_tokenizer(
+                    server_args.tokenizer_path,
+                    tokenizer_mode=server_args.tokenizer_mode,
+                    trust_remote_code=server_args.trust_remote_code,
+                    revision=server_args.revision,
+                )
+
+    def init_memory_pool_and_cache(self):
+        server_args = self.server_args
+
+        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
+            self.tp_worker.get_memory_pool()
+        )
+
+        if (
+            server_args.chunked_prefill_size is not None
+            and server_args.disable_radix_cache
+        ):
+            self.tree_cache = ChunkCache(
+                req_to_token_pool=self.req_to_token_pool,
+                token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            )
+        else:
+            if self.enable_hierarchical_cache:
+                self.tree_cache = HiRadixCache(
+                    req_to_token_pool=self.req_to_token_pool,
+                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                )
+            else:
+                self.tree_cache = RadixCache(
+                    req_to_token_pool=self.req_to_token_pool,
+                    token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+                    disable=server_args.disable_radix_cache,
+                )
+
+        self.decode_mem_cache_buf_multiplier = (
+            1
+            if self.spec_algorithm.is_none()
+            else (
+                server_args.speculative_num_draft_tokens
+                + (
+                    server_args.speculative_eagle_topk
+                    * server_args.speculative_num_steps
+                )
+            )
         )
-        # Wait for some time so that the parent process can print the error.
-        pyspy_dump_schedulers()
-        print(file=sys.stderr, flush=True)
-        print(file=sys.stdout, flush=True)
-        time.sleep(5)
-        self.parent_process.send_signal(signal.SIGQUIT)
+
+    def init_metrics(self):
+        # The largest prefill length of a single request
+        self._largest_prefill_len: int = 0
+        # The largest context length (prefill + generation) of a single request
+        self._largest_prefill_decode_len: int = 0
+        self.last_gen_throughput: float = 0.0
+        self.step_time_dict = defaultdict(list)  # Dict[batch size -> step time]
+        self.spec_num_total_accepted_tokens = 0
+        self.spec_num_total_forward_ct = 0
+        self.cum_spec_accept_length = 0
+        self.cum_spec_accept_count = 0
+        self.stats = SchedulerStats()
+        if self.enable_metrics:
+            engine_type = "unified"
+            self.metrics_collector = SchedulerMetricsCollector(
+                labels={
+                    "model_name": self.server_args.served_model_name,
+                    "engine_type": engine_type,
+                },
+            )
 
     @torch.no_grad()
     def event_loop_normal(self):
@@ -932,7 +910,7 @@ class Scheduler:
         ):
             # During idle time, also collect metrics every 30 seconds.
             num_used = self.max_total_num_tokens - (
-                self.token_to_kv_pool.available_size()
+                self.token_to_kv_pool_allocator.available_size()
                 + self.tree_cache.evictable_size()
             )
             num_running_reqs = len(self.running_batch.reqs) if self.running_batch else 0
@@ -1077,8 +1055,6 @@ class Scheduler:
                     else:
                         self.batch_is_full = True
                 break
-            if self.prefill_only_one_req:
-                break
 
         # Update waiting queue
         can_run_list: List[Req] = adder.can_run_list
@@ -1180,6 +1156,7 @@ class Scheduler:
         ):
             self.stop_profile()
 
+        # Run forward
         if self.is_generation:
             if self.spec_algorithm.is_none():
                 model_worker_batch = batch.get_model_worker_batch()
@@ -1200,6 +1177,7 @@ class Scheduler:
                 self.spec_num_total_forward_ct += batch.batch_size()
                 self.num_generated_tokens += num_accepted_tokens
             batch.output_ids = next_token_ids
+
             # These 2 values are needed for processing the output, but the values can be
             # modified by overlap schedule. So we have to copy them here so that
             # we can use the correct values in output processing.
@@ -1233,7 +1211,6 @@ class Scheduler:
         result: Union[GenerationBatchResult, EmbeddingBatchResult],
     ):
         if batch.forward_mode.is_decode():
-            assert isinstance(result, GenerationBatchResult)
             self.process_batch_result_decode(batch, result)
             if batch.is_empty():
                 self.running_batch = None
@@ -1485,6 +1462,7 @@ class Scheduler:
             batch.next_batch_sampling_info.update_regex_vocab_mask()
             self.current_stream.synchronize()
             batch.next_batch_sampling_info.sampling_info_done.set()
+
         self.stream_output(batch.reqs, batch.return_logprob)
 
         self.token_to_kv_pool_allocator.free_group_end()
@@ -1588,7 +1566,9 @@ class Scheduler:
                     req.temp_input_token_ids_logprobs_idx
                 )
                 for val, idx in zip(
-                    req.temp_input_top_logprobs_val, req.temp_input_top_logprobs_idx
+                    req.temp_input_top_logprobs_val,
+                    req.temp_input_top_logprobs_idx,
+                    strict=True,
                 ):
                     req.input_top_logprobs_val.extend(val)
                     req.input_top_logprobs_idx.extend(idx)
@@ -1813,14 +1793,18 @@ class Scheduler:
         else:  # embedding or reward model
             embeddings = []
             prompt_tokens = []
+            cached_tokens = []
             for req in reqs:
                 if req.finished():
                     rids.append(req.rid)
                     finished_reasons.append(req.finished_reason.to_json())
                     embeddings.append(req.embedding)
                     prompt_tokens.append(len(req.origin_input_ids))
+                    cached_tokens.append(req.cached_tokens)
             self.send_to_detokenizer.send_pyobj(
-                BatchEmbeddingOut(rids, finished_reasons, embeddings, prompt_tokens)
+                BatchEmbeddingOut(
+                    rids, finished_reasons, embeddings, prompt_tokens, cached_tokens
+                )
             )
 
     def prepare_dp_attn_batch(self, local_batch: ScheduleBatch):
@@ -1906,6 +1890,37 @@ class Scheduler:
         self._extend_requests_to_queue(self.grammar_queue[:num_ready_reqs])
         self.grammar_queue = self.grammar_queue[num_ready_reqs:]
 
+    def watchdog_thread(self):
+        """A watch dog thread that will try to kill the server itself if one forward batch takes too long."""
+        self.watchdog_last_forward_ct = 0
+        self.watchdog_last_time = time.time()
+
+        while True:
+            current = time.time()
+            if self.cur_batch is not None:
+                if self.watchdog_last_forward_ct == self.forward_ct:
+                    if current > self.watchdog_last_time + self.watchdog_timeout:
+                        logger.error(f"Watchdog timeout ({self.watchdog_timeout=})")
+                        break
+                else:
+                    self.watchdog_last_forward_ct = self.forward_ct
+                    self.watchdog_last_time = current
+            time.sleep(self.watchdog_timeout // 2)
+
+        # Print batch size and memory pool info to check whether there are de-sync issues.
+        logger.error(
+            f"{self.cur_batch.batch_size()=}, "
+            f"{self.cur_batch.reqs=}, "
+            f"{self.token_to_kv_pool_allocator.available_size()=}, "
+            f"{self.tree_cache.evictable_size()=}, "
+        )
+        # Wait for some time so that the parent process can print the error.
+        pyspy_dump_schedulers()
+        print(file=sys.stderr, flush=True)
+        print(file=sys.stdout, flush=True)
+        time.sleep(5)
+        self.parent_process.send_signal(signal.SIGQUIT)
+
     def flush_cache_wrapped(self, recv_req: FlushCacheReq):
         self.flush_cache()
 
@@ -1917,7 +1932,6 @@ class Scheduler:
             self.cur_batch = None
             self.last_batch = None
             self.tree_cache.reset()
-            self.tree_cache_metrics = {"total": 0, "hit": 0}
             if self.grammar_backend:
                 self.grammar_backend.reset()
             self.req_to_token_pool.clear()
@@ -2009,6 +2023,9 @@ class Scheduler:
                     req.to_abort = True
                     break
 
+    def _pause_engine(self) -> Tuple[List[Req], int]:
+        raise NotImplementedError()
+
     def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
         """In-place update of the weights from disk."""
         success, message = self.tp_worker.update_weights_from_disk(recv_req)

sglang/srt/managers/tokenizer_manager.py

@@ -1068,6 +1068,7 @@ class TokenizerManager:
             self.metrics_collector.observe_one_finished_request(
                 recv_obj.prompt_tokens[i],
                 completion_tokens,
+                recv_obj.cached_tokens[i],
                 state.finished_time - state.created_time,
             )
 

sglang/srt/mem_cache/memory_pool.py

@@ -20,9 +20,8 @@ Memory pool.
 
 SGLang has two levels of memory pool.
 ReqToTokenPool maps a a request to its token locations.
-TokenToKVPoolAllocator maps a token location to its KV cache data.
-KVCache actually holds the physical kv cache. Allocation indices are allocated
-by TokenToKVPoolAllocator
+TokenToKVPoolAllocator manages the indices to kv cache data.
+KVCache actually holds the physical kv cache.
 """
 
 import abc
@@ -92,14 +91,40 @@ class ReqToTokenPool:
         self.free_slots = list(range(self.size))
 
 
+class KVCache(abc.ABC):
+
+    @abc.abstractmethod
+    def get_key_buffer(self, layer_id: int) -> torch.Tensor:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def set_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+    ) -> None:
+        raise NotImplementedError()
+
+
 class TokenToKVPoolAllocator:
-    """A memory pool that maps a token location to its kv cache data."""
+    """An allocator managing the indices to kv cache data."""
 
     def __init__(
         self,
         size: int,
         dtype: torch.dtype,
         device: str,
+        kvcache: KVCache,
     ):
         self.size = size
         self.dtype = dtype
@@ -110,9 +135,14 @@ class TokenToKVPoolAllocator:
         self.free_group = []
         self.clear()
 
+        self._kvcache = kvcache
+
     def available_size(self):
         return len(self.free_slots)
 
+    def get_kvcache(self):
+        return self._kvcache
+
     def alloc(self, need_size: int):
         if need_size > len(self.free_slots):
             return None
@@ -147,31 +177,6 @@ class TokenToKVPoolAllocator:
         self.free_group = []
 
 
-class KVCache(abc.ABC):
-
-    @abc.abstractmethod
-    def get_key_buffer(self, layer_id: int) -> torch.Tensor:
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
-        raise NotImplementedError()
-
-    @abc.abstractmethod
-    def set_kv_buffer(
-        self,
-        layer: RadixAttention,
-        loc: torch.Tensor,
-        cache_k: torch.Tensor,
-        cache_v: torch.Tensor,
-    ) -> None:
-        raise NotImplementedError()
-
-
 class MHATokenToKVPool(KVCache):
 
     def __init__(

sglang/srt/metrics/collector.py

@@ -121,6 +121,12 @@ class TokenizerMetricsCollector:
             labelnames=labels.keys(),
         )
 
+        self.cached_tokens_total = Counter(
+            name="sglang:cached_tokens_total",
+            documentation="Number of cached prompt tokens.",
+            labelnames=labels.keys(),
+        )
+
         self.num_requests_total = Counter(
             name="sglang:num_requests_total",
             documentation="Number of requests processed.",
@@ -245,10 +251,12 @@ class TokenizerMetricsCollector:
         self,
         prompt_tokens: int,
         generation_tokens: int,
+        cached_tokens: int,
         e2e_latency: float,
     ):
         self.prompt_tokens_total.labels(**self.labels).inc(prompt_tokens)
         self.generation_tokens_total.labels(**self.labels).inc(generation_tokens)
+        self.cached_tokens_total.labels(**self.labels).inc(cached_tokens)
         self.num_requests_total.labels(**self.labels).inc(1)
         self._log_histogram(self.histogram_e2e_request_latency, e2e_latency)
         if generation_tokens >= 1:

sglang/srt/model_executor/cuda_graph_runner.py

@@ -396,16 +396,10 @@ class CudaGraphRunner:
 
             run_once()
 
-        torch.cuda.synchronize()
-        self.model_runner.tp_group.barrier()
-
         global global_graph_memory_pool
         with torch.cuda.graph(graph, pool=global_graph_memory_pool, stream=stream):
             out = run_once()
 
-        torch.cuda.synchronize()
-        self.model_runner.tp_group.barrier()
-
         global_graph_memory_pool = graph.pool()
         return graph, out
 
@@ -427,7 +421,7 @@ class CudaGraphRunner:
             self.capture_hidden_mode = hidden_mode_from_spec_info
             self.capture()
 
-    def replay(self, forward_batch: ForwardBatch):
+    def replay(self, forward_batch: ForwardBatch, skip_attn_backend_init: bool = False):
         self.recapture_if_needed(forward_batch)
 
         raw_bs = forward_batch.batch_size