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

sglang/srt/mem_cache/hiradix_cache.py

@@ -16,7 +16,6 @@ from sglang.srt.mem_cache.memory_pool import (
     TokenToKVPoolAllocator,
 )
 from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode
-from sglang.srt.mem_cache.radix_cache import _key_match_page_size1 as _key_match
 
 logger = logging.getLogger(__name__)
 
@@ -31,29 +30,25 @@ class HiRadixCache(RadixCache):
         page_size: int,
         hicache_ratio: float,
     ):
-        if page_size != 1:
-            raise ValueError(
-                "Page size larger than 1 is not yet supported in HiRadixCache."
-            )
         self.kv_cache = token_to_kv_pool_allocator.get_kvcache()
         if isinstance(self.kv_cache, MHATokenToKVPool):
             self.token_to_kv_pool_host = MHATokenToKVPoolHost(
-                self.kv_cache, hicache_ratio
+                self.kv_cache, hicache_ratio, page_size
             )
         elif isinstance(self.kv_cache, MLATokenToKVPool):
             self.token_to_kv_pool_host = MLATokenToKVPoolHost(
-                self.kv_cache, hicache_ratio
+                self.kv_cache, hicache_ratio, page_size
             )
         else:
-            raise ValueError(f"Only MHA and MLA supports swap kv_cache to host.")
+            raise ValueError(f"HiRadixCache only supports MHA and MLA yet")
 
         self.tp_group = tp_cache_group
-        self.page_size = page_size
 
         self.load_cache_event = threading.Event()
         self.cache_controller = HiCacheController(
             token_to_kv_pool_allocator,
             self.token_to_kv_pool_host,
+            page_size,
             load_cache_event=self.load_cache_event,
         )
 
@@ -65,7 +60,7 @@ class HiRadixCache(RadixCache):
         self.write_through_threshold = 1
         self.load_back_threshold = 10
         super().__init__(
-            req_to_token_pool, token_to_kv_pool_allocator, self.page_size, disable=False
+            req_to_token_pool, token_to_kv_pool_allocator, page_size, disable=False
         )
 
     def reset(self):
@@ -210,9 +205,9 @@ class HiRadixCache(RadixCache):
             # only evict the host value of evicted nodes
             if not x.evicted:
                 continue
-            assert x.lock_ref == 0 and x.host_value is not None
 
-            assert self.cache_controller.evict_host(x.host_value) > 0
+            num_evicted += self.cache_controller.evict_host(x.host_value)
+
             for k, v in x.parent.children.items():
                 if v == x:
                     break
@@ -299,18 +294,26 @@ class HiRadixCache(RadixCache):
 
         return last_node, prefix_indices
 
-    def read_to_load_cache(self):
+    def ready_to_load_cache(self):
         self.load_cache_event.set()
 
     def match_prefix(self, key: List[int], include_evicted=False, **kwargs):
-        if self.disable:
-            return [], self.root_node
+        empty_value = torch.empty((0,), dtype=torch.int64, device=self.device)
+        if self.disable or len(key) == 0:
+            if include_evicted:
+                return empty_value, self.root_node, self.root_node
+            else:
+                return empty_value, self.root_node
+
+        if self.page_size != 1:
+            page_aligned_len = len(key) // self.page_size * self.page_size
+            key = key[:page_aligned_len]
 
         value, last_node = self._match_prefix_helper(self.root_node, key)
         if value:
             value = torch.cat(value)
         else:
-            value = torch.tensor([], dtype=torch.int64)
+            value = empty_value
 
         last_node_global = last_node
         while last_node.evicted:
@@ -323,11 +326,13 @@ class HiRadixCache(RadixCache):
 
     def _match_prefix_helper(self, node: TreeNode, key: List):
         node.last_access_time = time.time()
+        child_key = self.get_child_key_fn(key)
         value = []
-        while len(key) > 0 and key[0] in node.children.keys():
-            child = node.children[key[0]]
+
+        while len(key) > 0 and child_key in node.children.keys():
+            child = node.children[child_key]
             child.last_access_time = time.time()
-            prefix_len = _key_match(child.key, key)
+            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)
                 if not new_node.evicted:
@@ -339,12 +344,16 @@ class HiRadixCache(RadixCache):
                     value.append(child.value)
                 node = child
                 key = key[prefix_len:]
+
+                if len(key):
+                    child_key = self.get_child_key_fn(key)
+
         return value, node
 
     def _split_node(self, key, child: TreeNode, split_len: int):
         # child node split into new_node -> child
         new_node = TreeNode()
