sglang 0.4.10__py3-none-any.whl → 0.4.10.post1__py3-none-any.whl

sglang/srt/mem_cache/memory_pool_host.py

@@ -8,6 +8,21 @@ import psutil
 import torch
 
 from sglang.srt.mem_cache.memory_pool import KVCache, MHATokenToKVPool, MLATokenToKVPool
+from sglang.srt.utils import is_npu
+
+_is_npu = is_npu()
+if not _is_npu:
+    from sgl_kernel.kvcacheio import (
+        transfer_kv_all_layer,
+        transfer_kv_all_layer_lf_pf,
+        transfer_kv_all_layer_mla,
+        transfer_kv_all_layer_mla_lf_pf,
+        transfer_kv_direct,
+        transfer_kv_per_layer,
+        transfer_kv_per_layer_mla,
+        transfer_kv_per_layer_mla_pf_lf,
+        transfer_kv_per_layer_pf_lf,
+    )
 
 logger = logging.getLogger(__name__)
 
@@ -42,15 +57,18 @@ class HostKVCache(abc.ABC):
         device_pool: KVCache,
         host_to_device_ratio: float,
         host_size: int,
+        page_size: int,
+        layout: str,
         pin_memory: bool,
         device: str,
-        page_size: int,
     ):
         self.device_pool = device_pool
-        self.dtype = device_pool.store_dtype
+        self.page_size = page_size
+        self.layout = layout
         self.pin_memory = pin_memory
         self.device = device
-        self.page_size = page_size
+
+        self.dtype = device_pool.store_dtype
         self.size_per_token = self.get_size_per_token()
         if host_size > 0:
             self.size = int(host_size * 1e9 // self.size_per_token)
@@ -98,6 +116,24 @@ class HostKVCache(abc.ABC):
     def init_kv_buffer(self):
         raise NotImplementedError()
 
+    @abc.abstractmethod
+    def load_to_device_per_layer(
+        self, device_pool, host_indices, device_indices, layer_id, io_backend
+    ) -> None:
+        """
+        Load KV data from the host memory pool to the device memory pool for a specific layer.
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def backup_from_device_all_layer(
+        self, device_pool, host_indices, device_indices, io_backend
+    ) -> None:
+        """
+        Backup KV data from the device memory pool to the host memory pool for all layers.
+        """
+        raise NotImplementedError()
+
     @abc.abstractmethod
     def get_flat_data_page(self, index) -> torch.Tensor:
         """
@@ -105,6 +141,14 @@ class HostKVCache(abc.ABC):
         """
         raise NotImplementedError()
 
+    @abc.abstractmethod
+    def get_dummy_flat_data_page(self) -> torch.Tensor:
+        """
+        Get a dummy flat data page from the host memory pool.
+        This is used for prefetching or initializing empty pages.
+        """
+        raise NotImplementedError()
+
     @abc.abstractmethod
     def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
         """
@@ -230,11 +274,30 @@ class MHATokenToKVPoolHost(HostKVCache):
         host_to_device_ratio: float,
         host_size: int,
         page_size: int,
+        layout: str,
         pin_memory: bool = True,
         device: str = "cpu",
     ):
         super().__init__(
-            device_pool, host_to_device_ratio, host_size, pin_memory, device, page_size
+            device_pool,
+            host_to_device_ratio,
+            host_size,
+            page_size,
+            layout,
+            pin_memory,
+            device,
+        )
+        self.k_data_refs = [self.k_buffer[i] for i in range(self.layer_num)]
+        self.v_data_refs = [self.v_buffer[i] for i in range(self.layer_num)]
+        self.k_data_ptrs = torch.tensor(
+            [x.data_ptr() for x in self.k_data_refs],
+            dtype=torch.uint64,
+            device=self.device_pool.device,
+        )
+        self.v_data_ptrs = torch.tensor(
+            [x.data_ptr() for x in self.v_data_refs],
+            dtype=torch.uint64,
+            device=self.device_pool.device,
         )
 
     def get_size_per_token(self):
@@ -245,25 +308,156 @@ class MHATokenToKVPoolHost(HostKVCache):
         return self.head_dim * self.head_num * self.layer_num * self.dtype.itemsize * 2
 
     def init_kv_buffer(self):
+        if self.layout == "layer_first":
+            dims = (2, self.layer_num, self.size, self.head_num, self.head_dim)
+        elif self.layout == "page_first":
+            dims = (2, self.size, self.layer_num, self.head_num, self.head_dim)
+        else:
+            raise ValueError(f"Unsupported layout: {self.layout}")
+        self.token_stride_size = self.head_num * self.head_dim * self.dtype.itemsize
+        self.layout_dim = self.token_stride_size * self.layer_num
         return torch.empty(
-            (2, self.layer_num, self.size, self.head_num, self.head_dim),
+            dims,
             dtype=self.dtype,
             device=self.device,
             pin_memory=self.pin_memory,
         )
 
