refactor(att_resources): unify attention state and cache transfers

docs/CN/source/index.rst

@@ -81,6 +81,7 @@ Lightllm 整合了众多的开源方案的优点，包括但不限于 FasterTran
    :caption: 架构介绍
 
    架构介绍 <framework/framework>
+   Attention Resources 架构设计 <framework/cache_manager_design>
    token attention介绍 <framework/token_attention>
    峰值显存调度器介绍 <framework/router>
 

docs/CN/source/models/add_new_model.md

@@ -109,7 +109,6 @@ class TpPartBaseModel:
         self._verify_must()
         self._verify_params()
         self._init_weights()
-        self._init_mem_manager()
         self._init_infer_layer()
         self._init_some_value()
         self._init_custom()
@@ -124,10 +123,14 @@ class TpPartBaseModel:
 | ---------------------------- | ------------------------------------------------------------ |
 | def _init_config(self)：     | 读取初始化模型的 config.json,  并进行一些 key 名的同名合法化操作 |
 | def _verify_params(self)：   | 校验参数                                                     |
-| def _init_mem_manager(self): | 初始化 token attention 使用的 mem manager 对象               |
+| `create_full_att_states(...)` / `create_bounded_att_states(...)` | 创建并返回当前模型的 full / bounded State 列表；默认返回空列表，公共流程统一注册 |
 | def _init_some_value(self):  | 初始化推理框架会使用的一些成员变量的值                       |
 | def _init_custom(self):      | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值  |
 
