sglang 0.4.6.post3__py3-none-any.whl → 0.4.6.post5__py3-none-any.whl

sglang/srt/managers/tp_worker.py

@@ -20,7 +20,7 @@ from typing import Optional, Tuple, Union
 import torch
 
 from sglang.srt.configs.model_config import ModelConfig
-from sglang.srt.distributed import get_pp_group, get_tp_group, get_world_group
+from sglang.srt.distributed import get_pp_group, get_world_group
 from sglang.srt.hf_transformers_utils import (
     get_processor,
     get_tokenizer,
@@ -183,8 +183,11 @@ class TpModelWorker:
     def forward_batch_generation(
         self,
         model_worker_batch: ModelWorkerBatch,
+        launch_done: Optional[threading.Event] = None,
         skip_sample: bool = False,
-    ) -> Tuple[Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor]]:
+    ) -> Tuple[
+        Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor], bool
+    ]:
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
 
         pp_proxy_tensors = None
@@ -196,11 +199,11 @@ class TpModelWorker:
             )
 
         if self.pp_group.is_last_rank:
-            logits_output = self.model_runner.forward(
+            logits_output, can_run_cuda_graph = self.model_runner.forward(
                 forward_batch, pp_proxy_tensors=pp_proxy_tensors
             )
-            if model_worker_batch.launch_done is not None:
-                model_worker_batch.launch_done.set()
+            if launch_done is not None:
+                launch_done.set()
 
             if skip_sample:
                 next_token_ids = None
@@ -209,17 +212,17 @@ class TpModelWorker:
                     logits_output, model_worker_batch
                 )
 
-            return logits_output, next_token_ids
+            return logits_output, next_token_ids, can_run_cuda_graph
         else:
-            pp_proxy_tensors = self.model_runner.forward(
+            pp_proxy_tensors, can_run_cuda_graph = self.model_runner.forward(
                 forward_batch,
                 pp_proxy_tensors=pp_proxy_tensors,
             )
-            return pp_proxy_tensors.tensors, None
+            return pp_proxy_tensors.tensors, None, can_run_cuda_graph
 
     def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
-        logits_output = self.model_runner.forward(forward_batch)
+        logits_output, _ = self.model_runner.forward(forward_batch)
         embeddings = logits_output.embeddings
         return embeddings
 

sglang/srt/managers/tp_worker_overlap_thread.py

@@ -18,7 +18,7 @@ import logging
 import signal
 import threading
 from queue import Queue
-from typing import Optional
+from typing import Optional, Tuple
 
 import psutil
 import torch
@@ -127,10 +127,12 @@ class TpModelWorkerClient:
         batch_lists = [None] * 2
 
         while True:
-            model_worker_batch, future_token_ids_ct = self.input_queue.get()
+            model_worker_batch, future_token_ids_ct, sync_event = self.input_queue.get()
             if not model_worker_batch:
                 break
 
+            sync_event.wait()
+
             # Keep a reference of model_worker_batch by storing it into a list.
             # Otherwise, the tensor members of model_worker_batch will be released
             # by pytorch and cause CUDA illegal memory access errors.
@@ -145,8 +147,10 @@ class TpModelWorkerClient:
             resolve_future_token_ids(input_ids, self.future_token_ids_map)
 
             # Run forward
-            logits_output, next_token_ids = self.worker.forward_batch_generation(
-                model_worker_batch
+            logits_output, next_token_ids, can_run_cuda_graph = (
+                self.worker.forward_batch_generation(
+                    model_worker_batch, model_worker_batch.launch_done
+                )
             )
 
             # Update the future token ids map
@@ -171,14 +175,18 @@ class TpModelWorkerClient:
             next_token_ids = next_token_ids.to("cpu", non_blocking=True)
             copy_done.record()
 
-            self.output_queue.put((copy_done, logits_output, next_token_ids))
+            self.output_queue.put(
+                (copy_done, logits_output, next_token_ids, can_run_cuda_graph)
+            )
 
     def resolve_last_batch_result(self, launch_done: Optional[threading.Event] = None):
         """
         This function is called to resolve the last batch result and
         wait for the current batch to be launched. Used in overlap mode.
         """
-        copy_done, logits_output, next_token_ids = self.output_queue.get()
+        copy_done, logits_output, next_token_ids, can_run_cuda_graph = (
+            self.output_queue.get()
+        )
 
         if launch_done is not None:
             launch_done.wait()
@@ -193,9 +201,11 @@ class TpModelWorkerClient:
                     logits_output.input_token_logprobs.tolist()
                 )
         next_token_ids = next_token_ids.tolist()
-        return logits_output, next_token_ids
+        return logits_output, next_token_ids, can_run_cuda_graph
 
-    def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
+    def forward_batch_generation(
+        self, model_worker_batch: ModelWorkerBatch
+    ) -> Tuple[None, torch.Tensor, bool]:
         # Create a new copy of sampling_info because it will be updated in-place by the scheduler for the next batch.
         sampling_info = model_worker_batch.sampling_info
         sampling_info.update_penalties()
@@ -206,10 +216,11 @@ class TpModelWorkerClient:
         )
 
         # A cuda stream sync here to avoid the cuda illegal memory access error.
