chore: import upstream snapshot with attribution

python/sglang/srt/mem_cache/hisparse_memory_pool.py

@@ -0,0 +1,122 @@
+# mapping on device memory, host memory and memory allocator
+
+import logging
+from typing import Optional
+
+import torch
+
+from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
+from sglang.srt.utils import is_cuda, is_hip
+
+logger = logging.getLogger(__name__)
+
+# sgl_kernel.kvcacheio is only available in CUDA/ROCm sgl-kernel builds (not XPU/MPS/NPU/CPU).
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+if _is_cuda or _is_hip:
+    from sgl_kernel.kvcacheio import transfer_kv_all_layer_mla
+else:
+
+    def transfer_kv_all_layer_mla(*args, **kwargs):
+        raise RuntimeError(
+            "HiSparse device KV transfer requires sgl_kernel.kvcacheio (CUDA/ROCm). "
+            "It is not available on this backend."
+        )
+
+
+class HiSparseDSATokenToKVPool(DSATokenToKVPool):
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        kv_lora_rank: int,
+        dtype: torch.dtype,
+        qk_rope_head_dim: int,
+        layer_num: int,
+        device: str,
+        index_head_dim: int,
+        enable_memory_saver: bool,
+        kv_cache_dim: int,
+        start_layer: Optional[int] = None,
+        end_layer: Optional[int] = None,
+        host_to_device_ratio: int = 2,
+    ):
+        super().__init__(
+            size=size,
+            page_size=page_size,
+            kv_lora_rank=kv_lora_rank,
+            dtype=dtype,
+            qk_rope_head_dim=qk_rope_head_dim,
+            layer_num=layer_num,
+            device=device,
+            index_head_dim=index_head_dim,
+            enable_memory_saver=enable_memory_saver,
+            kv_cache_dim=kv_cache_dim,
+            start_layer=start_layer,
+            end_layer=end_layer,
+            index_buf_size=size * host_to_device_ratio,
+        )
+        self.bytes_per_token = self.kv_cache_dim * self.dtype.itemsize
+
+    def register_mapping(self, full_to_hisparse_device_index_mapping: torch.Tensor):
+        self.full_to_hisparse_device_index_mapping = (
+            full_to_hisparse_device_index_mapping
+        )
+
+    def translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
+        return self.full_to_hisparse_device_index_mapping[compressed_indices]
+
+    def _translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
+        return self.full_to_hisparse_device_index_mapping[compressed_indices]
+
+    def translate_loc_from_full_to_hisparse_device(self, full_indices: torch.Tensor):
+        return self._translate_loc_to_hisparse_device(full_indices)
+
+    def translate_loc_from_full_to_compressed(self, full_indices: torch.Tensor):
+        return full_indices
+
+    def set_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+    ):
+        loc = self.translate_loc_to_hisparse_device(loc)
+        super().set_kv_buffer(layer, loc, cache_k, cache_v)
+
+    def set_mla_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        cache_k_nope: torch.Tensor,
+        cache_k_rope: torch.Tensor,
+    ):
+        loc = self.translate_loc_to_hisparse_device(loc)
+        super().set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope)
+
+    def get_mla_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        dst_dtype: Optional[torch.dtype] = None,
+    ):
+        loc = self.translate_loc_to_hisparse_device(loc)
+        return super().get_mla_kv_buffer(layer, loc, dst_dtype)
+
+    def transfer_values_on_device(self, dst_indices, src_indices):
+        transfer_kv_all_layer_mla(
+            src_layers=self.data_ptrs,
+            dst_layers=self.data_ptrs,
+            src_indices=src_indices,
+            dst_indices=dst_indices,
+            item_size=self.bytes_per_token,
+            num_layers=self.layer_num,
+        )
+
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        raise NotImplementedError("HiSparseDevicePool does not support get_cpu_copy")
+
+    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+        raise NotImplementedError("HiSparseDevicePool does not support load_cpu_copy")