sglang 0.5.1.post3__py3-none-any.whl → 0.5.2rc1__py3-none-any.whl

sglang/srt/mem_cache/hiradix_cache.py

@@ -102,10 +102,7 @@ class HiRadixCache(RadixCache):
         self.ongoing_backup = {}
         # todo: dynamically adjust the threshold
         self.write_through_threshold = (
-            1 if hicache_write_policy == "write_through" else 3
-        )
-        self.write_through_threshold_storage = (
-            1 if hicache_write_policy == "write_through" else 3
+            1 if hicache_write_policy == "write_through" else 2
         )
         self.load_back_threshold = 10
         super().__init__(
@@ -125,6 +122,15 @@ class HiRadixCache(RadixCache):
             height += 1
         return height
 
+    def clear_storage_backend(self):
+        if self.enable_storage:
+            self.cache_controller.storage_backend.clear()
+            logger.info("Hierarchical cache storage backend cleared successfully!")
+            return True
+        else:
+            logger.warning("Hierarchical cache storage backend is not enabled.")
+            return False
+
     def write_backup(self, node: TreeNode, write_back=False):
         host_indices = self.cache_controller.write(
             device_indices=node.value,
@@ -155,8 +161,9 @@ class HiRadixCache(RadixCache):
         self.ongoing_backup[operation_id] = node
         node.protect_host()
 
-    def inc_hit_count(self, node: TreeNode):
-        if self.cache_controller.write_policy == "write_back":
+    def _inc_hit_count(self, node: TreeNode, chunked=False):
+        # skip the hit count update for chunked requests
+        if self.cache_controller.write_policy == "write_back" or chunked:
             return
         node.hit_count += 1
 
@@ -164,14 +171,6 @@ class HiRadixCache(RadixCache):
             if node.hit_count >= self.write_through_threshold:
                 # write to host if the node is not backuped
                 self.write_backup(node)
-        else:
-            if (
-                self.enable_storage
-                and (not node.backuped_storage)
-                and node.hit_count >= self.write_through_threshold_storage
-            ):
-                # if the node is backuped on host memory but not on storage
-                self.write_backup_storage(node)
 
     def writing_check(self, write_back=False):
         if write_back:
@@ -192,8 +191,11 @@ class HiRadixCache(RadixCache):
             )
         for _ in range(queue_size.item()):
             ack_id = self.cache_controller.ack_write_queue.get()
-            self.dec_lock_ref(self.ongoing_write_through[ack_id])
+            backuped_node = self.ongoing_write_through[ack_id]
+            self.dec_lock_ref(backuped_node)
             del self.ongoing_write_through[ack_id]
+            if self.enable_storage:
+                self.write_backup_storage(backuped_node)
 
     def loading_check(self):
         while not self.cache_controller.ack_load_queue.empty():
@@ -376,57 +378,54 @@ class HiRadixCache(RadixCache):
         self.writing_check()
         self.loading_check()
         if self.enable_storage:
-            self.check_revoked_prefetch()
-            self.check_backup_progress()
-
-    def check_revoked_prefetch(self):
-        queue_size = torch.tensor(
-            self.cache_controller.prefetch_revoke_queue.qsize(), dtype=torch.int
+            self.drain_storage_control_queues()
+
+    def drain_storage_control_queues(self):
+        """
+        Combine prefetch revoke, backup ack, and host mem release checks
+        to minimize TP synchronization and Python overhead.
+        """
+        cc = self.cache_controller
+
+        qsizes = torch.tensor(
+            [
+                cc.prefetch_revoke_queue.qsize(),
+                cc.ack_backup_queue.qsize(),
+                cc.host_mem_release_queue.qsize(),
+            ],
+            dtype=torch.int,
         )
         if self.tp_world_size > 1:
-            # synchrnoize TP workers to make the same update to hiradix cache
             torch.distributed.all_reduce(
-                queue_size,
-                op=torch.distributed.ReduceOp.MIN,
-                group=self.tp_group,
+                qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
             )
-        for _ in range(queue_size.item()):
-            req_id = self.cache_controller.prefetch_revoke_queue.get()
-            if req_id in self.ongoing_prefetch:
-                last_host_node, token_ids, _, _ = self.ongoing_prefetch[req_id]
-                last_host_node.release_host()
-                del self.ongoing_prefetch[req_id]
-                self.cache_controller.prefetch_tokens_occupied -= len(token_ids)
-            else:
-                # the revoked operation already got terminated
-                pass
 
-    def check_backup_progress(self):
-        queue_size = torch.tensor(
-            self.cache_controller.ack_backup_queue.qsize(), dtype=torch.int
-        )
-        if self.tp_world_size > 1:
-            # synchrnoize TP workers to make the same update to hiradix cache
-            torch.distributed.all_reduce(
-                queue_size,
-                op=torch.distributed.ReduceOp.MIN,
-                group=self.tp_group,
-            )
-        for _ in range(queue_size.item()):
-            ack_id, completed_tokens = self.cache_controller.ack_backup_queue.get()
-            host_node = self.ongoing_backup[ack_id]
-
-            if completed_tokens > 0:
-                if completed_tokens < len(host_node.key):
-                    # backup is only partially successful, split the node
-                    new_node = self._split_node(
-                        host_node.key, host_node, completed_tokens
-                    )
-                    new_node.backuped_storage = True
-                else:
-                    host_node.backuped_storage = True
-            host_node.release_host()
-            del self.ongoing_backup[ack_id]
+        n_revoke, n_backup, n_release = map(int, qsizes.tolist())
+
+        # process prefetch revokes
+        for _ in range(n_revoke):
+            req_id = cc.prefetch_revoke_queue.get()
+            info = self.ongoing_prefetch.pop(req_id, None)
+            if info is not None:
+                last_host_node, token_ids, _, _ = info
+                last_host_node.release_host()
+                cc.prefetch_tokens_occupied -= len(token_ids)
+            # else: the revoked operation already got terminated, nothing to do
+
+        # process backup acks
+        for _ in range(n_backup):
+            ack_id = cc.ack_backup_queue.get()
+            entry = self.ongoing_backup.pop(ack_id, None)
+            if entry is not None:
+                entry.release_host()
+
+        # release host memory
+        host_indices_list = []
+        for _ in range(n_release):
+            host_indices_list.append(cc.host_mem_release_queue.get())
+        if host_indices_list:
+            host_indices = torch.cat(host_indices_list, dim=0)
+            cc.mem_pool_host.free(host_indices)
 
     def can_terminate_prefetch(self, operation: PrefetchOperation):
         can_terminate = True
@@ -469,9 +468,9 @@ class HiRadixCache(RadixCache):
 
         # todo: more policies for prefetch progress such as timeout
         # the current policy is to prefetch with best effort and terminate when queuing is over
-        last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch[
+        last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch.pop(
             req_id
-        ]
+        )
 
         if operation.host_indices is None:
             # prefetch has not been issued due to insufficient host memory
@@ -509,11 +508,10 @@ class HiRadixCache(RadixCache):
             self.cache_controller.mem_pool_host.update_prefetch(written_indices)
 
         self.cache_controller.mem_pool_host.free(host_indices[:matched_length])
-        self.cache_controller.mem_pool_host.free(
+        self.cache_controller.append_host_mem_release(
             host_indices[min_completed_tokens:completed_tokens]
         )
         last_host_node.release_host()
-        del self.ongoing_prefetch[req_id]
         self.cache_controller.prefetch_tokens_occupied -= len(token_ids)
 
         return True
@@ -565,7 +563,11 @@ class HiRadixCache(RadixCache):
             len(new_input_tokens) % self.page_size
         )
         new_input_tokens = new_input_tokens[:prefetch_length]
-        if not self.enable_storage or prefetch_length < self.prefetch_threshold:
+        if (
+            not self.enable_storage
+            or prefetch_length < self.prefetch_threshold
+            or self.cache_controller.prefetch_rate_limited()
+        ):
             return
 
         last_host_node.protect_host()
@@ -573,6 +575,10 @@ class HiRadixCache(RadixCache):
         if host_indices is None:
             self.evict_host(prefetch_length)
             host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
+        if host_indices is None:
+            last_host_node.release_host()
+            # no sufficient host memory for prefetch
+            return
         operation = self.cache_controller.prefetch(
             req_id, host_indices, new_input_tokens, last_hash
         )
@@ -672,11 +678,11 @@ class HiRadixCache(RadixCache):
         new_node.parent.children[self.get_child_key_fn(key)] = new_node
         return new_node
 
-    def _insert_helper(self, node: TreeNode, key: List, value):
-        node.last_access_time = time.monotonic()
+    def insert(self, key: List, value, chunked=False):
         if len(key) == 0:
             return 0
 
+        node = self.root_node
         child_key = self.get_child_key_fn(key)
         total_prefix_length = 0
 
@@ -693,7 +699,7 @@ class HiRadixCache(RadixCache):
                     self.token_to_kv_pool_host.update_synced(node.host_value)
                     self.evictable_size_ += len(node.value)
                 else:
-                    self.inc_hit_count(node)
+                    self._inc_hit_count(node, chunked)
                     total_prefix_length += prefix_len
             else:
                 # partial match, split the node
@@ -703,7 +709,7 @@ class HiRadixCache(RadixCache):
                     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)
+                    self._inc_hit_count(new_node, chunked)
                     total_prefix_length += prefix_len
                 node = new_node
 
@@ -737,7 +743,7 @@ class HiRadixCache(RadixCache):
                     last_hash = new_node.hash_value[-1]
 
             if self.cache_controller.write_policy != "write_back":
-                self.inc_hit_count(new_node)
+                self._inc_hit_count(new_node, chunked)
         return total_prefix_length
 
     def _collect_leaves_device(self):
@@ -764,3 +770,20 @@ class HiRadixCache(RadixCache):
                     if not cur_child.evicted:
                         stack.append(cur_child)
         return ret_list
+
+    def release_aborted_request(self, rid: str):
+        if rid not in self.ongoing_prefetch:
+            return
+
+        last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch.pop(
+            rid
+        )
+        if operation.host_indices is None:
+            return
+
+        completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
+        if self.tp_world_size > 1:
+            torch.distributed.barrier(group=self.tp_group)
+        last_host_node.release_host()
+        self.cache_controller.append_host_mem_release(host_indices[:completed_tokens])
+        self.cache_controller.prefetch_tokens_occupied -= len(token_ids)

sglang/srt/mem_cache/lora_radix_cache.py

@@ -183,7 +183,7 @@ class LoRARadixCache(BasePrefixCache):
         self.req_to_token_pool.free(req.req_pool_idx)
         self.dec_lock_ref(req.last_node)
 
-    def cache_unfinished_req(self, req: Req):
+    def cache_unfinished_req(self, req: Req, chunked=False):
         """Cache request when it is unfinished."""
         if self.disable:
             return

sglang/srt/mem_cache/memory_pool.py

@@ -918,6 +918,7 @@ class AscendMLAPagedTokenToKVPool(MLATokenToKVPool):
                     layer_num,
                     self.size // self.page_size + 1,
                     self.page_size,
+                    1,
                     self.kv_lora_rank,
                 ),
                 dtype=self.store_dtype,