+全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先按各模型 `get_req_manager_max_seq_len()` 的最大值创建一份 `ReqManager`，通过 `bind_req_manager()` 绑定给全部模型，再调用各模型的 `prepare_attention_states()` 准备自身 State；随后管理器依次完成 `finish_state_registration()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。资源准备和运行时初始化分别通过这两个单模型公共入口完成，均不接收 draft 列表。
+
+公共 `prepare_attention_states()` 为 main 创建空管理器，draft 引用 main 的实例。两个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。
+
 ## 2. 添加 bloom 模型的示例说明
 
 具体实现在 ***lightllm/models/bloom*** 目录下，下面的代码片段请对应源码进行阅读，其中 triton_kernel 目录下为推理类使用的一些 kernel，下文中不做详细介绍，同时 bloom 模型因为不需要传递特殊状态信息使用默认的状态类即可。如想更深入的理解整个框架，可以进一步参考 llama 和 llama2 等模型的接入实现源码。
@@ -319,9 +322,6 @@ class BloomPreLayerInfer(PreLayerInferTpl):
         self.vob_end_id_ = tp_vocab_size_ * (self.tp_rank_ + 1)
         return
 
-    def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor:
-        return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_)
-
     def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BloomPreAndPostLayerWeight):
         total_token_num = infer_state.total_token_num
         input_ids = input_ids[0:total_token_num]
@@ -346,6 +346,9 @@ class BloomPreLayerInfer(PreLayerInferTpl):
             dist.all_reduce(input_embdings, op=dist.ReduceOp.SUM, async_op=False)
         input_embdings = self._norm(input_embdings, infer_state, layer_weight)
         return input_embdings
+
+    def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor:
+        return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_)
 ~~~
 
 ***transformer_layer_infer.py***
@@ -416,9 +419,10 @@ class BloomTransformerLayerInfer(TransformerLayerInferTpl):
 
     def _token_attention_kernel(self, q, infer_state:InferStateInfo, layer_weight: BloomTransformerLayerWeight)->torch.Tensor:
         o_tensor = torch.empty_like(q)
+        k_cache, v_cache = infer_state.full_att_state.get_layer_view(self.layer_num_)
         token_attention_fwd(q.view(-1, self.tp_q_head_num_, self.head_dim_),
-                            infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0: self.tp_k_head_num_, :],
-                            infer_state.mem_manager.kv_buffer[self.layer_num_][:, self.tp_k_head_num_: self.tp_k_head_num_+ self.tp_v_head_num_, :],
+                            k_cache,
+                            v_cache,
                             o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
                             layer_weight.tp_alibi,
                             infer_state.b_loc,
@@ -470,9 +474,6 @@ class BloomPostLayerInfer(PostLayerInferTpl):
         self.embed_dim_ = network_config["n_embed"]
         return
 
-    def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor:
-        return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_)
-
     def soft_max(self, data):
         return torch.softmax(data.permute(1, 0).float(), dim=-1)
 
@@ -507,6 +508,9 @@ class BloomPostLayerInfer(PostLayerInferTpl):
             ans_logics = gather_data.permute(1, 0).float()
             gather_data = None
             return ans_logics
+
+    def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor:
+        return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_)
 ~~~
 
 ### （3） 实现模型的框架类

docs/EN/source/models/add_new_model.md

@@ -109,7 +109,6 @@ class TpPartBaseModel:
         self._verify_must()
         self._verify_params()
         self._init_weights()
-        self._init_mem_manager()
         self._init_infer_layer()
         self._init_some_value()
         self._init_custom()
@@ -124,10 +123,14 @@ Commonly used interfaces that need to be inherited and implemented
 | ---------------------------- | ------------------------------------------------------------ |
 | def _init_config(self)：     | Read the config.json for initializing the model and perform some key name legalization operations |
 | def _verify_params(self)：   | Validate parameters                                           |
-| def _init_mem_manager(self): | Initialize the mem manager object used by token attention    |
+| `create_full_att_states(...)` / `create_bounded_att_states(...)` | Return this model's full / bounded State lists; defaults return empty lists and the shared startup stage registers the results |
 | def _init_some_value(self):  | Initialize values of some member variables that the inference framework will use |
 | def _init_custom(self):      | Some personalized initialization of the model itself, such as llama initializing its own Rotary values |
 
+After all main/draft models are constructed, `ModeBackend.init_model` creates one `ReqManager` using the largest `get_req_manager_max_seq_len()` returned by the models, binds it to every model through `bind_req_manager()`, then calls each model's `prepare_attention_states()`, then calls `finish_state_registration()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization each use a single-model entry point; neither takes a draft-model list.
+
+The public `prepare_attention_states()` stage creates an empty manager for main and shares it with drafts. The two State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded factories also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none.
+
 ## 2. Example of Adding Bloom Model
 
 The specific implementation is under the ***lightllm/models/bloom*** directory. Please read the source code for the code snippets below. The triton_kernel directory contains some kernels used by inference classes, which will not be introduced in detail in this article. At the same time, the bloom model uses the default state class because it doesn't need to pass special state information. For a deeper understanding of the entire framework, you can further refer to the source code implementation of llama and llama2 model integration.
@@ -320,9 +323,6 @@ class BloomPreLayerInfer(PreLayerInferTpl):
         self.vob_end_id_ = tp_vocab_size_ * (self.tp_rank_ + 1)
         return
 
-    def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor:
-        return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_)
-
     def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BloomPreAndPostLayerWeight):
         total_token_num = infer_state.total_token_num
         input_ids = input_ids[0:total_token_num]
@@ -347,6 +347,9 @@ class BloomPreLayerInfer(PreLayerInferTpl):
             dist.all_reduce(input_embdings, op=dist.ReduceOp.SUM, async_op=False)
         input_embdings = self._norm(input_embdings, infer_state, layer_weight)
         return input_embdings
+
+    def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor:
+        return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_)
 ~~~
 
 ***transformer_layer_infer.py***
@@ -417,9 +420,10 @@ class BloomTransformerLayerInfer(TransformerLayerInferTpl):
 
     def _token_attention_kernel(self, q, infer_state:InferStateInfo, layer_weight: BloomTransformerLayerWeight)->torch.Tensor:
         o_tensor = torch.empty_like(q)
