[bugfixfix] correct _register function place for mooncacke (vllm-project#3747)

correct _register function place for mooncacke

- vLLM version: v0.11.0rc3
- vLLM main:
vllm-project/vllm@17c540a

Signed-off-by: fems14 <[email protected]>

Author: fems14

vllm_ascend/distributed/mooncake/mooncake_engine.py

@@ -126,16 +126,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                         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 @@ def _register(self, ptr, length):
                 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 = []