sglang 0.5.0rc1__py3-none-any.whl → 0.5.1__py3-none-any.whl

sglang/srt/mem_cache/allocator.py

@@ -20,7 +20,6 @@ Page-aligned memory pool.
 """
 
 import abc
-import weakref
 from typing import TYPE_CHECKING
 
 import torch
@@ -81,9 +80,6 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
         if self.free_group:
             self.free(torch.cat(self.free_group))
 
-    def estimated_num_new_pages(self, bs, extend_num_tokens):
-        return bs * ((extend_num_tokens + self.page_size - 1) // self.page_size)
-
     def merge_and_sort_free(self):
         if len(self.release_pages) > 0:
             self.free_pages = torch.cat((self.free_pages, self.release_pages))
@@ -149,6 +145,7 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def alloc(self, need_size: int):
         if self.need_sort and need_size > len(self.free_pages):
             self.merge_and_sort_free()
+
         if need_size > len(self.free_pages):
             return None
 
@@ -442,6 +439,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.num_pages = size // page_size
         self.debug_mode = get_bool_env_var("SGLANG_DEBUG_MEMORY_POOL")
         self.ret_values = torch.empty((), dtype=torch.int64, device=self.device)
+        self.seen_max_num_extend_tokens_next_power_of_2 = 1
         self.clear()
 
     def alloc(self, need_size: int):
@@ -479,8 +477,13 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
                 (last_loc + 1) % self.page_size == prefix_lens % self.page_size
             )
 
+        self.seen_max_num_extend_tokens_next_power_of_2 = max(
+            self.seen_max_num_extend_tokens_next_power_of_2,
+            next_power_of_2(extend_num_tokens),
+        )
+
         bs = len(prefix_lens)
-        if self.need_sort and self.estimated_num_new_pages(bs, extend_num_tokens) > len(
+        if self.need_sort and extend_num_tokens // self.page_size + bs + 1 > len(
             self.free_pages
         ):
             self.merge_and_sort_free()
@@ -497,7 +500,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             self.ret_values,
             next_power_of_2(bs),
             self.page_size,
-            next_power_of_2(extend_num_tokens),
+            self.seen_max_num_extend_tokens_next_power_of_2,
         )
 
         if self.debug_mode:
@@ -522,9 +525,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             )
 
         bs = len(seq_lens)
-        if self.need_sort and self.estimated_num_new_pages(bs, 1) > len(
-            self.free_pages
-        ):
+        if self.need_sort and bs > len(self.free_pages):
             self.merge_and_sort_free()
 
         out_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
@@ -578,151 +579,3 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
 
     def load_cpu_copy(self, kv_cache_cpu, indices):
         return self._kvcache.load_cpu_copy(kv_cache_cpu, indices)
-
-
-def alloc_extend_kernel_ascend(
-    prefix_lens,
-    seq_lens,
-    last_loc,
-    free_pages,
-    out_indices,
-    page_size,
-    device,
-):
-    extend_lens = seq_lens - prefix_lens
-    end_pos = torch.cumsum(extend_lens, 0)
-    start_pos = end_pos - extend_lens
-    num_new_pages = (seq_lens + page_size - 1) // page_size - (
-        prefix_lens + page_size - 1
-    ) // page_size
-    num_full_new_pages = (seq_lens) // page_size - (
-        prefix_lens + page_size - 1
-    ) // page_size
-    need_page = num_new_pages - num_full_new_pages
-    end_new_pages = torch.cumsum(num_new_pages, 0)
-    start_new_pages = end_new_pages - num_new_pages
-    pos_in_page = torch.arange(page_size, device=device, dtype=torch.int32)
-    for i in range(len(prefix_lens)):
-        num1 = (
-            min(
-                seq_lens[i],
-                (prefix_lens[i] + page_size - 1) // page_size * page_size,
-            )
-            - prefix_lens[i]
-        )
-        if num1:
-            out_indices[start_pos[i] : start_pos[i] + num1] = (
-                last_loc[i] + 1 + pos_in_page[:num1].view(-1)
-            )
-
-        num2 = (
-            seq_lens[i] // page_size - (prefix_lens[i] + page_size - 1) // page_size
-        ) * page_size
-        if num2:
-            pages = (
-                free_pages[start_new_pages[i] : end_new_pages[i] - need_page[i]]
-                * page_size
-            )
-            out_indices[start_pos[i] + num1 : start_pos[i] + num1 + num2] = (
-                pages.view(-1, 1) + pos_in_page.view(1, -1)
-            ).view(-1)
-
-        num3 = seq_lens[i] - seq_lens[i] // page_size * page_size
-        if num3:
-            out_indices[end_pos[i] - num3 : end_pos[i]] = (
-                free_pages[end_new_pages[i] - 1] * page_size + pos_in_page[:num3]
-            ).view(-1)
-
-
-class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
-
-    def __init__(
-        self,
-        size: int,
-        page_size: int,
-        dtype: torch.dtype,
-        device: str,
-        kvcache: KVCache,
-        need_sort: bool,
-    ):
-        super().__init__(size, page_size, dtype, device, kvcache, need_sort)
-
-    def alloc_extend(
-        self,
-        prefix_lens: torch.Tensor,
-        seq_lens: torch.Tensor,
-        last_loc: torch.Tensor,
-        extend_num_tokens: int,
-    ):
-        if self.debug_mode:
-            assert torch.all(
-                (last_loc + 1) % self.page_size == prefix_lens % self.page_size
-            )
-
-        estimated_num_new_pages = (
-            (
-                (seq_lens + self.page_size - 1) // self.page_size
-                - (prefix_lens + self.page_size - 1) // self.page_size
-            )
-            .sum()
-            .item()
-        )
-        if self.need_sort and estimated_num_new_pages > len(self.free_pages):
-            self.merge_and_sort_free()
-
-        if estimated_num_new_pages > len(self.free_pages):
-            return None
-
-        out_indices = torch.empty(
-            (extend_num_tokens,), dtype=torch.int32, device=self.device
-        )
-
-        alloc_extend_kernel_ascend(
-            prefix_lens,
-            seq_lens,
-            last_loc,
-            self.free_pages,
-            out_indices,
-            self.page_size,
-            self.device,
-        )
-
-        if self.debug_mode:
-            assert len(torch.unique(out_indices)) == len(out_indices)
-
-        self.free_pages = self.free_pages[estimated_num_new_pages:]
-        return out_indices
-
-    def alloc_decode(
-        self,
-        seq_lens: torch.Tensor,
-        last_loc: torch.Tensor,
-    ):
-        if self.debug_mode:
-            assert torch.all(
-                (last_loc + 2) % self.page_size == seq_lens % self.page_size
-            )
-
-        need_new_pages = (seq_lens % self.page_size == 1).int()
-        num_new_pages = need_new_pages.sum().item()
-
-        if num_new_pages > len(self.free_pages):
-            self.merge_and_sort_free()
-
-        if num_new_pages > len(self.free_pages):
-            return None
-
-        end_new_pages = torch.cumsum(need_new_pages, 0)
-        start_new_pages = end_new_pages - need_new_pages
-        if num_new_pages == 0:
-            out_indices = last_loc + 1
-        else:
-            out_indices = (last_loc + 1) * (1 - need_new_pages) + self.free_pages[
-                start_new_pages
-            ] * self.page_size * need_new_pages
-
-        if self.debug_mode:
-            assert len(torch.unique(out_indices)) == len(out_indices)
-
-        self.free_pages = self.free_pages[num_new_pages:]
-        return out_indices.int()

sglang/srt/mem_cache/allocator_ascend.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.mem_cache.allocator import PagedTokenToKVPoolAllocator
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.memory_pool import KVCache
+
+
+def alloc_extend_kernel_ascend(
+    prefix_lens,
+    seq_lens,
+    last_loc,
+    free_pages,
+    out_indices,
+    page_size,
+    device,
+):
+    extend_lens = seq_lens - prefix_lens
+    end_pos = torch.cumsum(extend_lens, 0)
+    start_pos = end_pos - extend_lens
+    num_new_pages = (seq_lens + page_size - 1) // page_size - (
+        prefix_lens + page_size - 1
+    ) // page_size
+    num_full_new_pages = (seq_lens) // page_size - (
+        prefix_lens + page_size - 1
+    ) // page_size
+    need_page = num_new_pages - num_full_new_pages
+    end_new_pages = torch.cumsum(num_new_pages, 0)
+    start_new_pages = end_new_pages - num_new_pages
+    pos_in_page = torch.arange(page_size, device=device, dtype=torch.int32)
+    for i in range(len(prefix_lens)):
+        num1 = (
+            min(
+                seq_lens[i],
+                (prefix_lens[i] + page_size - 1) // page_size * page_size,
+            )
+            - prefix_lens[i]
+        )
+        if num1:
+            out_indices[start_pos[i] : start_pos[i] + num1] = (
+                last_loc[i] + 1 + pos_in_page[:num1].view(-1)
+            )
+
+        num2 = (
+            seq_lens[i] // page_size - (prefix_lens[i] + page_size - 1) // page_size
+        ) * page_size
+        if num2:
+            pages = (
+                free_pages[start_new_pages[i] : end_new_pages[i] - need_page[i]]
+                * page_size
+            )
+            out_indices[start_pos[i] + num1 : start_pos[i] + num1 + num2] = (
+                pages.view(-1, 1) + pos_in_page.view(1, -1)
+            ).view(-1)
+
+        num3 = seq_lens[i] - seq_lens[i] // page_size * page_size
+        if num3:
+            out_indices[end_pos[i] - num3 : end_pos[i]] = (
+                free_pages[end_new_pages[i] - 1] * page_size + pos_in_page[:num3]
+            ).view(-1)
+
+
+class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
+
+    def alloc_extend(
+        self,
+        prefix_lens: torch.Tensor,
+        seq_lens: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+    ):
+        if self.debug_mode:
+            assert torch.all(
+                (last_loc + 1) % self.page_size == prefix_lens % self.page_size
+            )
+
+        num_new_pages = (
+            (
+                (seq_lens + self.page_size - 1) // self.page_size
+                - (prefix_lens + self.page_size - 1) // self.page_size
+            )
+            .sum()
+            .item()
+        )
+        if self.need_sort and num_new_pages > len(self.free_pages):
+            self.merge_and_sort_free()
+
+        if num_new_pages > len(self.free_pages):
+            return None
+
+        out_indices = torch.empty(
+            (extend_num_tokens,), dtype=torch.int32, device=self.device
+        )
+
+        alloc_extend_kernel_ascend(
+            prefix_lens,
+            seq_lens,
+            last_loc,
+            self.free_pages,
+            out_indices,
+            self.page_size,
+            self.device,
+        )
+
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+
+        self.free_pages = self.free_pages[num_new_pages:]
+        return out_indices
+
+    def alloc_decode(
+        self,
+        seq_lens: torch.Tensor,
+        last_loc: torch.Tensor,
+    ):
+        if self.debug_mode:
+            assert torch.all(
+                (last_loc + 2) % self.page_size == seq_lens % self.page_size
+            )
+
+        need_new_pages = (seq_lens % self.page_size == 1).int()
+        num_new_pages = need_new_pages.sum().item()
+
+        if num_new_pages > len(self.free_pages):
+            self.merge_and_sort_free()
+
+        if num_new_pages > len(self.free_pages):
+            return None
+
+        end_new_pages = torch.cumsum(need_new_pages, 0)
+        start_new_pages = end_new_pages - need_new_pages
+        if num_new_pages == 0:
+            out_indices = last_loc + 1
+        else:
+            out_indices = (last_loc + 1) * (1 - need_new_pages) + self.free_pages[
+                start_new_pages
+            ] * self.page_size * need_new_pages
+
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+
+        self.free_pages = self.free_pages[num_new_pages:]
+        return out_indices.int()

sglang/srt/mem_cache/chunk_cache.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 """Cache for chunked prefill, used when RadixCache is disabled."""
 
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 

sglang/srt/mem_cache/hicache_storage.py

@@ -13,6 +13,11 @@ from sglang.srt.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
+from sglang.srt.layers.dp_attention import (
+    get_attention_tp_rank,
+    get_attention_tp_size,
+    is_dp_attention_enabled,
+)
 
 
 def get_hash_str(token_ids: List[int], prior_hash: str = None) -> str:
@@ -101,11 +106,16 @@ class HiCacheStorage(ABC):
 
 class HiCacheFile(HiCacheStorage):
 
-    def __init__(self, file_path: str = "/tmp/hicache"):
+    def __init__(self, file_path: str = "/tmp/hicache", is_mla: bool = False):
         self.file_path = os.getenv("SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR", file_path)
-        tp_rank = get_tensor_model_parallel_rank()
-        tp_size = get_tensor_model_parallel_world_size()
-        self.tp_suffix = f"_{tp_rank}_{tp_size}" if tp_size > 1 else ""
+        if is_dp_attention_enabled():
+            tp_rank = get_attention_tp_rank()
+            tp_size = get_attention_tp_size()
+        else:
+            tp_rank = get_tensor_model_parallel_rank()
+            tp_size = get_tensor_model_parallel_world_size()
+
+        self.tp_suffix = f"_{tp_rank}_{tp_size}" if tp_size > 1 and not is_mla else ""
         if not os.path.exists(self.file_path) and tp_rank == 0:
             os.makedirs(self.file_path)
             logger.info(f"Created HiCacheFile storage directory at {self.file_path}")

sglang/srt/mem_cache/memory_pool.py

@@ -849,7 +849,7 @@ class MLATokenToKVPool(KVCache):
             cache_k_rope = cache_k_rope.view(self.store_dtype)
 
         set_mla_kv_buffer_triton(
-            self.kv_buffer[layer_id], loc, cache_k_nope, cache_k_rope
+            self.kv_buffer[layer_id - self.start_layer], loc, cache_k_nope, cache_k_rope
         )
 
     def get_cpu_copy(self, indices):
@@ -951,7 +951,7 @@ class AscendMLAPagedTokenToKVPool(MLATokenToKVPool):
             cache_k = cache_k.to(self.dtype)
 
         if self.store_dtype != self.dtype:
-            cache_k = cache_k.view(store_dtype)
+            cache_k = cache_k.view(self.store_dtype)
 
         import torch_npu
 
@@ -1070,7 +1070,7 @@ def copy_all_layer_kv_cache(
     num_loop = tl.cdiv(stride, BLOCK_SIZE)
     for i in range(num_loop):
         copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
-        mask = (num_locs_offset < num_locs)[:, None] and (copy_offset < stride)[None, :]
+        mask = (num_locs_offset < num_locs)[:, None] & (copy_offset < stride)[None, :]
         value = tl.load(
             data_ptr + src_locs[:, None] * stride + copy_offset[None, :], mask=mask
         )

sglang/srt/mem_cache/memory_pool_host.py

@@ -7,6 +7,7 @@ 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
 
@@ -307,6 +308,9 @@ class MHATokenToKVPoolHost(HostKVCache):
 
         return self.head_dim * self.head_num * self.layer_num * self.dtype.itemsize * 2
 
+    def get_ksize_per_token(self):
+        return self.get_size_per_token() // 2
+
     def init_kv_buffer(self):
         if self.layout == "layer_first":
             dims = (2, self.layer_num, self.size, self.head_num, self.head_dim)
@@ -484,8 +488,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_}_k")
-            key_list.append(f"{key_}_v")
+            key_list.append(f"{key_}_{get_tensor_model_parallel_rank()}_k")
+            key_list.append(f"{key_}_{get_tensor_model_parallel_rank()}_v")
         element_size = (
             self.layer_num
             * self.dtype.itemsize
@@ -496,6 +500,21 @@ 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):
+        assert self.layout == "page_first"
+        assert len(keys) == (len(indices) // self.page_size)
+
+        key_list = []
+        buf_list = []
+
+        for key, i in zip(keys, range(0, len(indices), self.page_size)):
+            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])
+
+        return key_list, buf_list
+
 
 class MLATokenToKVPoolHost(HostKVCache):
     device_pool: MLATokenToKVPool
@@ -538,6 +557,9 @@ class MLATokenToKVPoolHost(HostKVCache):
             * self.layer_num
         )
 
+    def get_ksize_per_token(self):
+        return self.get_size_per_token()
+
     def init_kv_buffer(self):
         if self.layout == "layer_first":
             dims = (
@@ -704,3 +726,14 @@ 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):
+        assert self.layout == "page_first"
+        assert 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])
+
+        return keys, buf_list

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

@@ -7,10 +7,15 @@ import signal
 import threading
 from abc import ABC, abstractmethod
 from functools import wraps
-from typing import List, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 import torch
 
+from sglang.srt.distributed import get_tensor_model_parallel_rank
+from sglang.srt.layers.dp_attention import (
+    get_attention_tp_rank,
+    is_dp_attention_enabled,
+)
 from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
 from sglang.srt.mem_cache.storage.hf3fs.client_hf3fs import Hf3fsClient
 
@@ -167,13 +172,20 @@ class HiCacheHF3FS(HiCacheStorage):
 
     @staticmethod
     def from_env_config(
-        rank: int, bytes_per_page: int, dtype: torch.dtype
+        bytes_per_page: int, dtype: torch.dtype, rank: int = None
     ) -> "HiCacheHF3FS":
         from sglang.srt.mem_cache.storage.hf3fs.mini_3fs_metadata_server import (
             Hf3fsGlobalMetadataClient,
             Hf3fsLocalMetadataClient,
         )
 
+        if rank is None:
+            rank = (
+                get_attention_tp_rank()
+                if is_dp_attention_enabled()
+                else get_tensor_model_parallel_rank()
+            )
+
         config_path = os.getenv(HiCacheHF3FS.default_env_var)
         if not config_path:
             return HiCacheHF3FS(
@@ -228,15 +240,23 @@ class HiCacheHF3FS(HiCacheStorage):
         )
 
     def get(
-        self, key: str, target_location: Optional[torch.Tensor] = None
+        self,
+        key: str,
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
     ) -> torch.Tensor | None:
-        return self.batch_get([key], [target_location] if target_location else None)[0]
+        return self.batch_get(
+            [key],
+            [target_location] if target_location is not None else None,
+            [target_sizes] if target_sizes is not None else None,
+        )[0]
 
     @synchronized()
     def batch_get(
         self,
         keys: List[str],
-        target_locations: Optional[List[torch.Tensor]] = None,
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
     ) -> List[torch.Tensor | None]:
         page_indices = self.metadata_client.get_page_indices(self.rank, keys)
 
@@ -246,9 +266,15 @@ class HiCacheHF3FS(HiCacheStorage):
                 batch_indices.append(i)
                 file_offsets.append(page_index * self.bytes_per_page)
 
-        file_results = [
-            torch.empty(self.numel, dtype=self.dtype) for _ in range(len(batch_indices))
-        ]
+        if target_locations is not None:
+            for target_location in target_locations:
+                assert target_location.is_contiguous()
+            file_results = target_locations
+        else:
+            file_results = [
+                torch.empty(self.numel, dtype=self.dtype)
+                for _ in range(len(batch_indices))
+            ]
 
         futures = [
             self.executor.submit(
@@ -273,10 +299,27 @@ class HiCacheHF3FS(HiCacheStorage):
 
         return results
 
-    def set(self, key: str, value: torch.Tensor) -> bool:
-        return self.batch_set([key], [value])
+    def set(
+        self,
+        key: str,
+        value: Optional[Any] = None,
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> bool:
+        return self.batch_set(
+            [key],
+            [value] if value is not None else None,
+            [target_location] if target_location is not None else None,
+            [target_sizes] if target_sizes is not None else None,
+        )
 
-    def batch_set(self, keys: List[str], values: List[torch.Tensor]) -> bool:
+    def batch_set(
+        self,
+        keys: List[str],
+        values: Optional[Any] = None,
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> bool:
         # Todo: Add prefix block's hash key
         key_with_prefix = [(key, "") for key in keys]
         indices = self.metadata_client.reserve_and_allocate_page_indices(
@@ -292,7 +335,8 @@ class HiCacheHF3FS(HiCacheStorage):
 
             batch_indices.append(i)
             file_offsets.append(page_index * self.bytes_per_page)
-            file_values.append(value.contiguous())
+            assert value.is_contiguous()
+            file_values.append(value)
 
         futures = [
             self.executor.submit(

sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -19,14 +19,13 @@ logger = logging.getLogger(__name__)
 
 
 def get_hash_str_mooncake(token_ids: List[int], prior_hash: str = None):
-    local_rank = get_tensor_model_parallel_rank()
     prefix_str = ""
     if prior_hash:
         prefix_str = hashlib.sha256(prior_hash.encode()).hexdigest()
     current_token_ids_bytes = np.array(token_ids).tobytes()
     current_hash_object = hashlib.sha256(current_token_ids_bytes)
     current_hash_hex = current_hash_object.hexdigest()
-    return f"{prefix_str}_{int(current_hash_hex[:16], 16)}_{local_rank}"
+    return f"{prefix_str}_{int(current_hash_hex[:16], 16)}"
 
 
 @dataclass
@@ -97,7 +96,7 @@ class MooncakeStoreConfig:
 
 
 class MooncakeStore(HiCacheStorage):
-    def __init__(self):
+    def __init__(self, is_mla: bool = False):
         try:
             from mooncake.store import MooncakeDistributedStore
         except ImportError as e:
@@ -127,6 +126,7 @@ class MooncakeStore(HiCacheStorage):
             logger.info("Connect to Mooncake store successfully.")
             self.warmup()
             logger.info("Mooncake store warmup successfully.")
+            self.is_mla = is_mla
 
         except ValueError as e:
             logger.error("Configuration loading failed: %s", e)
@@ -223,11 +223,15 @@ class MooncakeStore(HiCacheStorage):
 
     def exists(self, keys) -> bool | dict:
         _keys = []
+        local_rank = get_tensor_model_parallel_rank()
         for key in keys:
             if key is None:
                 return None
 
-            _keys.append(f"{key}_k")
+            if self.is_mla:
+                _keys.append(f"{key}_k")
+            else:
+                _keys.append(f"{key}_{local_rank}_k")
         result = {k: v for k, v in zip(keys, self.store.batch_is_exist(_keys))}
         return result