EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[0.11.0] cherry-pick from #3747 (#3746)

Browse Source 
cherry-pick from #3747

correct _register function place for mooncacke

Signed-off-by: fems14 <1804143737@qq.com>
This commit is contained in:
fems14
2025-10-25 14:21:30 +08:00
committed by GitHub
parent fed8145aea
commit 99e154dc84
1 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
vllm_ascend/distributed/mooncake/mooncake_engine.py
View File

@@ -126,16 +126,6 @@ class MooncakeEngine:
region_len = self.num_blocks * self.block_len[0] region_len = self.num_blocks * self.block_len[0]
self._register(base_addr, region_len) self._register(base_addr, region_len)
def _register(self, ptr, length):
logger.debug(
"Registering KV cache: ptr=0x%x, length=%d, num_blocks=%d, "
"block_lens=%s", ptr, length, self.num_blocks, self.block_len)
try:
self.m_store.register_buffer(ptr, length)
except Exception as e:
raise RuntimeError(
f"Mooncake memory registration failed. Error is: {e}")
if self.use_layerwise: if self.use_layerwise:
self.get_event = threading.Event() self.get_event = threading.Event()
if self.kv_role in ['kv_producer', 'kv_both']: if self.kv_role in ['kv_producer', 'kv_both']:
@@ -170,6 +160,16 @@ class MooncakeEngine:
self.kv_recv_thread.start() self.kv_recv_thread.start()
ready_event.wait() ready_event.wait()
def _register(self, ptr, length):
logger.debug(
"Registering KV cache: ptr=0x%x, length=%d, num_blocks=%d, "
"block_lens=%s", ptr, length, self.num_blocks, self.block_len)
try:
self.m_store.register_buffer(ptr, length)
except Exception as e:
raise RuntimeError(
f"Mooncake memory registration failed. Error is: {e}")
def start_load_kv(self, metadata: MooncakeConnectorMetadata): def start_load_kv(self, metadata: MooncakeConnectorMetadata):
self.current_layer = 0 self.current_layer = 0
self.layerwise_retrievers = [] self.layerwise_retrievers = []