-    # todo, page first memory layout
+    @property
+    def k_buffer(self):
+        return self.kv_buffer[0]
+
+    @property
+    def v_buffer(self):
+        return self.kv_buffer[1]
+
+    def load_to_device_per_layer(
+        self,
+        device_pool,
+        host_indices,
+        device_indices,
+        layer_id,
+        io_backend,
+    ):
+        if io_backend == "kernel":
+            if self.layout == "layer_first":
+                transfer_kv_per_layer(
+                    src_k=self.k_buffer[layer_id],
+                    dst_k=device_pool.k_buffer[layer_id],
+                    src_v=self.v_buffer[layer_id],
+                    dst_v=device_pool.v_buffer[layer_id],
+                    src_indices=host_indices,
+                    dst_indices=device_indices,
+                    item_size=self.token_stride_size,
+                )
+            elif self.layout == "page_first":
+                transfer_kv_per_layer_pf_lf(
+                    src_k=self.k_buffer,
+                    dst_k=device_pool.k_buffer[layer_id],
+                    src_v=self.v_buffer,
+                    dst_v=device_pool.v_buffer[layer_id],
+                    src_indices=host_indices,
+                    dst_indices=device_indices,
+                    item_size=self.token_stride_size,
+                    src_layout_dim=self.layout_dim,
+                )
+            else:
+                raise ValueError(f"Unsupported layout: {self.layout}")
+        elif io_backend == "direct":
+            assert (
+                self.layout == "layer_first"
+            ), f"Direct IO backend only supports layer_first layout."
+            transfer_kv_direct(
+                src_layers=[self.k_buffer[layer_id], self.v_buffer[layer_id]],
+                dst_layers=[
+                    device_pool.k_buffer[layer_id],
+                    device_pool.v_buffer[layer_id],
+                ],
+                src_indices=host_indices,
+                dst_indices=device_indices,
+                page_size=self.page_size,
+            )
+        else:
+            raise ValueError(f"Unsupported IO backend: {io_backend}")
+
+    def backup_from_device_all_layer(
+        self, device_pool, host_indices, device_indices, io_backend
+    ):
+        if io_backend == "kernel":
+            if self.layout == "layer_first":
+                transfer_kv_all_layer(
+                    src_k_layers=device_pool.k_data_ptrs,
+                    dst_k_layers=self.k_data_ptrs,
+                    src_v_layers=device_pool.v_data_ptrs,
+                    dst_v_layers=self.v_data_ptrs,
+                    src_indices=device_indices,
+                    dst_indices=host_indices,
+                    item_size=self.token_stride_size,
+                    num_layers=self.layer_num,
+                )
+            elif self.layout == "page_first":
+                transfer_kv_all_layer_lf_pf(
+                    src_k_layers=device_pool.k_data_ptrs,
+                    dst_k=self.k_buffer,
+                    src_v_layers=device_pool.v_data_ptrs,
+                    dst_v=self.v_buffer,
+                    src_indices=device_indices,
+                    dst_indices=host_indices,
+                    item_size=self.token_stride_size,
+                    dst_layout_dim=self.layout_dim,
+                    num_layers=self.layer_num,
+                )
+            else:
+                raise ValueError(f"Unsupported layout: {self.layout}")
+        elif io_backend == "direct":
+            assert (
+                self.layout == "layer_first"
+            ), f"Direct IO backend only supports layer_first layout."
+            transfer_kv_direct(
+                src_layers=device_pool.k_buffer + device_pool.v_buffer,
+                dst_layers=self.k_data_refs + self.v_data_refs,
+                src_indices=device_indices,
+                dst_indices=host_indices,
+                page_size=self.page_size,
+            )
+        else:
+            raise ValueError(f"Unsupported IO backend: {io_backend}")
+
     def get_flat_data_page(self, index) -> torch.Tensor:
-        return self.kv_buffer[:, :, index : index + self.page_size, :, :].flatten()
+        if self.layout == "layer_first":
+            return self.kv_buffer[:, :, index : index + self.page_size, :, :].flatten()
+        elif self.layout == "page_first":
+            return self.kv_buffer[:, index : index + self.page_size, :, :, :].flatten()
+        else:
+            raise ValueError(f"Unsupported layout: {self.layout}")
+
+    def get_dummy_flat_data_page(self) -> torch.Tensor:
+        return torch.zeros(
+            (2, self.layer_num, self.page_size, self.head_num, self.head_dim),
+            dtype=self.dtype,
+            device=self.device,
+            pin_memory=self.pin_memory,
+        ).flatten()
 
     def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