-        self.scheduler_stream.synchronize()
+        sync_event = torch.get_device_module(self.device).Event()
+        sync_event.record(self.scheduler_stream)
 
         # Push a new batch to the queue
-        self.input_queue.put((model_worker_batch, self.future_token_ids_ct))
+        self.input_queue.put((model_worker_batch, self.future_token_ids_ct, sync_event))
 
         # Allocate output future objects
         bs = len(model_worker_batch.seq_lens)
@@ -223,7 +234,7 @@ class TpModelWorkerClient:
         self.future_token_ids_ct = (
             self.future_token_ids_ct + bs
         ) % self.future_token_ids_limit
-        return None, future_next_token_ids
+        return None, future_next_token_ids, False
 
     def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
         success, message = self.worker.update_weights_from_disk(recv_req)

sglang/srt/mem_cache/base_prefix_cache.py

@@ -48,3 +48,6 @@ class BasePrefixCache(ABC):
 
     def pretty_print(self):
         raise NotImplementedError()
+
+    def take_events(self):
+        return []

sglang/srt/mem_cache/chunk_cache.py

@@ -38,7 +38,9 @@ class ChunkCache(BasePrefixCache):
 
     def cache_finished_req(self, req: Req):
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, : len(req.origin_input_ids) + len(req.output_ids) - 1
+            req.req_pool_idx,
+            # For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
+            : len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0),
         ]
         self.req_to_token_pool.free(req.req_pool_idx)
         self.token_to_kv_pool_allocator.free(kv_indices)

sglang/srt/mem_cache/hiradix_cache.py

@@ -335,13 +335,13 @@ class HiRadixCache(RadixCache):
             return value, last_node
 
     def _match_prefix_helper(self, node: TreeNode, key: List):
-        node.last_access_time = time.time()
+        node.last_access_time = time.monotonic()
         child_key = self.get_child_key_fn(key)
         value = []
 
         while len(key) > 0 and child_key in node.children.keys():
             child = node.children[child_key]
-            child.last_access_time = time.time()
+            child.last_access_time = time.monotonic()
             prefix_len = self.key_match_fn(child.key, key)
             if prefix_len < len(child.key):
                 new_node = self._split_node(child.key, child, prefix_len)
@@ -386,7 +386,7 @@ class HiRadixCache(RadixCache):
         return new_node
 
     def _insert_helper(self, node: TreeNode, key: List, value):
-        node.last_access_time = time.time()
+        node.last_access_time = time.monotonic()
         if len(key) == 0:
             return 0
 
@@ -395,7 +395,7 @@ class HiRadixCache(RadixCache):
 
         while len(key) > 0 and child_key in node.children.keys():
             node = node.children[child_key]
-            node.last_access_time = time.time()
+            node.last_access_time = time.monotonic()
             prefix_len = self.key_match_fn(node.key, key)
 
             if prefix_len == len(node.key):

sglang/srt/mem_cache/memory_pool.py

@@ -38,11 +38,17 @@ import triton
 import triton.language as tl
 
 from sglang.srt.layers.radix_attention import RadixAttention
-from sglang.srt.utils import debug_timing, get_compiler_backend
+from sglang.srt.utils import (
+    debug_timing,
+    get_compiler_backend,
+    is_cuda,
+    next_power_of_2,
+)
 
 logger = logging.getLogger(__name__)
 
 GB = 1024 * 1024 * 1024
+_is_cuda = is_cuda()
 
 
 class ReqToTokenPool:
@@ -94,6 +100,33 @@ class ReqToTokenPool:
 
 
 class KVCache(abc.ABC):