@@ -928,6 +929,7 @@ class AscendMLAPagedTokenToKVPool(MLATokenToKVPool):
                     layer_num,
                     self.size // self.page_size + 1,
                     self.page_size,
+                    1,
                     self.qk_rope_head_dim,
                 ),
                 dtype=self.store_dtype,
@@ -1000,9 +1002,11 @@ class AscendMLAPagedTokenToKVPool(MLATokenToKVPool):
         layer_id = layer.layer_id
         if cache_k.dtype != self.dtype:
             cache_k = cache_k.to(self.dtype)
+            cache_v = cache_v.to(self.dtype)
 
         if self.store_dtype != self.dtype:
             cache_k = cache_k.view(self.store_dtype)
+            cache_v = cache_v.view(self.store_dtype)
 
         if cache_v is None:
             cache_k, cache_v = cache_k.split(

sglang/srt/mem_cache/memory_pool_host.py

@@ -7,7 +7,6 @@ from functools import wraps
 import psutil
 import torch
 
-from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.mem_cache.memory_pool import KVCache, MHATokenToKVPool, MLATokenToKVPool
 from sglang.srt.utils import is_npu
 
@@ -464,11 +463,11 @@ class MHATokenToKVPoolHost(HostKVCache):
         else:
             raise ValueError(f"Unsupported layout: {self.layout}")
 
-    def get_buffer_meta(self, keys, indices):
-        local_rank = get_tensor_model_parallel_rank()
+    def get_buffer_meta(self, keys, indices, local_rank):
         ptr_list = []
         key_list = []
         kv_buffer_data_ptr = self.kv_buffer.data_ptr()
+        indices = indices.tolist()
         v_offset = (
             self.layer_num
             * self.size
@@ -704,10 +703,11 @@ class MLATokenToKVPoolHost(HostKVCache):
         else:
             raise ValueError(f"Unsupported layout: {self.layout}")
 
-    def get_buffer_meta(self, keys, indices):
+    def get_buffer_meta(self, keys, indices, local_rank):
         ptr_list = []
         key_list = []
         kv_buffer_data_ptr = self.kv_buffer.data_ptr()
+        indices = indices.tolist()
         for index in range(0, len(indices), self.page_size):
             k_ptr = (
                 kv_buffer_data_ptr

sglang/srt/mem_cache/radix_cache.py

@@ -62,7 +62,6 @@ class TreeNode:
         self.host_value: Optional[torch.Tensor] = None
         # store hash values of each pages
         self.hash_value: Optional[List[str]] = None
-        self.backuped_storage = False
 
         self.id = TreeNode.counter if id is None else id
         TreeNode.counter += 1
@@ -195,7 +194,7 @@ class RadixCache(BasePrefixCache):
             last_host_node=last_node,
         )
 
-    def insert(self, key: List, value=None):
+    def insert(self, key: List, value=None, chunked=False):
         if self.disable:
             return 0
 
@@ -240,7 +239,7 @@ class RadixCache(BasePrefixCache):
         self.req_to_token_pool.free(req.req_pool_idx)
         self.dec_lock_ref(req.last_node)
 
-    def cache_unfinished_req(self, req: Req):
+    def cache_unfinished_req(self, req: Req, chunked=False):
         """Cache request when it is unfinished."""
         if self.disable:
             return
@@ -261,7 +260,9 @@ class RadixCache(BasePrefixCache):
         page_aligned_token_ids = token_ids[:page_aligned_len]
 
         # Radix Cache takes one ref in memory pool
-        new_prefix_len = self.insert(page_aligned_token_ids, page_aligned_kv_indices)
+        new_prefix_len = self.insert(
+            page_aligned_token_ids, page_aligned_kv_indices, chunked=chunked
+        )
         self.token_to_kv_pool_allocator.free(
             kv_indices[len(req.prefix_indices) : new_prefix_len]
         )

sglang/srt/mem_cache/radix_cache_cpp.py

@@ -181,7 +181,7 @@ class RadixCacheCpp(BasePrefixCache):
         self.dec_lock_ref(req.last_node)
         self.req_to_token_pool.free(req.req_pool_idx)
 
-    def cache_unfinished_req(self, req: Req):
+    def cache_unfinished_req(self, req: Req, chunked=False):
         """Cache request when it is unfinished."""
         assert req.req_pool_idx is not None
         token_ids = req.fill_ids

sglang/srt/mem_cache/storage/hf3fs/mini_3fs_metadata_server.py

@@ -4,10 +4,12 @@ import json
 import logging
 import threading
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, OrderedDict, Tuple
 
+import orjson
 import requests
-from fastapi import FastAPI, HTTPException, Request, status
+from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi.responses import ORJSONResponse
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
@@ -24,10 +26,10 @@ class RankMetadata:
     """Holds all metadata for a single rank."""
 
     def __init__(self, num_pages: int):
-        self.lock = threading.RLock()
+        self.lock = threading.Lock()
         self.num_pages = num_pages
         self.free_pages: List[int] = list(range(num_pages))
-        self.key_to_index: Dict[str, int] = {}
+        self.key_to_index: OrderedDict[str, int] = OrderedDict()
         # Todo: Support multi files for HF3FS
 
     def exists_keys(self, keys: List[str]) -> List[bool]:
@@ -46,16 +48,18 @@ class RankMetadata:
             for i, (key, prefix_key) in enumerate(keys):
                 if key in self.key_to_index:
                     results[i] = (True, self.key_to_index[key])
+                    self.key_to_index.move_to_end(key)
                 else:
                     new_keys_to_process.append((i, key, prefix_key))
 
             # Todo: Implementing data eviction logic after HiCache supports prefix information pass-through
             for i, key, prefix_key in new_keys_to_process:
                 if len(self.free_pages) > 0:
-                    page_idx = self.free_pages.pop()
-                    results[i] = (False, page_idx)
+                    page_index = self.free_pages.pop()
                 else:
-                    results[i] = (False, -1)
+                    page_index = self.key_to_index.popitem(last=False)[1]
+
+                results[i] = (False, page_index)
 
             return results
 
@@ -68,6 +72,7 @@ class RankMetadata:
         with self.lock:
             for key, page_index in written_keys_to_confirm:
                 self.key_to_index[key] = page_index
+                self.key_to_index.move_to_end(key)
 
             for page_index in pages_to_release:
                 if page_index not in self.free_pages:
@@ -94,7 +99,14 @@ class RankMetadata:
     def get_page_indices(self, keys: List[str]) -> List[Optional[int]]:
         """Get page indices for keys."""
         with self.lock:
-            return [self.key_to_index.get(key) for key in keys]
+            results = []
+            for key in keys:
+                if key in self.key_to_index:
+                    results.append(self.key_to_index[key])
+                    self.key_to_index.move_to_end(key)
+                else:
+                    results.append(None)
+            return results
 
 
 class GlobalMetadataState:
@@ -182,7 +194,8 @@ class Hf3fsMetadataServer:
 
     def __init__(self, persistence_path: Optional[str] = None, save_interval: int = 60):
         self.state = GlobalMetadataState(persistence_path, save_interval)
-        self.app = FastAPI()
+        self.app = FastAPI(default_response_class=ORJSONResponse)
+
         self._setup_routes()
 
     def _setup_routes(self):
@@ -199,17 +212,25 @@ class Hf3fsMetadataServer:
 
     def get_rank_metadata(self, rank: int) -> RankMetadata:
         """Get rank metadata with proper error handling."""
-        with self.state.global_lock:
-            if rank not in self.state.ranks:
-                raise HTTPException(
-                    status_code=404,
-                    detail=f"Rank {rank} not initialized. Please call /{{rank}}/initialize first.",
-                )
-            return self.state.ranks[rank]
+        if rank not in self.state.ranks:
+            raise HTTPException(
+                status_code=404,
+                detail=f"Rank {rank} not initialized. Please call /{rank}/initialize first.",
+            )
+        return self.state.ranks[rank]
+
+    async def _read_json(self, request: Request) -> dict:
+        """Parse request JSON using orjson if available."""
+        body = await request.body()
+        return orjson.loads(body)
+
+    def _json_response(self, content: dict):
+        """Return ORJSONResponse when available to bypass jsonable_encoder."""
+        return ORJSONResponse(content)
 
     async def initialize(self, rank: int, request: Request):
         """Initialize a rank with specified number of pages."""
-        data = await request.json()
+        data = await self._read_json(request)
         num_pages = data["num_pages"]
         with self.state.global_lock:
             if rank in self.state.ranks:
@@ -223,57 +244,55 @@ class Hf3fsMetadataServer:
             else:
                 logging.info(f"Initializing new Rank {rank} with {num_pages} pages.")
                 self.state.ranks[rank] = RankMetadata(num_pages)
-        return {"message": f"Rank {rank} is ready."}
+        return Response(status_code=204)
 
     async def exists(self, rank: int, request: Request):
         """Check if keys exist in metadata."""
-        data = await request.json()
+        data = await self._read_json(request)
         keys = data["keys"]
         metadata = self.get_rank_metadata(rank)
         results = metadata.exists_keys(keys)
-        return {"exists": results}
+        return self._json_response({"exists": results})
 
     async def reserve_and_allocate_page_indices(self, rank: int, request: Request):
         """Reserve and allocate page indices for keys."""
-        data = await request.json()
+        data = await self._read_json(request)
         metadata = self.get_rank_metadata(rank)
         keys = data["keys"]
         results = metadata.reserve_and_allocate_page_indices(keys)
-        return {"indices": results}
+        return self._json_response({"indices": results})
 
     async def confirm_write(self, rank: int, request: Request):
         """Confirm write operations and release pages."""
-        data = await request.json()
+        data = await self._read_json(request)
         metadata = self.get_rank_metadata(rank)
         success_written_keys = data.get("written_keys_to_confirm", [])
         released_pages = data.get("pages_to_release", [])
 
         metadata.confirm_write(success_written_keys, released_pages)
 
-        return {
-            "message": f"Rank {rank}: Write confirmed for {len(success_written_keys)} keys. {len(released_pages)} pages released."
-        }
+        return Response(status_code=204)
 
     async def delete_keys(self, rank: int, request: Request):
         """Delete keys from metadata."""
-        data = await request.json()
+        data = await self._read_json(request)
         metadata = self.get_rank_metadata(rank)
         count = metadata.delete_keys(data["keys"])
-        return {"message": f"Rank {rank}: {count} keys deleted."}
+        return Response(status_code=204)
 
     async def clear(self, rank: int):
         """Clear all metadata for a rank."""
         metadata = self.get_rank_metadata(rank)
         metadata.clear_all()
-        return {"message": f"Rank {rank}: Metadata cleared."}
+        return Response(status_code=204)
 
     async def get_page_indices(self, rank: int, request: Request):
         """Get page indices for keys."""
-        data = await request.json()
+        data = await self._read_json(request)
         metadata = self.get_rank_metadata(rank)
         keys = data["keys"]
         results = metadata.get_page_indices(keys)
-        return {"indices": results}
+        return self._json_response({"indices": results})
 
     def run(self, host: str = "0.0.0.0", port: int = 18000):
         """Run the metadata server."""
@@ -309,14 +328,22 @@ class Hf3fsGlobalMetadataClient(Hf3fsMetadataInterface):
             status_forcelist=[500, 502, 503, 504],
             allowed_methods=["GET", "POST"],
         )
-        adapter = HTTPAdapter(max_retries=retry_strategy)
+        adapter = HTTPAdapter(
+            max_retries=retry_strategy, pool_connections=256, pool_maxsize=256
+        )
         self._session.mount("http://", adapter)
 
     def _post(self, endpoint: str, json_data: dict) -> dict:
         try:
-            response = self._session.post(f"{self.base_url}/{endpoint}", json=json_data)
+            url = f"{self.base_url}/{endpoint}"
+            headers = {"Content-Type": "application/json"}
+            payload = orjson.dumps(json_data)  # type: ignore[union-attr]
+            response = self._session.post(url, data=payload, headers=headers)
             response.raise_for_status()
-            return response.json()
+
+            if response.status_code == 204 or not response.content:
+                return {}
+            return orjson.loads(response.content)  # type: ignore[union-attr]
         except requests.exceptions.RequestException as e:
             logging.error(f"Failed to POST to {endpoint} after retries: {e}")
             raise RuntimeError(f"Failed to connect to metadata server: {e}") from e