perf(cache): make CPU cache reads and writes asynchronous

lightllm/common/attention_state/allocators/state_page_allocator.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import List
+from typing import Callable, List
 from collections import deque
 
 
@@ -15,13 +15,17 @@ def __init__(self, size: int, reserved_slots: int = 0):
         self.capacity = self.size - self.reserved_slots
         self._free = deque(range(self.capacity))
         self._allocated = bytearray(self.capacity)
+        self._retained = {}
+        self.on_allocation_pressure: Callable[[StatePageAllocator, int], None] | None = None
 
     @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 and self.on_allocation_pressure is not None:
+            self.on_allocation_pressure(self, num_pages)
         if num_pages > self.free_count:
             return None
         slots = [self._free.popleft() for _ in range(num_pages)]
@@ -35,5 +39,17 @@ def free(self, page_slot_ids: List[int]) -> None:
         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:
+            if slot in self._retained:
+                self._retained[slot] -= 1
+                if self._retained[slot] == 0:
+                    del self._retained[slot]
+                continue
             self._allocated[slot] = 0
-        self._free.extend(page_slot_ids)
+            self._free.append(slot)
+
+    def retain(self, page_slot_ids: List[int]) -> None:
+        """异步复制持有独立引用，原请求释放后仍不能复用对应的 pinned 状态页。"""
+        assert len(set(page_slot_ids)) == len(page_slot_ids), "cannot retain duplicate slots"
+        assert all(0 <= slot < self.capacity and self._allocated[slot] for slot in page_slot_ids)
+        for slot in page_slot_ids:
+            self._retained[slot] = self._retained.get(slot, 0) + 1

lightllm/common/attention_state/bounded_att/gdn/gdn_state.py

@@ -207,10 +207,10 @@ def write_state_pages_to_cpu(
         layout: StateTransferLayout,
         cpu_page_ready_flags: list[bool] | None = None,
     ) -> None:
-        """将指定页池槽位的快照写入 CPU 页。"""
+        """将指定页池槽位的快照写入 CPU 页或 CUDA 暂存页。"""
         assert layout.state_name == self.state_name, "transfer layout does not describe this resource"
         assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval"
-        assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU"
+        assert cpu_page_buffers.device.type in ("cpu", "cuda"), "CPU cache pages require CPU or CUDA staging"
         pool = self._pool(page_kind)
         assert isinstance(page_slot_ids, list), "slot indices must be a list"
         assert isinstance(cpu_page_indices, list), "CPU page indices must be a list"

lightllm/common/attention_state/bounded_att/kda/kda_state.py

@@ -207,10 +207,10 @@ def write_state_pages_to_cpu(
         layout: StateTransferLayout,
         cpu_page_ready_flags: list[bool] | None = None,
     ) -> None:
-        """将指定页池槽位的快照写入 CPU 页。"""
+        """将指定页池槽位的快照写入 CPU 页或 CUDA 暂存页。"""
         assert layout.state_name == self.state_name, "transfer layout does not describe this resource"
         assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval"
-        assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU"
+        assert cpu_page_buffers.device.type in ("cpu", "cuda"), "CPU cache pages require CPU or CUDA staging"
         pool = self._pool(page_kind)
         assert isinstance(page_slot_ids, list), "slot indices must be a list"
         assert isinstance(cpu_page_indices, list), "CPU page indices must be a list"

lightllm/common/attention_state/cpu_cache_operator.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping
 from dataclasses import replace
 from functools import lru_cache
 from typing import TYPE_CHECKING, List
 
 import torch
 
+from .bounded_att.copy_kernel import copy_bounded_att_tensors
 from lightllm.common.attention_state.types import StateTransferLayout
 from lightllm.utils.envs_utils import get_env_start_args
 
 if TYPE_CHECKING:
+    from lightllm.common.basemodel.attention.base_att import BaseAttBackend
     from lightllm.common.attention_state.manager import AttentionStateManager
 
 
