From 99e154dc84de7f3ac7b43898aec1779c27eb60dd Mon Sep 17 00:00:00 2001 From: fems14 <74094523+fems14@users.noreply.github.com> Date: Sat, 25 Oct 2025 14:21:30 +0800 Subject: [PATCH] [0.11.0] cherry-pick from #3747 (#3746) cherry-pick from #3747 correct _register function place for mooncacke Signed-off-by: fems14 <1804143737@qq.com> --- .../distributed/mooncake/mooncake_engine.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/vllm_ascend/distributed/mooncake/mooncake_engine.py b/vllm_ascend/distributed/mooncake/mooncake_engine.py index bff3f3e..cc58877 100644 --- a/vllm_ascend/distributed/mooncake/mooncake_engine.py +++ b/vllm_ascend/distributed/mooncake/mooncake_engine.py @@ -126,16 +126,6 @@ class MooncakeEngine: region_len = self.num_blocks * self.block_len[0] 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: self.get_event = threading.Event() if self.kv_role in ['kv_producer', 'kv_both']: @@ -170,6 +160,16 @@ class MooncakeEngine: self.kv_recv_thread.start() 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): self.current_layer = 0 self.layerwise_retrievers = []