-        new_node.children = {key[split_len]: child}
+        new_node.children = {self.get_child_key_fn(key[split_len:]): child}
         new_node.parent = child.parent
         new_node.lock_ref = child.lock_ref
         new_node.key = child.key[:split_len]
@@ -361,7 +370,7 @@ class HiRadixCache(RadixCache):
             child.host_value = child.host_value[split_len:]
         child.parent = new_node
         child.key = child.key[split_len:]
-        new_node.parent.children[key[0]] = new_node
+        new_node.parent.children[self.get_child_key_fn(key)] = new_node
         return new_node
 
     def _insert_helper(self, node: TreeNode, key: List, value):
@@ -369,52 +378,53 @@ class HiRadixCache(RadixCache):
         if len(key) == 0:
             return 0
 
-        if key[0] in node.children.keys():
-            child = node.children[key[0]]
-            prefix_len = _key_match(child.key, key)
+        child_key = self.get_child_key_fn(key)
+        total_prefix_length = 0
 
-            if prefix_len == len(child.key):
-                if child.evicted:
+        while len(key) > 0 and child_key in node.children.keys():
+            node = node.children[child_key]
+            node.last_access_time = time.time()
+            prefix_len = self.key_match_fn(node.key, key)
+
+            if prefix_len == len(node.key):
+                if node.evicted:
                     # change the reference if the node is evicted
                     # this often happens in the case of KV cache recomputation
-                    child.value = value[:prefix_len]
-                    self.token_to_kv_pool_host.update_synced(child.host_value)
-                    self.evictable_size_ += len(value[:prefix_len])
-                    return self._insert_helper(
-                        child, key[prefix_len:], value[prefix_len:]
-                    )
+                    node.value = value[:prefix_len]
+                    self.token_to_kv_pool_host.update_synced(node.host_value)
+                    self.evictable_size_ += len(node.value)
                 else:
-                    self.inc_hit_count(child)
-                    return prefix_len + self._insert_helper(
-                        child, key[prefix_len:], value[prefix_len:]
-                    )
-
-            # partial match, split the node
-            new_node = self._split_node(child.key, child, prefix_len)
-            if new_node.evicted:
-                new_node.value = value[:prefix_len]
-                self.token_to_kv_pool_host.update_synced(new_node.host_value)
-                self.evictable_size_ += len(new_node.value)
-                return self._insert_helper(
-                    new_node, key[prefix_len:], value[prefix_len:]
-                )
+                    self.inc_hit_count(node)
+                    total_prefix_length += prefix_len
             else:
-                self.inc_hit_count(new_node)
-                return prefix_len + self._insert_helper(
-                    new_node, key[prefix_len:], value[prefix_len:]
-                )
+                # partial match, split the node
+                new_node = self._split_node(node.key, node, prefix_len)
+                if new_node.evicted:
+                    new_node.value = value[:prefix_len]
+                    self.token_to_kv_pool_host.update_synced(new_node.host_value)
+                    self.evictable_size_ += len(new_node.value)
+                else:
+                    self.inc_hit_count(new_node)
+                    total_prefix_length += prefix_len
+                node = new_node
+
+            key = key[prefix_len:]
+            value = value[prefix_len:]
+
+            if len(key):
+                child_key = self.get_child_key_fn(key)
 
         if len(key):
             new_node = TreeNode()
             new_node.parent = node
             new_node.key = key
             new_node.value = value
-            node.children[key[0]] = new_node
+            node.children[child_key] = new_node
             self.evictable_size_ += len(value)
 
             if self.cache_controller.write_policy == "write_through":
                 self.write_backup(new_node)
-        return 0
+        return total_prefix_length
 
     def _collect_leaves_device(self):
         def is_leaf(node):

sglang/srt/mem_cache/memory_pool.py

@@ -185,6 +185,12 @@ class TokenToKVPoolAllocator:
         if self.free_group:
             self.free(torch.cat(self.free_group))
 
+    def backup_state(self):
+        return self.free_slots
+
+    def restore_state(self, free_slots):
+        self.free_slots = free_slots
+
     def clear(self):
         # The padded slot 0 is used for writing dummy outputs from padded tokens.
         self.free_slots = torch.arange(
@@ -602,8 +608,9 @@ class HostKVCache(abc.ABC):
         self,
         device_pool: MHATokenToKVPool,
         host_to_device_ratio: float,
-        pin_memory: bool = False,  # no need to use pin memory with the double buffering
-        device: str = "cpu",
+        pin_memory: bool,
+        device: str,
+        page_size: int,
     ):
         assert (
             host_to_device_ratio >= 1
@@ -614,8 +621,11 @@ class HostKVCache(abc.ABC):
         self.host_to_device_ratio = host_to_device_ratio
         self.pin_memory = pin_memory
         self.device = device
+        self.page_size = page_size
 
         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.dtype = device_pool.store_dtype
         self.size_per_token = self.get_size_per_token()
 
@@ -769,10 +779,13 @@ class MHATokenToKVPoolHost(HostKVCache):
         self,
         device_pool: MHATokenToKVPool,
         host_to_device_ratio: float,
-        pin_memory: bool = False,  # no need to use pin memory with the double buffering
+        page_size: int,
+        pin_memory: bool = True,
         device: str = "cpu",
     ):
-        super().__init__(device_pool, host_to_device_ratio, pin_memory, device)
+        super().__init__(
+            device_pool, host_to_device_ratio, pin_memory, device, page_size
+        )
 
     def get_size_per_token(self):
         self.head_num = self.device_pool.head_num
@@ -805,16 +818,48 @@ class MHATokenToKVPoolHost(HostKVCache):
     def assign_flat_data(self, indices, flat_data):
         self.kv_buffer[:, :, indices] = flat_data
 
+    def write_page_all_layers(self, host_indices, device_indices, device_pool):
+        device_indices_cpu = device_indices[:: self.page_size].cpu()
+        for i in range(len(device_indices_cpu)):
+            h_index = host_indices[i * self.page_size]
+            d_index = device_indices_cpu[i]
+            for j in range(self.layer_num):
+                self.kv_buffer[0, j, h_index : h_index + self.page_size].copy_(
+                    device_pool.k_buffer[j][d_index : d_index + self.page_size],
+                    non_blocking=True,
+                )
+                self.kv_buffer[1, j, h_index : h_index + self.page_size].copy_(
+                    device_pool.v_buffer[j][d_index : d_index + self.page_size],
+                    non_blocking=True,
+                )
+
+    def load_page_per_layer(self, host_indices, device_indices, device_pool, layer_id):
+        device_indices_cpu = device_indices[:: self.page_size].cpu()
+        for i in range(len(device_indices_cpu)):
+            h_index = host_indices[i * self.page_size]
+            d_index = device_indices_cpu[i]
+            device_pool.k_buffer[layer_id][d_index : d_index + self.page_size].copy_(
+                self.kv_buffer[0, layer_id, h_index : h_index + self.page_size],
+                non_blocking=True,
+            )
+            device_pool.v_buffer[layer_id][d_index : d_index + self.page_size].copy_(
+                self.kv_buffer[1, layer_id, h_index : h_index + self.page_size],
+                non_blocking=True,
+            )
+
 
 class MLATokenToKVPoolHost(HostKVCache):
     def __init__(
         self,
         device_pool: MLATokenToKVPool,
         host_to_device_ratio: float,
-        pin_memory: bool = False,  # no need to use pin memory with the double buffering
+        page_size: int,
+        pin_memory: bool = True,
         device: str = "cpu",
     ):
-        super().__init__(device_pool, host_to_device_ratio, pin_memory, device)
+        super().__init__(
+            device_pool, host_to_device_ratio, pin_memory, device, page_size
+        )
 
     def get_size_per_token(self):
         self.kv_lora_rank = self.device_pool.kv_lora_rank
@@ -851,3 +896,24 @@ class MLATokenToKVPoolHost(HostKVCache):
 
     def assign_flat_data(self, indices, flat_data):
         self.kv_buffer[:, indices] = flat_data
+
+    def write_page_all_layers(self, host_indices, device_indices, device_pool):
+        device_indices_cpu = device_indices[:: self.page_size].cpu()
+        for i in range(len(device_indices_cpu)):
+            h_index = host_indices[i * self.page_size]
+            d_index = device_indices_cpu[i]
+            for j in range(self.layer_num):
+                self.kv_buffer[j, h_index : h_index + self.page_size].copy_(
+                    device_pool.kv_buffer[j][d_index : d_index + self.page_size],
+                    non_blocking=True,
+                )
+
+    def load_page_per_layer(self, host_indices, device_indices, device_pool, layer_id):
+        device_indices_cpu = device_indices[:: self.page_size].cpu()
+        for i in range(len(device_indices_cpu)):
+            h_index = host_indices[i * self.page_size]
+            d_index = device_indices_cpu[i]
+            device_pool.kv_buffer[layer_id][d_index : d_index + self.page_size].copy_(
+                self.kv_buffer[layer_id, h_index : h_index + self.page_size],
+                non_blocking=True,
+            )

sglang/srt/mem_cache/paged_allocator.py