@@ -21,8 +23,19 @@ class CpuCacheOperator:
     以及 full-attention token 与 bounded 运行态快照如何写回/读回。
     """
 
+    STAGING_SLOTS = 2
+    STAGING_PAGES = 4
+
     def __init__(self, owner: AttentionStateManager):
         self.owner = owner
+        self.async_cache = False
+
+    def configure_backends(self, backends: Iterable[BaseAttBackend | None]) -> None:
+        """容量规划和实际 staging 分配共用已选定的 main/draft backend 能力。"""
+        self.async_cache = not any(backend.requires_cpu_cache_sync for backend in backends if backend is not None)
+
+    def get_staging_bytes(self) -> int:
+        return self.STAGING_SLOTS * self.STAGING_PAGES * self.get_page_bytes() if self.async_cache else 0
 
     @lru_cache(maxsize=None)
     def get_page_layouts(self) -> Mapping[str, StateTransferLayout]:
@@ -92,13 +105,27 @@ def write_req(
             small_page_readies = cpu_page_ready_flags[big_page_count:] if cpu_page_ready_flags is not None else None
 
         for state in full_states:
-            state.write_to_cpu(
-                token_indices=token_indices,
-                cpu_page_buffers=cpu_page_buffers,
-                cpu_page_indices=cpu_page_indices,
-                layout=layouts[state.state_name],
-                cpu_page_ready_flags=cpu_page_ready_flags,
-            )
+            layout = layouts[state.state_name]
+            if cpu_page_buffers.is_cuda:
+                # 暂存页沿用 CPU 布局；full State 的 CUDA 页写入与 CPU 页共用搬运内核。
+                for index, page in enumerate(cpu_page_indices):
+                    if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]:
+                        continue
+                    start = index * layout.tokens_per_page
+                    state.write_to_pd_page(
+                        token_indices=token_indices[start : start + layout.tokens_per_page],
+                        pd_page_buffers=cpu_page_buffers,
+                        pd_page_index=page,
+                        layout=layout,
+                    )
+            else:
+                state.write_to_cpu(
+                    token_indices=token_indices,
+                    cpu_page_buffers=cpu_page_buffers,
+                    cpu_page_indices=cpu_page_indices,
+                    layout=layout,
+                    cpu_page_ready_flags=cpu_page_ready_flags,
+                )
 
         for state in bounded_states:
             layout = layouts.get(state.state_name)
@@ -125,6 +152,33 @@ def write_req(
                     cpu_page_ready_flags=small_page_readies,
                 )
 
+    def scatter_staging_pages(
+        self, staging: torch.Tensor, cpu_page_buffers: torch.Tensor, cpu_page_indices: List[int]
+    ) -> None:
+        """只写当前 TP rank 负责的字段，避免暂存页的空白区域覆盖其他 rank。"""
+        for layout in self.get_page_layouts().values():
+            for region in layout.regions:
+                if not region.write_enabled:
+                    continue
+                slices = [slice(None)] * len(region.axes)
+                for mapping in region.axis_mappings:
+                    assert not mapping.external_circular
+                    slices[region.axes.index(mapping.axis)] = slice(
+                        mapping.external_start, mapping.external_start + mapping.length
+                    )
+                copies = []
+                for index, page in enumerate(cpu_page_indices):
+                    source = staging[index, region.byte_start : region.byte_end].view(region.dtype).view(region.shape)
+                    target = (
+                        cpu_page_buffers[page]
+                        .view(torch.uint8)
+                        .view(-1)[region.byte_start : region.byte_end]
+                        .view(region.dtype)
+                        .view(region.shape)
+                    )
+                    copies.append((source[tuple(slices)], target[tuple(slices)]))
+                copy_bounded_att_tensors(copies)
+
     def read_req(
         self,
         cpu_page_buffers: torch.Tensor,

lightllm/common/attention_state/manager.py

@@ -202,12 +202,13 @@ def initialize_storage(
         pd_reserved_bytes = 0
         if args.run_mode in ("prefill", "decode"):
             pd_reserved_bytes = self.get_pd_operator().get_page_bytes() * args.pd_kv_page_num
+        cpu_staging_bytes = self.get_cpu_cache_operator().get_staging_bytes() if args.enable_cpu_cache else 0
         if token_capacity is None:
             # 权重和固定运行态已分配，用当前设备剩余显存扣除保留比例及 PD buffer。
             torch.cuda.empty_cache()
             gpu_device = torch.device("cuda", get_current_device_id())
             free_bytes, total_bytes = torch.cuda.mem_get_info(gpu_device)
-            budget = int(free_bytes - total_bytes * (1 - mem_fraction)) - pd_reserved_bytes
+            budget = int(free_bytes - total_bytes * (1 - mem_fraction)) - pd_reserved_bytes - cpu_staging_bytes
 
             # 按每页开销直接计算容量，预留一个 HOLD 页；静态量化参数占比较小，不单独扣减。
             assert args.page_size > 0, "capacity planning requires a positive page size"
@@ -224,6 +225,7 @@ def initialize_storage(
             logger.info(
                 f"Attention storage budget: {budget / 1024**3:.3f} GiB, "
                 f"transfer reserve: {pd_reserved_bytes / 1024**2:.3f} MiB, "
+                f"CPU cache staging reserve: {cpu_staging_bytes / 1024**2:.3f} MiB, "
                 f"profiled token capacity: {token_capacity} (mem_fraction={mem_fraction})"
             )
         token_capacity = token_capacity // args.page_size * args.page_size

lightllm/common/basemodel/attention/base_att.py

@@ -20,6 +20,8 @@ class BaseAttBackend:
     每个 model 复用一个 backend 实例。
     """
 
