EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

support mooncake store dp attention (#9684)

Browse Source
This commit is contained in:
huangtingwei
2025-08-28 12:31:31 +08:00
committed by GitHub
parent e1f7cf57dc
commit 55349e361d
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/managers/cache_controller.py
View File

@@ -636,6 +636,7 @@ class HiCacheController:
key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta( key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta(
hash_values, hash_values,
host_indices, host_indices,
self.storage_config.tp_rank,
) )
get_result = self.storage_backend.batch_get( get_result = self.storage_backend.batch_get(
key_strs, key_strs,
@@ -838,6 +839,7 @@ class HiCacheController:
key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta( key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta(
hash_values, hash_values,
host_indices, host_indices,
self.storage_config.tp_rank,
) )
success = self.storage_backend.batch_set( success = self.storage_backend.batch_set(
key_strs, key_strs,

6
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -7,7 +7,6 @@ from functools import wraps
import psutil import psutil
import torch 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.mem_cache.memory_pool import KVCache, MHATokenToKVPool, MLATokenToKVPool
from sglang.srt.utils import is_npu from sglang.srt.utils import is_npu
@@ -464,8 +463,7 @@ class MHATokenToKVPoolHost(HostKVCache):
else: else:
raise ValueError(f"Unsupported layout: {self.layout}") raise ValueError(f"Unsupported layout: {self.layout}")
def get_buffer_meta(self, keys, indices): def get_buffer_meta(self, keys, indices, local_rank):
local_rank = get_tensor_model_parallel_rank()
ptr_list = [] ptr_list = []
key_list = [] key_list = []
kv_buffer_data_ptr = self.kv_buffer.data_ptr() kv_buffer_data_ptr = self.kv_buffer.data_ptr()
@@ -704,7 +702,7 @@ class MLATokenToKVPoolHost(HostKVCache):
else: else:
raise ValueError(f"Unsupported layout: {self.layout}") 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 = [] ptr_list = []
key_list = [] key_list = []
kv_buffer_data_ptr = self.kv_buffer.data_ptr() kv_buffer_data_ptr = self.kv_buffer.data_ptr()