feat: optimize hybrid radix cache buffer insertion strategy

Add mamba_model_match_len based optimization for buffer insertion:
- Only insert buffer at actual branch points instead of fixed intervals
- Use threshold (chunked_prefill_size // 2) to decide strategy
- Reduce buffer storage overhead while maintaining cache hit rate

Author: sufubao

lightllm/server/router/dynamic_prompt/hybrid_radix_cache.py

@@ -80,14 +80,38 @@ def _evict_buffer(self, need_evict_buffer_num, evict_buffer_callback, evict_toke
                     self.evict_tree_set.add(parent_node)
         return
 
+    def _should_insert_buffer(self, req) -> bool:
+        """决定是否需要在当前位置插入 buffer"""
+        # 情况1：prefill 完成（即将进入 decode），必须插入
+        if req.cur_kv_len >= req.get_cur_total_len():
+            return True
+
+        # 情况2：使用优化策略时
+        if req.use_mamba_match_len_strategy:
+            # 只在 mamba_model_match_len 位置插入
+            if req.cur_kv_len == req.mamba_model_match_len and not req.mamba_buffer_inserted:
+                return True
+            return False
+
+        # 情况3：原策略（每个 chunk 后都插入）
+        return True
+
     def insert_for_hybrid_radix_cache(self, reqs):
         from lightllm.server.router.model_infer.infer_batch import g_infer_context
 
-        self.free_radix_cache_to_get_enough_buffer(len(reqs))
-        new_buffer_indexes = self.mem_manager.alloc_buffer(len(reqs))
-        # req_ids_gpu = req_ids.cuda()
+        # 过滤需要插入的请求
+        reqs_to_insert = []
+        for req in reqs:
+            if self._should_insert_buffer(req):
+                reqs_to_insert.append(req)
+
+        if len(reqs_to_insert) == 0:
+            return
 
-        for i, req in enumerate(reqs):
+        self.free_radix_cache_to_get_enough_buffer(len(reqs_to_insert))
+        new_buffer_indexes = self.mem_manager.alloc_buffer(len(reqs_to_insert))
+
+        for i, req in enumerate(reqs_to_insert):
             input_token_ids = req.get_input_token_ids()
             key = torch.tensor(input_token_ids[0 : req.cur_kv_len], dtype=torch.int64, device="cpu")
             value = g_infer_context.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].cpu()
@@ -105,6 +129,10 @@ def insert_for_hybrid_radix_cache(self, reqs):
             # free_a_req_mem 中会释放 [prompt_cache_len:prefix_len]，更新后这个范围为空
             req.shm_req.prompt_cache_len = req.cur_kv_len
 
+            # 标记已在 mamba_model_match_len 位置插入
+            if req.cur_kv_len == req.mamba_model_match_len:
+                req.mamba_buffer_inserted = True
+
     def match_prefix(self, key, update_refs=False):
         assert len(key) != 0
         self.match_count = (self.match_count + 1) % self.log_interval

lightllm/server/router/model_infer/infer_batch.py

@@ -409,6 +409,10 @@ def __init__(
 
         # 在开启radix cache的情况下，用于标记命中情况，用于插入算法
         self.mamba_model_match_len = 0
+        # 是否使用基于 mamba_model_match_len 的优化策略
+        self.use_mamba_match_len_strategy = False
+        # 是否已在 mamba_model_match_len 位置插入 buffer
+        self.mamba_buffer_inserted = False
 
         # 在开启 enable_cpu_cache 的情况下，当请求结束后，会将请求的 kv cache
         # 卸载到 cpu cache 中，该标志变量用于标记请求的卸载任务的状态
@@ -471,6 +475,12 @@ def _match_radix_cache(self):
                 self.cur_kv_len = int(ready_cache_len)  # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换
                 self.shm_req.prompt_cache_len = self.cur_kv_len  # 记录 prompt cache 的命中长度
 
+                # 判断是否使用基于 mamba_model_match_len 的优化策略
+                # 当需要重新计算的增量足够大时，值得单独在分支点保存 buffer
+                increment = self.mamba_model_match_len - ready_cache_len
+                threshold = self.shm_req.chunked_prefill_size // 2
+                self.use_mamba_match_len_strategy = increment >= threshold
+
         self.shm_req.shm_cur_kv_len = self.cur_kv_len
         return
 
@@ -518,13 +528,24 @@ def get_input_token_ids(self):
         return self.shm_req.shm_prompt_ids.arr[0 : self.get_cur_total_len()]
 
     def get_chuncked_input_token_ids(self):
-        chunked_start = self.cur_kv_len
-        chunked_end = min(self.get_cur_total_len(), chunked_start + self.shm_req.chunked_prefill_size)
+        # 复用 get_chuncked_input_token_len 的逻辑，保持一致性
+        chunked_end = self.get_chuncked_input_token_len()
         return self.shm_req.shm_prompt_ids.arr[0:chunked_end]
 
     def get_chuncked_input_token_len(self):
         chunked_start = self.cur_kv_len
         chunked_end = min(self.get_cur_total_len(), chunked_start + self.shm_req.chunked_prefill_size)
+
+        # 优化策略：第一个 chunk 直接到 mamba_model_match_len（分支点）
+        # 这样可以在分支点位置保存 buffer，提升后续请求的缓存命中率
+        if (
+            self.use_mamba_match_len_strategy
+            and not self.mamba_buffer_inserted
+            and self.mamba_model_match_len > chunked_start
+            and self.mamba_model_match_len <= self.get_cur_total_len()
+        ):
+            chunked_end = self.mamba_model_match_len
+
         return chunked_end
 
     def set_next_gen_token_id(self, next_token_id: int, logprob: float, output_len: int):