+    @abc.abstractmethod
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        layer_num: int,
+        device: str,
+        enable_memory_saver: bool,
+        start_layer: Optional[int] = None,
+        end_layer: Optional[int] = None,
+    ):
+        self.size = size
+        self.page_size = page_size
+        self.dtype = dtype
+        self.device = device
+        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
+            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
+            self.store_dtype = torch.uint8
+        else:
+            self.store_dtype = dtype
+        self.layer_num = layer_num
+        self.start_layer = start_layer or 0
+        self.end_layer = end_layer or layer_num - 1
+        self.memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+        )
 
     @abc.abstractmethod
     def get_key_buffer(self, layer_id: int) -> torch.Tensor:
@@ -217,30 +250,24 @@ class MHATokenToKVPool(KVCache):
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
     ):
-        self.size = size
-        self.page_size = page_size
-        self.dtype = dtype
-        self.device = device
-        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
-            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
-            self.store_dtype = torch.uint8
-        else:
-            self.store_dtype = dtype
-        self.memory_saver_adapter = TorchMemorySaverAdapter.create(
-            enable=enable_memory_saver
+        super().__init__(
+            size,
+            page_size,
+            dtype,
+            layer_num,
+            device,
+            enable_memory_saver,
+            start_layer,
+            end_layer,
         )
 
         self.head_num = head_num
         self.head_dim = head_dim
-        self.layer_num = layer_num
         self._create_buffers()
-        self.start_layer = start_layer or 0
-        self.end_layer = end_layer or layer_num - 1
 
         self.layer_transfer_counter = None
-        self.capture_mode = False
         self.device_module = torch.get_device_module(self.device)
-        self.alt_stream = self.device_module.Stream()
+        self.alt_stream = self.device_module.Stream() if is_cuda else None
 
         k_size, v_size = self.get_kv_size_bytes()
         logger.info(
@@ -357,6 +384,8 @@ class MHATokenToKVPool(KVCache):
         k_scale: Optional[float] = None,
         v_scale: Optional[float] = None,
     ):
+        from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
+
         layer_id = layer.layer_id
         if cache_k.dtype != self.dtype:
             if k_scale is not None:
@@ -370,7 +399,7 @@ class MHATokenToKVPool(KVCache):
             cache_k = cache_k.view(self.store_dtype)
             cache_v = cache_v.view(self.store_dtype)
 
-        if self.capture_mode and cache_k.shape[0] < 4:
+        if get_is_capture_mode() and self.alt_stream is not None:
             # Overlap the copy of K and V cache for small batch size
             current_stream = self.device_module.current_stream()
             self.alt_stream.wait_stream(current_stream)
@@ -493,26 +522,21 @@ class MLATokenToKVPool(KVCache):
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
     ):
-        self.size = size
-        self.page_size = page_size
-        self.dtype = dtype
-        self.device = device
-        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
-            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
-            self.store_dtype = torch.uint8
-        else:
-            self.store_dtype = dtype
+        super().__init__(
+            size,
+            page_size,
+            dtype,
+            layer_num,
+            device,
+            enable_memory_saver,
+            start_layer,
+            end_layer,
+        )
+
         self.kv_lora_rank = kv_lora_rank
         self.qk_rope_head_dim = qk_rope_head_dim
-        self.layer_num = layer_num
-        self.start_layer = start_layer or 0
-        self.end_layer = end_layer or layer_num - 1
-
-        memory_saver_adapter = TorchMemorySaverAdapter.create(
-            enable=enable_memory_saver
-        )
 
