sglang 0.5.1.post2__py3-none-any.whl → 0.5.2__py3-none-any.whl

sglang/srt/mem_cache/memory_pool_host.py

@@ -3,16 +3,17 @@ import logging
 import threading
 from enum import IntEnum
 from functools import wraps
+from typing import Optional
 
 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
+from sglang.srt.utils import is_npu, is_xpu
 
 _is_npu = is_npu()
-if not _is_npu:
+_is_xpu = is_xpu()
+if not (_is_npu or _is_xpu):
     from sgl_kernel.kvcacheio import (
         transfer_kv_all_layer,
         transfer_kv_all_layer_lf_pf,
@@ -169,7 +170,7 @@ class HostKVCache(abc.ABC):
         return len(self.free_slots)
 
     @synchronized()
-    def alloc(self, need_size: int) -> torch.Tensor:
+    def alloc(self, need_size: int) -> Optional[torch.Tensor]:
         assert (
             need_size % self.page_size == 0
         ), "The requested size should be a multiple of the page size."
@@ -464,10 +465,11 @@ class MHATokenToKVPoolHost(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()
         v_offset = (
             self.layer_num
             * self.size
@@ -488,8 +490,8 @@ class MHATokenToKVPoolHost(HostKVCache):
             ptr_list.append(k_ptr)
             ptr_list.append(v_ptr)
             key_ = keys[index // self.page_size]
-            key_list.append(f"{key_}_{get_tensor_model_parallel_rank()}_k")
-            key_list.append(f"{key_}_{get_tensor_model_parallel_rank()}_v")
+            key_list.append(f"{key_}_{local_rank}_k")
+            key_list.append(f"{key_}_{local_rank}_v")
         element_size = (
             self.layer_num
             * self.dtype.itemsize
@@ -500,20 +502,23 @@ class MHATokenToKVPoolHost(HostKVCache):
         element_size_list = [element_size] * len(key_list)
         return key_list, ptr_list, element_size_list
 
-    def get_buffer_with_hash(self, keys, indices):
+    def get_buffer_with_hash(self, keys, indices=None):
         assert self.layout == "page_first"
-        assert len(keys) == (len(indices) // self.page_size)
+        assert indices is None or (len(keys) == (len(indices) // self.page_size))
 
         key_list = []
         buf_list = []
 
-        for key, i in zip(keys, range(0, len(indices), self.page_size)):
+        for i in range(len(keys)):
+            key = keys[i]
             key_list.append(f"{key}-k")
-            buf_list.append(self.k_buffer[i : i + self.page_size])
             key_list.append(f"{key}-v")
-            buf_list.append(self.v_buffer[i : i + self.page_size])
+            if indices is not None:
+                index = indices[i * self.page_size]
+                buf_list.append(self.k_buffer[index : index + self.page_size])
+                buf_list.append(self.v_buffer[index : index + self.page_size])
 
-        return key_list, buf_list
+        return key_list, buf_list, 2
 
 
 class MLATokenToKVPoolHost(HostKVCache):
@@ -703,10 +708,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
@@ -727,13 +733,15 @@ class MLATokenToKVPoolHost(HostKVCache):
         element_size_list = [element_size] * len(key_list)
         return key_list, ptr_list, element_size_list
 
-    def get_buffer_with_hash(self, keys, indices):
+    def get_buffer_with_hash(self, keys, indices=None):
         assert self.layout == "page_first"
-        assert len(keys) == (len(indices) // self.page_size)
+        assert indices is None or (len(keys) == (len(indices) // self.page_size))
 
         buf_list = []
 
-        for i in range(0, len(indices), self.page_size):
-            buf_list.append(self.kv_buffer[i : i + self.page_size])
+        if indices is not None:
+            for i in range(len(keys)):
+                index = indices[i * self.page_size]
+                buf_list.append(self.kv_buffer[index : index + self.page_size])
 
-        return keys, buf_list
+        return keys, buf_list, 1

sglang/srt/mem_cache/radix_cache.py

@@ -53,8 +53,6 @@ class TreeNode:
         self.last_access_time = time.monotonic()
 
         self.hit_count = 0
-        # indicating the node is loading KV cache from host
-        self.loading = False
         # indicating the node is locked to protect from eviction
         # incremented when the node is referenced by a storage operation
         self.host_ref_counter = 0
@@ -62,7 +60,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
@@ -152,6 +149,7 @@ class RadixCache(BasePrefixCache):
         self.root_node = TreeNode()
         self.root_node.key = []
         self.root_node.value = []
+        self.root_node.host_value = []
         self.root_node.lock_ref = 1
         self.evictable_size_ = 0
         self.protected_size_ = 0
@@ -194,7 +192,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
 
@@ -239,7 +237,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
@@ -260,7 +258,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/hf3fs_client.py

@@ -0,0 +1,164 @@
+import logging
+import os
+import threading
+from abc import ABC, abstractmethod
+from typing import List
+
+import torch
+
+
+class Hf3fsClient(ABC):
+    """Abstract interface for HF3FS clients."""
+
+    @abstractmethod
+    def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
+        """Initialize the HF3FS client.
+
+        Args:
+            path: File path for storage
+            size: Total size of storage file
+            bytes_per_page: Bytes per page
+            entries: Number of entries for batch operations
+        """
+        pass
+
+    @abstractmethod
+    def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        """Batch read from storage."""
+        pass
+
+    @abstractmethod
+    def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        """Batch write to storage."""
+        pass
+
+    @abstractmethod
+    def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
+        """Validate batch operation parameters."""
+        pass
+
+    @abstractmethod
+    def get_size(self) -> int:
+        """Get total storage size."""
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the client and cleanup resources."""
+        pass
+
+    @abstractmethod
+    def flush(self) -> None:
+        """Flush data to disk."""
+        pass
+
+
+logger = logging.getLogger(__name__)
+
+
+class Hf3fsMockClient(Hf3fsClient):
+    """Mock implementation of Hf3fsClient for CI testing purposes."""
+
+    def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
+        """Initialize mock HF3FS client."""
+        self.path = path
+        self.size = size
+        self.bytes_per_page = bytes_per_page
+        self.entries = entries
+
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(self.path), exist_ok=True)
+
+        # Create and initialize the file
+        self.file = os.open(self.path, os.O_RDWR | os.O_CREAT)
+        os.ftruncate(self.file, size)
+
+        logger.info(
+            f"Hf3fsMockClient initialized: path={path}, size={size}, "
+            f"bytes_per_page={bytes_per_page}, entries={entries}"
+        )
+
+    def batch_read(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        """Batch read from mock storage."""
+        self.check(offsets, tensors)
+
+        results = []
+
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+
+            try:
+                os.lseek(self.file, offset, os.SEEK_SET)
+                bytes_read = os.read(self.file, size)
+
+                if len(bytes_read) == size:
+                    # Convert bytes to tensor and copy to target
+                    bytes_tensor = torch.frombuffer(bytes_read, dtype=torch.uint8)
+                    typed_tensor = bytes_tensor.view(tensor.dtype).view(tensor.shape)
+                    tensor.copy_(typed_tensor)
+                    results.append(size)
+                else:
+                    logger.warning(
+                        f"Short read: expected {size}, got {len(bytes_read)}"
+                    )
+                    results.append(len(bytes_read))
+
+            except Exception as e:
+                logger.error(f"Error reading from offset {offset}: {e}")
+                results.append(0)
+
+        return results
+
+    def batch_write(self, offsets: List[int], tensors: List[torch.Tensor]) -> List[int]:
+        """Batch write to mock storage."""
+        self.check(offsets, tensors)
+
+        results = []
+
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+
+            try:
+                # Convert tensor to bytes and write directly to file
+                tensor_bytes = tensor.contiguous().view(torch.uint8).flatten()
+                data = tensor_bytes.numpy().tobytes()
+
+                os.lseek(self.file, offset, os.SEEK_SET)
+                bytes_written = os.write(self.file, data)
+
+                if bytes_written == size:
+                    results.append(size)
+                else:
+                    logger.warning(f"Short write: expected {size}, got {bytes_written}")
+                    results.append(bytes_written)
+
+            except Exception as e:
+                logger.error(f"Error writing to offset {offset}: {e}")
+                results.append(0)
+
+        return results
+
+    def check(self, offsets: List[int], tensors: List[torch.Tensor]) -> None:
+        """Validate batch operation parameters."""
+        pass
+
+    def get_size(self) -> int:
+        """Get total storage size."""
+        return self.size
+
+    def close(self) -> None:
+        """Close the mock client and cleanup resources."""
+        try:
+            if hasattr(self, "file") and self.file >= 0:
+                os.close(self.file)
+                self.file = -1  # Mark as closed
+            logger.info(f"MockHf3fsClient closed: {self.path}")
+        except Exception as e:
+            logger.error(f"Error closing MockHf3fsClient: {e}")
+
+    def flush(self) -> None:
+        """Flush data to disk."""
+        try:
+            os.fsync(self.file)
+        except Exception as e:
+            logger.error(f"Error flushing MockHf3fsClient: {e}")

sglang/srt/mem_cache/storage/hf3fs/{client_hf3fs.py → hf3fs_usrbio_client.py}

@@ -9,6 +9,8 @@ from typing import List
 import torch
 from torch.utils.cpp_extension import load
 
+from sglang.srt.mem_cache.storage.hf3fs.hf3fs_client import Hf3fsClient
+
 root = Path(__file__).parent.resolve()
 hf3fs_utils = load(name="hf3fs_utils", sources=[f"{root}/hf3fs_utils.cpp"])
 
@@ -51,7 +53,9 @@ def wsynchronized():
     return _decorator
 
 
-class Hf3fsClient:
+class Hf3fsUsrBioClient(Hf3fsClient):
+    """HF3FS client implementation using usrbio."""
+
     def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
         if not HF3FS_AVAILABLE:
             raise ImportError(

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