-        self.kv_buffer[:, :, index : index + self.page_size, :, :] = data_page.reshape(
-            2,
-            self.layer_num,
-            self.page_size,
-            self.head_num,
-            self.head_dim,
-        )
+        if self.layout == "layer_first":
+            self.kv_buffer[:, :, index : index + self.page_size, :, :] = (
+                data_page.reshape(
+                    2,
+                    self.layer_num,
+                    self.page_size,
+                    self.head_num,
+                    self.head_dim,
+                )
+            )
+        elif self.layout == "page_first":
+            self.kv_buffer[:, index : index + self.page_size, :, :, :] = (
+                data_page.reshape(
+                    2, self.page_size, self.layer_num, self.head_num, self.head_dim
+                )
+            )
+        else:
+            raise ValueError(f"Unsupported layout: {self.layout}")
 
     def get_buffer_meta(self, keys, indices):
         ptr_list = []
@@ -302,14 +496,6 @@ class MHATokenToKVPoolHost(HostKVCache):
         element_size_list = [element_size] * len(key_list)
         return key_list, ptr_list, element_size_list
 
-    @property
-    def k_buffer(self):
-        return self.kv_buffer[0]
-
-    @property
-    def v_buffer(self):
-        return self.kv_buffer[1]
-
 
 class MLATokenToKVPoolHost(HostKVCache):
     device_pool: MLATokenToKVPool
@@ -320,11 +506,24 @@ class MLATokenToKVPoolHost(HostKVCache):
         host_to_device_ratio: float,
         host_size: int,
         page_size: int,
+        layout: str,
         pin_memory: bool = True,
         device: str = "cpu",
     ):
         super().__init__(
-            device_pool, host_to_device_ratio, host_size, pin_memory, device, page_size
+            device_pool,
+            host_to_device_ratio,
+            host_size,
+            page_size,
+            layout,
+            pin_memory,
+            device,
+        )
+        self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)]
+        self.data_ptrs = torch.tensor(
+            [x.data_ptr() for x in self.data_refs],
+            dtype=torch.uint64,
+            device=self.device_pool.device,
         )
 
     def get_size_per_token(self):
@@ -340,28 +539,146 @@ class MLATokenToKVPoolHost(HostKVCache):
         )
 
     def init_kv_buffer(self):
-        return torch.empty(
-            (
+        if self.layout == "layer_first":
+            dims = (
                 self.layer_num,
                 self.size,
                 1,
                 self.kv_lora_rank + self.qk_rope_head_dim,
-            ),
+            )
+        elif self.layout == "page_first":
+            dims = (
+                self.size,
+                self.layer_num,
+                1,
+                self.kv_lora_rank + self.qk_rope_head_dim,
+            )
+        else:
+            raise ValueError(f"Unsupported layout: {self.layout}")
+        self.token_stride_size = (
+            self.kv_lora_rank + self.qk_rope_head_dim
+        ) * self.dtype.itemsize
+        self.layout_dim = self.token_stride_size * self.layer_num
+
+        return torch.empty(
+            dims,
             dtype=self.dtype,
             device=self.device,
             pin_memory=self.pin_memory,
         )
 
+    def load_to_device_per_layer(
+        self, device_pool, host_indices, device_indices, layer_id, io_backend
+    ):
+        if io_backend == "kernel":
+            if self.layout == "layer_first":
+                transfer_kv_per_layer_mla(
+                    src=self.kv_buffer[layer_id],
+                    dst=device_pool.kv_buffer[layer_id],
+                    src_indices=host_indices,
+                    dst_indices=device_indices,
+                    item_size=self.token_stride_size,
+                )
+            elif self.layout == "page_first":
+                transfer_kv_per_layer_mla_pf_lf(
+                    src=self.kv_buffer,
+                    dst=device_pool.kv_buffer[layer_id],
+                    src_indices=host_indices,
+                    dst_indices=device_indices,
+                    item_size=self.token_stride_size,
+                    src_layout_dim=self.layout_dim,
+                )
+            else:
+                raise ValueError(f"Unsupported layout: {self.layout}")
+        elif io_backend == "direct":
+            assert (
+                self.layout == "layer_first"
+            ), f"Direct IO backend only supports layer_first layout."
+            transfer_kv_direct(
+                src_layers=[self.kv_buffer[layer_id]],
+                dst_layers=[device_pool.kv_buffer[layer_id]],
+                src_indices=host_indices,
+                dst_indices=device_indices,
+                page_size=self.page_size,
+            )
+
+    def backup_from_device_all_layer(
+        self, device_pool, host_indices, device_indices, io_backend
+    ):
+        if io_backend == "kernel":
+            if self.layout == "layer_first":
+                transfer_kv_all_layer_mla(
+                    src_layers=device_pool.data_ptrs,
+                    dst_layers=self.data_ptrs,
+                    src_indices=device_indices,
+                    dst_indices=host_indices,
+                    item_size=self.token_stride_size,
+                    num_layers=self.layer_num,
+                )
+            elif self.layout == "page_first":
+                transfer_kv_all_layer_mla_lf_pf(
+                    src_layers=device_pool.data_ptrs,
+                    dst_k=self.kv_buffer,
+                    src_indices=device_indices,
+                    dst_indices=host_indices,
+                    item_size=self.token_stride_size,
+                    dst_layout_dim=self.layout_dim,
+                    num_layers=self.layer_num,
+                )
+            else:
+                raise ValueError(f"Unsupported layout: {self.layout}")
+        elif io_backend == "direct":
+            assert (
+                self.layout == "layer_first"
+            ), f"Direct IO backend only supports layer_first layout."
+            transfer_kv_direct(
+                src_layers=device_pool.kv_buffer,
+                dst_layers=self.data_refs,
+                src_indices=device_indices,
+                dst_indices=host_indices,
+                page_size=self.page_size,
+            )
+        else:
+            raise ValueError(f"Unsupported IO backend: {io_backend}")
+
     def get_flat_data_page(self, index) -> torch.Tensor:
-        return self.kv_buffer[:, index : index + self.page_size, :, :].flatten()
+        if self.layout == "layer_first":
+            return self.kv_buffer[:, index : index + self.page_size, :, :].flatten()
+        elif self.layout == "page_first":
+            return self.kv_buffer[index : index + self.page_size, :, :, :].flatten()
+        else:
+            raise ValueError(f"Unsupported layout: {self.layout}")
+
+    def get_dummy_flat_data_page(self) -> torch.Tensor:
+        return torch.zeros(
+            (
+                self.layer_num,
+                self.page_size,
+                1,
+                self.kv_lora_rank + self.qk_rope_head_dim,
+            ),
+            dtype=self.dtype,
+            device=self.device,
+            pin_memory=self.pin_memory,
+        ).flatten()
 
     def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
-        self.kv_buffer[:, index : index + self.page_size, :, :] = data_page.reshape(
-            self.layer_num,
-            self.page_size,
-            1,
-            self.kv_lora_rank + self.qk_rope_head_dim,
-        )
+        if self.layout == "layer_first":
+            self.kv_buffer[:, index : index + self.page_size, :, :] = data_page.reshape(
+                self.layer_num,
+                self.page_size,
+                1,
+                self.kv_lora_rank + self.qk_rope_head_dim,
+            )
+        elif self.layout == "page_first":
+            self.kv_buffer[index : index + self.page_size, :, :, :] = data_page.reshape(
+                self.page_size,
+                self.layer_num,
+                1,
+                self.kv_lora_rank + self.qk_rope_head_dim,
+            )
+        else:
+            raise ValueError(f"Unsupported layout: {self.layout}")
 
     def get_buffer_meta(self, keys, indices):
         ptr_list = []

sglang/srt/mem_cache/nixl/hicache_nixl.py

@@ -0,0 +1,163 @@
+import hashlib
+import logging
+import os
+import time
+import uuid
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
+
+from .nixl_utils import NixlBackendSelection, NixlFileManager, NixlRegistration
+
+try:
+    from nixl._api import nixl_agent, nixl_agent_config
+except ImportError as e:
+    raise ImportError(
+        "Please install NIXL by following the instructions at "
+        "https://github.com/ai-dynamo/nixl/blob/main/README.md "
+        "to use HiCacheNixl storage backend."
+    ) from e
+
+logger = logging.getLogger(__name__)
+
+
+class HiCacheNixl(HiCacheStorage):
+    """HiCacheNixl provides high-performance storage using NIXL plugins."""
+
+    def __init__(self, file_path: str = "/tmp/hicache_storage", plugin: str = "auto"):
+        """Initialize NIXL storage connector."""
+        self.file_manager = (
+            NixlFileManager(file_path)
+            if plugin not in NixlBackendSelection.OBJ_PLUGINS
+            else None
+        )
+
+        agent_config = nixl_agent_config(backends=[])
+        self.agent_name = f"hicache_nixl_{str(uuid.uuid4())}"
+        self.agent = nixl_agent(self.agent_name, agent_config)
+
+        self.backend_selector = NixlBackendSelection(plugin)
+        if not self.backend_selector.create_backend(self.agent):
+            raise RuntimeError("Failed to create NIXL backend")
+
+        self.registration = NixlRegistration(self.agent)
+
+    def _execute_transfer(
+        self, tensors: List[torch.Tensor], keys: List[str], direction: str
+    ) -> bool:
+        if len(tensors) != len(keys):
+            logger.error("Mismatch between number of tensors and files/objects")
+            return False
+
+        if not self.registration.register_buffers(tensors):
+            logger.error("Failed to register tensors")
+            return False
+
+        # Get transfer tuples based on backend type
+        tensor_sizes = [tensor.element_size() * tensor.numel() for tensor in tensors]
+        if self.backend_selector.mem_type == "FILE":
+            file_tuples = self.file_manager.files_to_nixl_tuples(keys)
+            if not file_tuples or not self.registration.register_files(file_tuples):
+                logger.error("Failed to prepare files for transfer")
+                return False
+            transfer_tuples = [
+                (x[0], s, x[2]) for x, s in zip(file_tuples, tensor_sizes)
+            ]
+        else:
+            if not self.registration.register_objects(keys, tensors):
+                logger.error("Failed to register objects")
+                return False
+            transfer_tuples = [(0, s, key) for s, key in zip(tensor_sizes, keys)]
+
+        try:
+            # Get transfer descriptors
+            if (tensor_descs := self.agent.get_xfer_descs(tensors)) is None or (
+                file_descs := self.agent.get_xfer_descs(
+                    transfer_tuples, self.backend_selector.mem_type
+                )
+            ) is None:
+                logger.error("Failed to get transfer descriptors")
+                return False
+
+            # Initialize and execute transfer
+            if (
+                xfer_req := self.agent.initialize_xfer(
+                    direction, tensor_descs, file_descs, self.agent_name
+                )
+            ) is None:
+                logger.error("Failed to create transfer request")
+                return False
+
+            state = self.agent.transfer(xfer_req)
+            while state != "DONE":
+                state = self.agent.check_xfer_state(xfer_req)
+                if state == "ERR":
+                    logger.error("Transfer failed")
+                    return False
+            time.sleep(0.0001)  # Can be changed to os.sched_yield() or parametrized
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to execute transfer: {e}")
+            import traceback
+
+            logger.error(f"Traceback: {traceback.format_exc()}")
+            return False
+
+    def batch_set(self, keys: List[str], values: List[torch.Tensor]) -> bool:
+        if not keys:
+            return True
+
+        if self.backend_selector.mem_type == "FILE":
+            file_paths = []
+            for key in keys:
+                tensor_path = self.file_manager.get_file_path(key)
+                if not self.file_manager.create_file(tensor_path):
+                    logger.error(f"Failed to create file {tensor_path}")
+                    return False
+                file_paths.append(tensor_path)
+            return self._execute_transfer(values, file_paths, "WRITE")
+        else:
+            return self._execute_transfer(values, keys, "WRITE")
+
+    def set(self, key: str, value: torch.Tensor) -> bool:
+        return self.batch_set([key], [value])
+
+    def get(
+        self, key: str, dst_tensor: Optional[torch.Tensor] = None
+    ) -> torch.Tensor | None:
+        if dst_tensor is None:  # To be removed, being compatible with the current API
+            return None
+        result = self.batch_get([key], [dst_tensor])
+        return result[0] if result else None
+
+    def batch_get(
+        self, keys: List[str], dst_tensors: List[torch.Tensor]
+    ) -> List[Optional[torch.Tensor]]:
+        if not keys:
+            return []
+
+        if self.backend_selector.mem_type == "FILE":
+            file_paths = [self.file_manager.get_file_path(key) for key in keys]
+            success = self._execute_transfer(dst_tensors, file_paths, "READ")
+        else:
+            success = self._execute_transfer(dst_tensors, keys, "READ")
+        return dst_tensors if success else [None] * len(keys)
+
+    def exists(self, key: str) -> bool:
+        tuples = self.registration.create_query_tuples(
+            key,
+            self.backend_selector.mem_type,
+            self.file_manager if self.backend_selector.mem_type == "FILE" else None,
+        )
+        if not tuples:
+            return False
+
+        query_res = self.agent.query_memory(
+            tuples,
+            self.backend_selector.backend_name,
+            mem_type=self.backend_selector.mem_type,
+        )
+        return query_res[0] is not None  # can be expanded to multiple keys