+        k_cache, v_cache = infer_state.full_att_state.get_layer_view(self.layer_num_)
         token_attention_fwd(q.view(-1, self.tp_q_head_num_, self.head_dim_),
-                            infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0: self.tp_k_head_num_, :],
-                            infer_state.mem_manager.kv_buffer[self.layer_num_][:, self.tp_k_head_num_: self.tp_k_head_num_+ self.tp_v_head_num_, :],
+                            k_cache,
+                            v_cache,
                             o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_),
                             layer_weight.tp_alibi,
                             infer_state.b_loc,
@@ -471,9 +475,6 @@ class BloomPostLayerInfer(PostLayerInferTpl):
         self.embed_dim_ = network_config["n_embed"]
         return
 
-    def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor:
-        return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_)
-
     def soft_max(self, data):
         return torch.softmax(data.permute(1, 0).float(), dim=-1)
 
@@ -508,6 +509,9 @@ class BloomPostLayerInfer(PostLayerInferTpl):
             ans_logics = gather_data.permute(1, 0).float()
             gather_data = None
             return ans_logics
+
+    def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor:
+        return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_)
 ~~~
 
 ### （3） Add implementation model class

lightllm/common/attention_state/__init__.py

@@ -0,0 +1,5 @@
+"""Attention data resources, independent of request and cache policy."""
+
+from .manager import AttentionStateManager
+
+__all__ = ["AttentionStateManager"]

lightllm/common/attention_state/allocators/__init__.py

@@ -0,0 +1,4 @@
+from .state_page_allocator import StatePageAllocator
+from .full_att_kv_allocator import KVCapacityReader, FullAttKVAllocator
+
+__all__ = ["StatePageAllocator", "KVCapacityReader", "FullAttKVAllocator"]

lightllm/common/attention_state/allocators/full_att_kv_allocator.py