@@ -190,6 +190,30 @@ class PagedTokenToKVPoolAllocator:
     def available_size(self):
         return len(self.free_pages) * self.page_size
 
+    def get_kvcache(self):
+        return self._kvcache
+
+    def alloc(self, need_size: int):
+        # page-aligned allocation, returning contiguous indices of pages
+        if self.debug_mode:
+            assert (
+                need_size % self.page_size == 0
+            ), "The allocation size should be page-aligned"
+
+        num_pages = need_size // self.page_size
+        if num_pages > len(self.free_pages):
+            return None
+
+        out_pages = self.free_pages[:num_pages]
+        self.free_pages = self.free_pages[num_pages:]
+
+        out_indices = (
+            out_pages[:, None] * self.page_size
+            + torch.arange(self.page_size, device=self.device)
+        ).reshape(-1)
+
+        return out_indices
+
     def alloc_extend(
         self,
         prefix_lens: torch.Tensor,
@@ -218,6 +242,9 @@ class PagedTokenToKVPoolAllocator:
             next_power_of_2(extend_num_tokens),
         )
 
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+
         merged_value = self.ret_values.item()
         num_new_pages = merged_value >> 32
         if num_new_pages > len(self.free_pages):
@@ -248,6 +275,9 @@ class PagedTokenToKVPoolAllocator:
             self.page_size,
         )
 
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+
         num_new_pages = self.ret_values.item()
         if num_new_pages > len(self.free_pages):
             return None
@@ -265,6 +295,9 @@ class PagedTokenToKVPoolAllocator:
         else:
             self.free_group.append(free_index)
 
+        if self.debug_mode:
+            assert len(torch.unique(self.free_pages)) == len(self.free_pages)
+
     def free_group_begin(self):
         self.is_not_in_free_group = False
         self.free_group = []
@@ -274,6 +307,12 @@ class PagedTokenToKVPoolAllocator:
         if self.free_group:
             self.free(torch.cat(self.free_group))
 
+    def backup_state(self):
+        return self.free_pages
+
+    def restore_state(self, free_pages):
+        self.free_pages = free_pages
+
     def clear(self):
         # The padded slot 0 is used for writing dummy outputs from padded tokens.
         self.free_pages = torch.arange(

sglang/srt/metrics/collector.py

@@ -33,7 +33,7 @@ class SchedulerMetricsCollector:
 
     def __init__(self, labels: Dict[str, str]) -> None:
         # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
-        from prometheus_client import Gauge
+        from prometheus_client import Gauge, Histogram
 
         self.labels = labels
         self.last_log_time = time.time()
@@ -139,10 +139,10 @@ class TokenizerMetricsCollector:
             labelnames=labels.keys(),
             buckets=[
                 0.1,
-                0.3,
-                0.5,
-                0.7,
-                0.9,
+                0.2,
+                0.4,
+                0.6,
+                0.8,
                 1,
                 2,
                 4,
@@ -153,36 +153,9 @@ class TokenizerMetricsCollector:
                 40,
                 60,
                 80,
-                120,
-                160,
-            ],
-        )
-
-        self.histogram_time_per_output_token = Histogram(
-            name="sglang:time_per_output_token_seconds",
-            documentation="Histogram of time per output token in seconds.",
-            labelnames=labels.keys(),
-            buckets=[
-                0.002,
-                0.005,
-                0.010,
-                0.020,
-                0.030,
-                0.040,
-                0.050,
-                0.060,
-                0.070,
-                0.080,
-                0.090,
-                0.100,
-                0.150,
-                0.200,
-                0.300,
-                0.400,
-                0.600,
-                0.800,
-                1.000,
-                2.000,
+                100,
+                200,
+                400,
             ],
         )
 
@@ -202,17 +175,18 @@ class TokenizerMetricsCollector:
                 0.030,
                 0.035,
                 0.040,
-                0.050,
-                0.075,
+                0.060,
+                0.080,
                 0.100,
-                0.150,
                 0.200,
-                0.300,
                 0.400,
-                0.500,
-                0.750,
+                0.600,
+                0.800,
                 1.000,
                 2.000,
+                4.000,
+                6.000,
+                8.000,
             ],
         )
 
@@ -224,23 +198,22 @@ class TokenizerMetricsCollector:
                 0.1,
                 0.2,
                 0.4,
+                0.6,
                 0.8,
                 1,
                 2,
-                5,
+                4,
+                6,
+                8,
                 10,
                 20,
                 40,
                 60,
                 80,
                 100,
-                150,
                 200,
-                250,
-                300,
-                350,
-                500,
-                1000,
+                400,
+                800,
             ],
         )
 