-        with memory_saver_adapter.region():
+        with self.memory_saver_adapter.region():
             # The padded slot 0 is used for writing dummy outputs from padded tokens.
             self.kv_buffer = [
                 torch.zeros(
@@ -524,7 +548,6 @@ class MLATokenToKVPool(KVCache):
             ]
 
         self.layer_transfer_counter = None
-        self.page_size = page_size
 
         kv_size = self.get_kv_size_bytes()
         logger.info(
@@ -637,20 +660,18 @@ class DoubleSparseTokenToKVPool(KVCache):
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
     ):
-        self.size = size
-        self.page_size = page_size
-        self.dtype = dtype
-        self.device = device
-        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
-            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
-            self.store_dtype = torch.uint8
-        else:
-            self.store_dtype = dtype
-        memory_saver_adapter = TorchMemorySaverAdapter.create(
-            enable=enable_memory_saver
+        super().__init__(
+            size,
+            page_size,
+            dtype,
+            layer_num,
+            device,
+            enable_memory_saver,
+            start_layer,
+            end_layer,
         )
 
-        with memory_saver_adapter.region():
+        with self.memory_saver_adapter.region():
             # [size, head_num, head_dim] for each layer
             self.k_buffer = [
                 torch.zeros(
@@ -673,9 +694,6 @@ class DoubleSparseTokenToKVPool(KVCache):
                 for _ in range(layer_num)
             ]
 
-        self.start_layer = start_layer or 0
-        self.end_layer = end_layer or layer_num - 1
-
     def get_key_buffer(self, layer_id: int):
         return self.k_buffer[layer_id - self.start_layer]
 
@@ -743,7 +761,7 @@ class HostKVCache(abc.ABC):
 
     def __init__(
         self,
-        device_pool: MHATokenToKVPool,
+        device_pool: KVCache,
         host_to_device_ratio: float,
         host_size: int,
         pin_memory: bool,
@@ -762,6 +780,8 @@ class HostKVCache(abc.ABC):
             self.size = int(device_pool.size * host_to_device_ratio)
         # Align the host memory pool size to the page size
         self.size = self.size - (self.size % self.page_size)
+        self.start_layer = device_pool.start_layer
+        self.end_layer = device_pool.end_layer
 
         assert (
             self.size > device_pool.size
@@ -913,6 +933,8 @@ class HostKVCache(abc.ABC):
 
 
 class MHATokenToKVPoolHost(HostKVCache):
+    device_pool: MHATokenToKVPool
+
     def __init__(
         self,
         device_pool: MHATokenToKVPool,
@@ -996,6 +1018,8 @@ class MHATokenToKVPoolHost(HostKVCache):
 
 
 class MLATokenToKVPoolHost(HostKVCache):
+    device_pool: MLATokenToKVPool
+
     def __init__(
         self,
         device_pool: MLATokenToKVPool,

sglang/srt/mem_cache/multimodal_cache.py

@@ -0,0 +1,45 @@
+from typing import Dict
+
+import torch
+
+
+class MultiModalCache:
+    """MultiModalCache is used to store vlm encoder results"""
+
+    def __init__(
+        self,
+        max_size: int,
+    ):
+        self.max_size = max_size
+        self.mm_cache: Dict[int, torch.Tensor] = {}
+        self.current_size = 0
+
+    def put(self, mm_hash: int, embedding: torch.Tensor) -> bool:
+        if mm_hash in self.mm_cache:
+            return True
+        data_size = self._get_tensor_size(embedding)
+        if self.current_size + data_size > self.max_size:
+            return False
+        self.mm_cache[mm_hash] = embedding
+        self.current_size += data_size
+        return True
+
+    def get(self, mm_hash: int) -> torch.Tensor:
+        return self.mm_cache.get(mm_hash)
+
+    def free(self, mm_hash: int) -> bool:
+        if mm_hash not in self.mm_cache:
+            return False
+        old_embedding = self.mm_cache.pop(mm_hash)
+        self.current_size -= self._get_tensor_size(old_embedding)
+        return True
+
+    def clear(self):
+        self.mm_cache.clear()
+        self.current_size = 0
+
+    def _get_tensor_size(self, embedding: torch.Tensor):
+        return embedding.element_size() * embedding.numel()
+
+    def __len__(self):
+        return len(self.mm_cache)

sglang/srt/mem_cache/radix_cache.py

@@ -27,6 +27,12 @@ from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 
+from sglang.srt.disaggregation.kv_events import (
+    AllBlocksCleared,
+    BlockRemoved,
+    BlockStored,
+    KVCacheEvent,
+)
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPoolAllocator
@@ -45,7 +51,7 @@ class TreeNode:
         self.key = None
         self.value = None
         self.lock_ref = 0
-        self.last_access_time = time.time()
+        self.last_access_time = time.monotonic()
 
         self.hit_count = 0
         # indicating the node is loading KV cache from host
@@ -96,11 +102,14 @@ class RadixCache(BasePrefixCache):
         token_to_kv_pool_allocator: TokenToKVPoolAllocator,
         page_size: int,
         disable: bool = False,
+        enable_kv_cache_events: bool = False,
     ):
         self.req_to_token_pool = req_to_token_pool
         self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
         self.page_size = page_size
         self.disable = disable
+        self.enable_kv_cache_events = enable_kv_cache_events
+        self.kv_event_queue = []
 
         if self.token_to_kv_pool_allocator:
             self.device = self.token_to_kv_pool_allocator.device
@@ -124,6 +133,7 @@ class RadixCache(BasePrefixCache):
         self.root_node.lock_ref = 1
         self.evictable_size_ = 0
         self.protected_size_ = 0
+        self._record_all_cleared_event()
 
     def match_prefix(self, key: List[int], **kwargs) -> Tuple[torch.Tensor, int]:
         """Find the matching prefix from the radix tree.
@@ -273,6 +283,8 @@ class RadixCache(BasePrefixCache):
             if len(x.parent.children) == 0:
                 heapq.heappush(leaves, x.parent)
 
+            self._record_remove_event(x)
+
     def inc_lock_ref(self, node: TreeNode):
         if self.disable:
             return 0
@@ -322,14 +334,14 @@ class RadixCache(BasePrefixCache):
     ##### Internal Helper Functions #####
 
     def _match_prefix_helper(self, node: TreeNode, key: List):
-        node.last_access_time = time.time()
+        node.last_access_time = time.monotonic()
 
         child_key = self.get_child_key_fn(key)
 
         value = []
         while len(key) > 0 and child_key in node.children.keys():
             child = node.children[child_key]
-            child.last_access_time = time.time()
+            child.last_access_time = time.monotonic()
             prefix_len = self.key_match_fn(child.key, key)
             if prefix_len < len(child.key):
                 new_node = self._split_node(child.key, child, prefix_len)
@@ -348,6 +360,7 @@ class RadixCache(BasePrefixCache):
 
     def _split_node(self, key, child: TreeNode, split_len: int):
         # new_node -> child
+        self._record_remove_event(child)
         new_node = TreeNode()
         new_node.children = {self.get_child_key_fn(key[split_len:]): child}
         new_node.parent = child.parent
@@ -358,10 +371,14 @@ class RadixCache(BasePrefixCache):
         child.key = child.key[split_len:]
         child.value = child.value[split_len:]
         new_node.parent.children[self.get_child_key_fn(key)] = new_node
+
+        self._record_store_event(new_node)
+        self._record_store_event(child)
+
         return new_node
 
     def _insert_helper(self, node: TreeNode, key: List, value):
-        node.last_access_time = time.time()
+        node.last_access_time = time.monotonic()
         if len(key) == 0:
             return 0
 
@@ -370,7 +387,7 @@ class RadixCache(BasePrefixCache):
         total_prefix_length = 0
         while len(key) > 0 and child_key in node.children.keys():
             node = node.children[child_key]
-            node.last_access_time = time.time()
+            node.last_access_time = time.monotonic()
             prefix_len = self.key_match_fn(node.key, key)
             total_prefix_length += prefix_len
             key = key[prefix_len:]
@@ -390,6 +407,7 @@ class RadixCache(BasePrefixCache):
             new_node.value = value
             node.children[child_key] = new_node
             self.evictable_size_ += len(value)
+            self._record_store_event(new_node)
         return total_prefix_length
 
     def _print_helper(self, node: TreeNode, indent: int):
@@ -442,6 +460,41 @@ class RadixCache(BasePrefixCache):
 
         return ret_list
 
+    def _record_store_event(self, node: TreeNode):
+        if self.enable_kv_cache_events:
+            block_hash = hash(tuple(node.key))
+            parent_block_hash = hash(tuple(node.parent.key))
+            self.kv_event_queue.append(
+                BlockStored(
+                    block_hashes=[block_hash],
+                    parent_block_hash=parent_block_hash,
+                    token_ids=node.key,
+                    block_size=len(node.key),
+                    lora_id=None,
+                )
+            )
+
+    def _record_remove_event(self, node: TreeNode):
+        if self.enable_kv_cache_events:
+            block_hash = hash(tuple(node.key))
+            self.kv_event_queue.append(BlockRemoved(block_hashes=[block_hash]))
+
+    def _record_all_cleared_event(self):
+        if self.enable_kv_cache_events:
+            self.kv_event_queue.append(AllBlocksCleared())
+
+    def take_events(self):
+        """Atomically takes all events and clears the queue.
+
+        Returns:
+            A list of KV cache events.
+        """
+        if not self.enable_kv_cache_events:
+            return []
+        events = self.kv_event_queue
+        self.kv_event_queue = []
+        return events
+
 
 if __name__ == "__main__":
     tree = RadixCache(None, None, page_size=1, disable=False)