@@ -0,0 +1,110 @@
+"""Full attention KV 分配器：接口使用 token 数量和索引，分配与回收须满足整页约束。"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+import numpy as np
+import torch
+
+from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
+from lightllm.utils.dist_utils import get_current_rank_in_node
+from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.log_utils import init_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+logger = init_logger(__name__)
+
+
+class FullAttKVAllocator:
+    def __init__(self, size: int, page_size: int = 1, pin_memory: bool = True):
+        assert size >= 0 and page_size > 0 and size % page_size == 0, "KV capacity must contain complete pages"
+        assert size + page_size - 1 <= np.iinfo(np.int32).max, "KV token indices including HOLD must fit in int32"
+        self.size = size
+        self.page_size = page_size
+        self.mem_state = torch.arange(
+            end=self.size,
+            dtype=torch.int32,
+            device="cpu",
+            pin_memory=pin_memory,
+        )
+        self.mark_start = 0
+        self.mark_end = self.size
+        self._mem_state_return = torch.empty(
+            size=(self.size * 3,),
+            dtype=torch.int32,
+            device="cpu",
+            pin_memory=pin_memory,
+        )
+        self._return_start = 0
+        self.can_use_mem_size = self.size
+        rank_in_node = get_current_rank_in_node()
+        self.shared_can_use_token_num = SharedInt(f"kv_available_tokens_{rank_in_node}")
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
+
+    @property
+    def hold_token_indices(self) -> tuple[int, ...]:
+        """One reserved page after the ordinary allocation domain."""
+        return tuple(range(self.size, self.size + self.page_size))
+
+    def alloc(self, need_size: int) -> torch.Tensor:
+        assert need_size >= 0 and need_size % self.page_size == 0, "allocation must contain complete KV pages"
+        assert (
+            need_size <= self.mark_end - self.mark_start
+        ), f"need {need_size} KV tokens, only {self.can_use_mem_size} available"
+        start = self.mark_start
+        end = start + need_size
+        self.mark_start = end
+        self.can_use_mem_size -= need_size
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
+
+        # Reuse a rolling buffer so frees do not overwrite indices
+        # still being consumed by asynchronous H2D copies.
+        if self._return_start + need_size > self._mem_state_return.shape[0]:
+            self._return_start = 0
+        ans = self._mem_state_return[self._return_start : self._return_start + need_size]
+        ans.copy_(self.mem_state[start:end])
+        self._return_start += need_size
+        return ans
+
+    def free(self, token_indices: Sequence[int] | torch.Tensor) -> None:
+        num_tokens = len(token_indices)
+        if num_tokens == 0:
+            return
+        assert num_tokens % self.page_size == 0, "free must contain complete KV pages"
+        end = self.mark_start
+        start = end - num_tokens
+        assert start >= 0, f"free {num_tokens} KV tokens exceeds allocated {end}"
+        if isinstance(token_indices, torch.Tensor):
+            self.mem_state[start:end].copy_(token_indices)
+        else:
+            self.mem_state.numpy()[start:end] = token_indices
+        self.mark_start = start
+        self.can_use_mem_size += num_tokens
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
+        if self.can_use_mem_size == self.size:
+            logger.debug(f"freed all gpu mem size {self.can_use_mem_size}")
+
+    def free_all(self) -> None:
+        """Reset after all users of this allocation domain have finished."""
+        self.mem_state.numpy()[:] = np.arange(self.size, dtype=np.int32)
+        self.mark_start = 0
+        self.mark_end = self.size
+        self.can_use_mem_size = self.size
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
+
+
+class KVCapacityReader:
+    """Read each local DP group's allocator capacity without modifying it."""
+
+    def __init__(self):
+        args = get_env_start_args()
+        node_world_size = args.tp // args.nnodes
+        tp_group_size = args.tp // args.dp
+        self.is_multinode_tp = args.dp == 1 and args.nnodes > 1
+        self.counters = [SharedInt(f"kv_available_tokens_{rank}") for rank in range(0, node_world_size, tp_group_size)]
+
+    def get_available_tokens(self, dp_rank_in_node: int) -> int:
+        return self.counters[0 if self.is_multinode_tp else dp_rank_in_node].get_value()

lightllm/common/attention_state/allocators/state_page_allocator.py

@@ -0,0 +1,39 @@
+"""状态页槽位分配器，仅管理槽位，不持有状态 buffer 或缓存策略。"""
+
+from __future__ import annotations
+
+from typing import List
+from collections import deque
+
+
+class StatePageAllocator:
+    def __init__(self, size: int, reserved_slots: int = 0):
+        self.size = size
+        self.reserved_slots = reserved_slots
+        if not 0 <= self.reserved_slots <= self.size:
+            raise ValueError("reserved_slots must be between zero and capacity")
+        self.capacity = self.size - self.reserved_slots
+        self._free = deque(range(self.capacity))
+        self._allocated = bytearray(self.capacity)
+
+    @property
+    def free_count(self) -> int:
+        return len(self._free)
+
+    def alloc(self, num_pages: int) -> list[int] | None:
+        assert num_pages >= 0, "num_pages must be nonnegative"
+        if num_pages > self.free_count:
+            return None
+        slots = [self._free.popleft() for _ in range(num_pages)]
+        for slot in slots:
+            self._allocated[slot] = 1
+        return slots
+
+    def free(self, page_slot_ids: List[int]) -> None:
+        assert isinstance(page_slot_ids, list), "slot indices must be a list"
+        assert len(set(page_slot_ids)) == len(page_slot_ids), "cannot free duplicate slots"
+        assert all(0 <= slot < self.capacity for slot in page_slot_ids), "slot outside the allocatable domain"
+        assert all(self._allocated[slot] for slot in page_slot_ids), "cannot free an unallocated slot"
+        for slot in page_slot_ids:
+            self._allocated[slot] = 0
+        self._free.extend(page_slot_ids)