@@ -256,13 +229,10 @@ class TokenizerMetricsCollector:
     ):
         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)
+        if cached_tokens > 0:
+            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:
-            self.histogram_time_per_output_token.labels(**self.labels).observe(
-                e2e_latency / generation_tokens
-            )
 
     def observe_time_to_first_token(self, value: float):
         self.histogram_time_to_first_token.labels(**self.labels).observe(value)

sglang/srt/model_executor/cuda_graph_runner.py

@@ -116,16 +116,18 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
     if capture_bs is None:
         if server_args.speculative_algorithm is None:
             if server_args.disable_cuda_graph_padding:
-                capture_bs = list(range(1, 33)) + [64, 96, 128, 160]
+                capture_bs = list(range(1, 33)) + range(40, 161, 16)
             else:
-                capture_bs = [1, 2, 4] + [i * 8 for i in range(1, 21)]
+                capture_bs = [1, 2, 4, 8] + list(range(16, 161, 8))
         else:
             # Since speculative decoding requires more cuda graph memory, we
             # capture less.
-            capture_bs = list(range(1, 9)) + list(range(9, 33, 2)) + [64, 96, 128, 160]
+            capture_bs = (
+                list(range(1, 9)) + list(range(10, 33, 2)) + list(range(40, 161, 16))
+            )
 
         if _is_hip:
-            capture_bs += [i * 8 for i in range(21, 33)]
+            capture_bs += list(range(160, 257, 8))
 
     if max(capture_bs) > model_runner.req_to_token_pool.size:
         # In some case (e.g., with a small GPU or --max-running-requests), the #max-running-requests
@@ -489,10 +491,10 @@ class CudaGraphRunner:
         self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
         self.out_cache_loc[:raw_num_token].copy_(forward_batch.out_cache_loc)
         self.positions[:raw_num_token].copy_(forward_batch.positions)
-        if forward_batch.decode_seq_lens_cpu is not None:
+        if forward_batch.seq_lens_cpu is not None:
             if bs != raw_bs:
                 self.seq_lens_cpu.fill_(1)
-            self.seq_lens_cpu[:raw_bs].copy_(forward_batch.decode_seq_lens_cpu)
+            self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
 
         if self.is_encoder_decoder:
             self.encoder_lens[:raw_bs].copy_(forward_batch.encoder_lens)

sglang/srt/model_executor/forward_batch_info.py

@@ -104,6 +104,9 @@ class ForwardMode(IntEnum):
             or self == ForwardMode.IDLE
         )
 
+    def is_extend_or_draft_extend(self):
+        return self == ForwardMode.EXTEND or self == ForwardMode.DRAFT_EXTEND
+
     def is_dummy_first(self):
         return self == ForwardMode.DUMMY_FIRST
 
@@ -148,6 +151,9 @@ class ForwardBatch:
     # The sum of all sequence lengths
     seq_lens_sum: int
 
+    # Optional seq_lens on cpu
+    seq_lens_cpu: Optional[torch.Tensor] = None
+
     # For logprob
     return_logprob: bool = False
     top_logprobs_nums: Optional[List[int]] = None
@@ -162,9 +168,6 @@ class ForwardBatch:
     # Position information
     positions: torch.Tensor = None
 
-    # For decode
-    decode_seq_lens_cpu: Optional[torch.Tensor] = None
-
     # For extend
     extend_num_tokens: Optional[int] = None
     extend_seq_lens: Optional[torch.Tensor] = None
@@ -293,12 +296,14 @@ class ForwardBatch:
         ):
             ret.positions = ret.spec_info.positions
 
+        # Get seq_lens_cpu if needed
+        if ret.seq_lens_cpu is None:
+            ret.seq_lens_cpu = batch.seq_lens_cpu
+
         # Init position information
         if ret.forward_mode.is_decode():
             if ret.positions is None:
                 ret.positions = clamp_position(batch.seq_lens)
-            if ret.decode_seq_lens_cpu is None:
-                ret.decode_seq_lens_cpu = batch.decode_seq_lens
         else:
             ret.extend_seq_lens = torch.tensor(
                 batch.extend_seq_lens, dtype=torch.int32
@@ -353,11 +358,6 @@ class ForwardBatch:
         for mm_input in valid_inputs[1:]:
             merged.merge(mm_input)
 
-        if isinstance(merged.pixel_values, np.ndarray):
-            merged.pixel_values = torch.from_numpy(merged.pixel_values)
-        if isinstance(merged.audio_features, np.ndarray):
-            merged.audio_features = torch.from_numpy(merged.audio_features)
-
         return merged
 
     def contains_image_inputs(self) -> bool: