refactor: unify cache lifecycle and isolate KV storage resources

lightllm/common/basemodel/attention/linear/gdn.py

@@ -139,7 +139,7 @@ def prefill_att(
         layer_num = linear_att_dict["layer_num"]
         backend: LinearAttBackend = self.backend
 
-        conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num)
+        conv_states, ssm_states = self.infer_state.linear_state_manager.get_layer_buffers(layer_num)
         # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。
         # prefill 模式下，使用不到这几个维度，所以需要扣除掉，
         if backend.mtp_step > 0:
@@ -227,7 +227,7 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int):
         ) = build_dynamic_mtp_linear_att_state_params(
             b_req_idx=self.infer_state.b_req_idx,
             b_mtp_index=self.infer_state.b_mtp_index,
-            req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index,
+            req_to_mtp_state_index=self.infer_state.linear_state_manager.mtp_state_index,
             hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID,
         )
         self._init_mtp_ssm_buffer_idx(mtp_size)
@@ -249,7 +249,7 @@ def _init_fixed_mtp_decode_state(self, draft_step: int):
             device=self.infer_state.b_req_idx.device,
         )
         self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous()
-        self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1
+        self.b_num_accepted_tokens = self.infer_state.linear_state_manager.mtp_state_index[self.b_conv_buffer_idx] + 1
         self._init_mtp_ssm_buffer_idx(mtp_size)
 
     def _init_mtp_ssm_buffer_idx(self, mtp_size: int):
@@ -281,7 +281,7 @@ def decode_att(
         backend: LinearAttBackend = self.backend
 
         mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba)
-        conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num)
+        conv_states, ssm_states = self.infer_state.linear_state_manager.get_layer_buffers(layer_num)
 
         draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model)
         if draft_step > 0:

lightllm/common/basemodel/attention/linear/kda.py

@@ -83,7 +83,7 @@ def prefill_att(
         raw_beta = params["raw_beta"]
         backend: KDALinearAttBackend = self.backend
 
-        conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num)
+        conv_states, ssm_states = self.infer_state.linear_state_manager.get_layer_buffers(layer_num)
         conv_states = conv_states[..., : backend.conv_kernel_size - 1]
         mixed_qkv = causal_conv1d_fn(
             mixed_qkv.transpose(0, 1),
@@ -151,7 +151,7 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int):
         ) = build_dynamic_mtp_linear_att_state_params(
             b_req_idx=self.infer_state.b_req_idx,
             b_mtp_index=self.infer_state.b_mtp_index,
-            req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index,
+            req_to_mtp_state_index=self.infer_state.linear_state_manager.mtp_state_index,
             hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID,
         )
         self._init_mtp_ssm_buffer_idx(mtp_size)
@@ -173,7 +173,7 @@ def _init_fixed_mtp_decode_state(self, mtp_step: int):
             device=self.infer_state.b_req_idx.device,
         )
         self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous()
-        self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1
+        self.b_num_accepted_tokens = self.infer_state.linear_state_manager.mtp_state_index[self.b_conv_buffer_idx] + 1
         self._init_mtp_ssm_buffer_idx(mtp_size)
 
     def _init_mtp_ssm_buffer_idx(self, mtp_size: int):
@@ -204,7 +204,7 @@ def decode_att(
         raw_beta = params["raw_beta"]
         backend: KDALinearAttBackend = self.backend
 
-        conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num)
+        conv_states, ssm_states = self.infer_state.linear_state_manager.get_layer_buffers(layer_num)
         conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx)
         conv_update = causal_conv1d_update
         if backend.mtp_step > 0:

lightllm/common/basemodel/basemodel.py

@@ -16,6 +16,7 @@
 from lightllm.common.kv_cache_mem_manager import MemoryManager
 from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class
 from lightllm.common.req_manager import ReqManager
+from lightllm.common.cache_coordinator import CacheCoordinator
 from lightllm.common.build_utils import repair_config
 from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import (
     select_kv_index_from_req,
@@ -114,12 +115,14 @@ def __init__(self, kvargs):
                 self._init_weights()
         with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE):
             self._init_req_manager()
+            if getattr(self, "is_mtp_draft_model", False):
+                self.cache_coordinator = self.main_model.cache_coordinator
+            else:
+                self._init_cache_coordinator()
             self._init_mem_manager()
 
-        # Qwen3.5 等 linear attention 模型会在 req_manager 中保存大量运行时 state。先初始化
-        # req_manager、再初始化 mem_manager，可以让 KV cache 显存评估包含这些 state 的实际占用；
-        # 因此 req_manager 创建时暂不传入 mem_manager，需要在两者初始化完成后再进行绑定。
         self.req_manager.bind_mem_manager(self.mem_manager)
+        self.cache_coordinator.bind_mem_manager(self.mem_manager)
         self._check_mem_size()
         self._init_infer_layer()
         self._init_some_value()
@@ -196,9 +199,13 @@ def load_weights(self, weight_dict: dict):
         [weight.verify_load() for weight in self.trans_layers_weight]
         return
 
+    @classmethod
+    def get_mem_manager_class(cls):
+        return select_mem_manager_class()
+
     def _init_mem_manager(self):
         assert self.config["num_attention_heads"] % self.tp_world_size_ == 0
-        self.mem_manager: MemoryManager = select_mem_manager_class()(
+        self.mem_manager: MemoryManager = self.get_mem_manager_class()(
             self.max_total_token_num,
             dtype=self.data_type,
             head_num=self.config["num_attention_heads"] // self.tp_world_size_,
@@ -229,6 +236,9 @@ def _check_mem_size(self):
 
         return
 
+    def _init_cache_coordinator(self):
+        self.cache_coordinator = CacheCoordinator()
+
     def _init_req_manager(self):
         create_max_seq_len = 0
 
@@ -376,6 +386,9 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0)
 
         infer_state.mem_manager = self.mem_manager
         infer_state.req_manager = self.req_manager
+        infer_state.state_manager = self.cache_coordinator.state_manager
+        infer_state.linear_state_manager = self.cache_coordinator.linear_state_manager
+        infer_state.indexer_state_manager = self.cache_coordinator.indexer_state_manager
 
         infer_state.mem_index = self._select_mem_indexes(model_input)
         infer_state.microbatch_index = microbatch_index
@@ -631,7 +644,6 @@ def _decode(
 
     @final
     def _context_forward(self, infer_state: InferStateInfo):
-
         input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
         if self.args.enable_dp_prefill_balance:
             assert not self.args.enable_prefill_cudagraph, "not support now"

lightllm/common/basemodel/infer_struct.py

@@ -49,6 +49,8 @@ def __init__(self):
 
         self.mem_manager: MemoryManager = None
         self.req_manager: ReqManager = None
+        self.linear_state_manager = None
+        self.indexer_state_manager = None
 
         self.mem_index: torch.Tensor = None