+    requires_cpu_cache_sync = False
+
     _instances = {}
     _workspace_buffers = {}
     _workspace_buffer_lock = threading.Lock()

lightllm/common/basemodel/attention/fa3/fp.py

@@ -16,6 +16,7 @@
 class Fa3AttBackend(BaseAttBackend):
     """Common fixed page-table storage for FA3 attention backends."""
 
+    requires_cpu_cache_sync = True
     page_table_buffers = None
 
     def __init__(self, model):

lightllm/common/basemodel/basemodel.py

@@ -179,11 +179,12 @@ def prepare_attention_states(self) -> None:
                 self.att_state_manager.register_state(self.att_state_role, state)
 
     def init_runtime(self) -> None:
-        """Initialize this model's backends, warmup and graphs after shared allocation."""
+        """Finish runtime initialization, warmup and graphs after shared allocation."""
         self.req_manager.bind_hold_page(self.att_state_manager.full_att_allocator.hold_token_indices)
         self._check_mem_size()
-        self._init_att_backend()
-        self._init_att_backend1()
+        if not self.args.enable_cpu_cache:
+            self._init_att_backend()
+            self._init_att_backend1()
         self._init_hidden_collector()
         role = self.att_state_role
         logger.info(

lightllm/server/multi_level_kv_cache/cpu_cache_client.py

@@ -1,5 +1,6 @@
 import ctypes
 import torch
+from dataclasses import dataclass
 from lightllm.utils.envs_utils import get_env_start_args, get_disk_cache_prompt_limit_length
 from typing import List, Optional, Tuple
 from lightllm.utils.log_utils import init_logger
@@ -11,6 +12,13 @@
 logger = init_logger(__name__)
 
 
+@dataclass
+class CpuPageAllocation:
+    page_indices: List[int]
+    ready_flags: List[bool]
+    writer_flags: List[bool]
+
+
 class CpuKvCacheClient(object):
     """
     This class is responsible for handling cpu kv cache meta data.
@@ -85,28 +93,36 @@ def allocate_one_page(
                 return None, False
 
     def allocate_pages(self, hash_keys: List[int], disk_offload_enable: bool) -> Tuple[List[int], List[bool]]:
-        """
-        allocate_pages will add _CpuPageStaus ref_count
-        """
+        """保持原接口：增加页面引用，返回页面和 READY 标记；调用方持有 cache lock。"""
+        allocation = self.allocate_pages_for_write(hash_keys, disk_offload_enable)
+        return allocation.page_indices, allocation.ready_flags
+
+    def allocate_pages_for_write(self, hash_keys: List[int], disk_offload_enable: bool) -> CpuPageAllocation:
+        """调用方持有 cache lock；增加引用，只对新建页授予 writer，READY 与 writer 分别返回。"""
         page_list = []
         ready_list = []
+        writer_list = []
         page_items = self.page_items.linked_items
         for hash_key in hash_keys:
+            is_owner = self.page_hash_dict.get(hash_key) is None
             page_index, ready = self.allocate_one_page(
                 page_items=page_items, hash_key=hash_key, disk_offload_enable=disk_offload_enable
             )
             if page_index is not None:
                 page_list.append(page_index)
                 ready_list.append(ready)
+                writer_list.append(is_owner)
             else:
                 page_list.append(-1)
                 ready_list.append(False)
+                writer_list.append(False)
                 break
 
         left_num = len(hash_keys) - len(page_list)
         page_list.extend([-1 for _ in range(left_num)])
         ready_list.extend([False for _ in range(left_num)])
-        return page_list, ready_list
+        writer_list.extend([False for _ in range(left_num)])
+        return CpuPageAllocation(page_list, ready_list, writer_list)
 
     def update_pages_status_to_ready(
         self,

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -170,6 +170,21 @@ def init_model(self, kvargs):
         with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE):
             manager.finish_state_registration()
             manager.load_full_att_calibration(args=self.model.args)
+            if self.args.enable_cpu_cache:
+                # 后端只依赖模型配置和已注册 State；先确定能力，再规划 CPU staging 预算。
+                for item in models:
+                    item._init_att_backend()
+                    item._init_att_backend1()
+                manager.get_cpu_cache_operator().configure_backends(
+                    backend
+                    for item in models
+                    for backend in (
+                        item.prefill_att_backend,
+                        item.decode_att_backend,
+                        item.prefill_att_backend1,
+                        item.decode_att_backend1,
+                    )
+                )
             manager.initialize_storage(
                 self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction
             )
@@ -730,7 +745,7 @@ def _get_classed_reqs(
         # 定期对 radix cache 进行 merge，防止查询插入的操作效率下降
         self._timer_merge_radix_tree()
 
-        if self.args.enable_cpu_cache and len(g_infer_context.infer_req_ids) > 0:
+        if self.args.enable_cpu_cache:
             self.multi_level_cache_module.update_cpu_cache_task_states()
 
         if req_ids is None:
@@ -859,6 +874,8 @@ def _get_classed_reqs(
         else:
             true_finished_reqs = finished_reqs
 
+        if self.args.enable_cpu_cache and self.multi_level_cache_module.async_cache:
+            self.multi_level_cache_module.close_cache_sessions(true_finished_reqs)
         g_infer_context.filter_reqs(finished_reqs=true_finished_reqs)
         g_infer_context.pause_reqs(wait_pause_reqs, is_master_in_dp=self.is_master_in_dp)
 
@@ -887,10 +904,14 @@ def _get_classed_reqs(
     # 一些可以复用的通用功能函数
     def _pre_post_handle(self, run_reqs: List[InferReq], is_chuncked_mode: bool) -> List[InferReqUpdatePack]:
         update_func_objs: List[InferReqUpdatePack] = []
+        cpu_store_reqs = []
         # 通用状态预先填充
         is_master_in_dp = self.is_master_in_dp
         for req_obj in run_reqs:
             req_obj: InferReq = req_obj
+            if self.args.enable_cpu_cache and self.multi_level_cache_module.incremental_offload:
+                if req_obj.cur_kv_len < req_obj.shm_req.input_len:
+                    cpu_store_reqs.append(req_obj)
             if is_chuncked_mode:
                 new_kv_len = req_obj.get_chuncked_input_token_len()
             else:
@@ -910,6 +931,10 @@ def _pre_post_handle(self, run_reqs: List[InferReq], is_chuncked_mode: bool) ->
             req_obj.cur_output_len += 1
             pack = InferReqUpdatePack(req_obj=req_obj, output_len=req_obj.cur_output_len)
             update_func_objs.append(pack)
+        if cpu_store_reqs:
+            self.multi_level_cache_module.store_completed_prefill_pages(
+                cpu_store_reqs, producer_stream=g_infer_context.get_overlap_stream()
+            )
         return update_func_objs
 
     # 一些可以复用的通用功能函数