fix

Author: wangzaijun

lightllm/server/router/dynamic_prompt/radix_cache.py

@@ -333,6 +333,18 @@ def add_node_ref_counter(self, node: TreeNode):
             self.evict_tree_set.add(old_node)
         return
 
+    def get_mem_index_value_by_node(self, node: TreeNode) -> Optional[torch.Tensor]:
+        if node is None:
+            return None
+
+        ans_list = []
+        while node is not None:
+            ans_list.append(node.token_mem_index_value)
+            node = node.parent
+
+        ans_list.reverse()
+        return torch.concat(ans_list, dim=0)
+
     def get_refed_tokens_num(self):
         return self.refed_tokens_num.arr[0]
 

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_impl.py

@@ -85,6 +85,11 @@ def _prefill_req_frozen_tokens_and_put_to_kvmove_taskqueue(self, finished_reqs:
                 self.radix_cache.add_node_ref_counter(new_shared_kv_node)
                 req.shared_kv_node = new_shared_kv_node
 
+                _kv_len = req.cur_kv_len
+                _value = self.radix_cache.get_mem_index_value_by_node(new_shared_kv_node)
+                assert len(_value) == _kv_len
+                self.model.req_manager.req_to_token_indexs[req.req_idx][0:_kv_len] = _value
+
                 assert new_shared_kv_node.node_prefix_total_len == req.cur_kv_len
 
                 if req.shm_req.sample_params.move_kv_to_decode_node.exists: