diff --git a/lightllm/common/attention_state/allocators/state_page_allocator.py b/lightllm/common/attention_state/allocators/state_page_allocator.py index 9096bd9a96..d9130bd567 100644 --- a/lightllm/common/attention_state/allocators/state_page_allocator.py +++ b/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,6 +15,8 @@ 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: @@ -22,6 +24,8 @@ def free_count(self) -> int: 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 diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index b8a641a73d..52e87cf525 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/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" diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index e5480a3ab0..6aa1f5d904 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/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" diff --git a/lightllm/common/attention_state/cpu_cache_operator.py b/lightllm/common/attention_state/cpu_cache_operator.py index 814d340481..1968de4748 100644 --- a/lightllm/common/attention_state/cpu_cache_operator.py +++ b/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, diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 6fb18c7ea1..989309e0f0 100644 --- a/lightllm/common/attention_state/manager.py +++ b/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 diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index 1122df3551..ae37385e68 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/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() diff --git a/lightllm/common/basemodel/attention/fa3/fp.py b/lightllm/common/basemodel/attention/fa3/fp.py index 6fcee4312a..40d9194e8a 100644 --- a/lightllm/common/basemodel/attention/fa3/fp.py +++ b/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): diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 0edbdadb94..350c031d6f 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/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( diff --git a/lightllm/server/multi_level_kv_cache/cpu_cache_client.py b/lightllm/server/multi_level_kv_cache/cpu_cache_client.py index 8138b798f6..d2ee94467c 100644 --- a/lightllm/server/multi_level_kv_cache/cpu_cache_client.py +++ b/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, diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 7997f49934..7bd81358be 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/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 # 一些可以复用的通用功能函数 diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 261f090b0b..d3664dcdf2 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -2,10 +2,11 @@ import torch import dataclasses import bisect -from functools import lru_cache +from itertools import islice from typing import Optional, List, Deque from collections import deque from lightllm.server.multi_level_kv_cache import CacheTier +from lightllm.server.multi_level_kv_cache.cache_placement_controller import LegacyCachePlacementController from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient from lightllm.utils.envs_utils import get_env_start_args from ..infer_batch import InferReq @@ -18,6 +19,8 @@ class MultiLevelKvCacheModule(object): + async_cache = False + def __init__(self, backend): self.args = get_env_start_args() assert self.args.cpu_cache_token_page_size % self.args.page_size == 0 @@ -30,39 +33,54 @@ def __init__(self, backend): dist.barrier(group=self.init_sync_group) self.offload_sync_group = create_new_group_for_current_dp("nccl") dist.barrier(group=self.offload_sync_group) - self.offload_sync_tensor = torch.empty((1,), dtype=torch.int32, device="cuda") + self.offload_sync_tensor = torch.zeros((1,), dtype=torch.int32, device="cuda") + self.load_sync_tensor = torch.zeros((1,), dtype=torch.int32, device="cuda") self.cpu_cache_handle_queue: Deque[TransTask] = deque() + self.load_tasks = deque() + self.checkpoint_tasks = deque() + self.store_tasks = deque() + self.cache_sessions = {} + self.staging_slots = [] self.cpu_cache_client = CpuKvCacheClient(only_create_meta_data=False, init_shm_data=False) resources = self.backend.model.att_state_manager self.cpu_cache_operator = resources.get_cpu_cache_operator() meta = self.cpu_cache_client.kv_cache_tensor_meta if meta.page_bytes != self.cpu_cache_operator.get_page_bytes(): raise ValueError("CPU cache service and model resource page sizes differ") + self.async_cache = self.cpu_cache_operator.async_cache + logger.info("MultiLevelKvCacheModule: async CPU cache=%s", self.async_cache) + self.incremental_offload = self.async_cache and isinstance( + g_infer_context.cache_placement_controller, LegacyCachePlacementController + ) + if self.async_cache: + for allocator in (resources.bounded_att_big_page_allocator, resources.bounded_att_small_page_allocator): + if allocator is not None: + allocator.on_allocation_pressure = self._reclaim_checkpoint_slots + self.staging_slots = [ + StagingSlot( + torch.empty( + (self.cpu_cache_operator.STAGING_PAGES, meta.page_bytes), dtype=torch.uint8, device="cuda" + ) + ) + for _ in range(self.cpu_cache_operator.STAGING_SLOTS) + ] - @lru_cache() def need_sync_compute_stream(self) -> bool: - """ - fa3 在 offload 和 load kv cache 的时候，需要等待计算流完成，否则可能会概率崩溃。 - """ + return not self.cpu_cache_operator.async_cache - model = self.backend.model - att_backends = [ - model.prefill_att_backend, - model.decode_att_backend, - model.prefill_att_backend1, - model.decode_att_backend1, - ] - for att_backend in att_backends: - if att_backend is not None and "fa3" in att_backend.__class__.__name__.lower(): - logger.info("MultiLevelKvCacheModule: need sync compute stream for fa3 backend.") - return True - logger.info("MultiLevelKvCacheModule: no need sync compute stream.") - return False + def _reclaim_checkpoint_slots(self, allocator, num_pages: int) -> None: + # 各 TP rank 的队列和分配顺序一致；只在槽位不足时等待必要的源快照搬运。 + while allocator.free_count < num_pages and self.checkpoint_tasks: + task = self.checkpoint_tasks.popleft() + task.event.synchronize() + for owner, slots in task.leases: + owner.free(slots) def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): idle_token_num = g_infer_context.get_can_alloc_token_num() all_page_list = [] + checkpoint_leases = [] is_master_in_dp = self.backend.is_master_in_dp for req in reqs: page_list = req.shm_req.cpu_cache_match_page_indexes.get_all() @@ -111,27 +129,44 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): assert ready_page_num <= matched_page_num need_pages = page_list[ready_page_num:matched_page_num] - self.read_req_from_cpu( + slots = self.read_req_from_cpu( req=req, cpu_page_indices=need_pages, page_lengths=page_len_list[ready_page_num:matched_page_num], start_token_len=page_len_list[ready_page_num - 1] if ready_page_num else 0, ) + if self.async_cache and slots: + allocator = self.backend.model.att_state_manager.bounded_att_big_page_allocator + allocator.retain(slots) + checkpoint_leases.append((allocator, slots)) idle_token_num -= need_token_num - torch.cuda.current_stream().synchronize() + if not self.async_cache: + torch.cuda.current_stream().synchronize() if self.backend.is_master_in_dp: req.shm_req.shm_cur_kv_len = req.cur_kv_len all_page_list.extend(page_list) - dist.barrier(group=self.init_sync_group) - - if self.backend.is_master_in_dp: - self.cpu_cache_client.lock.acquire_sleep1ms() - self.cpu_cache_client.deref_pages(page_list=all_page_list) - self.cpu_cache_client.lock.release() + if self.async_cache: + # 同一计算流上的后续 forward 自然等待恢复；CPU 页与状态槽位由完成事件持有。 + if all_page_list: + if self.backend.dp_world_size > 1: + dist.all_reduce(self.load_sync_tensor, group=self.init_sync_group) + event = torch.cuda.Event() + event.record() + self.load_tasks.append(LoadTask(all_page_list, event)) + if checkpoint_leases: + self.checkpoint_tasks.append(CheckpointTask(event, checkpoint_leases)) + for req in reqs: + self.cache_sessions[req.req_id] = CacheSession(req.req_id) + else: + dist.barrier(group=self.init_sync_group) + if self.backend.is_master_in_dp: + self.cpu_cache_client.lock.acquire_sleep1ms() + self.cpu_cache_client.deref_pages(page_list=all_page_list) + self.cpu_cache_client.lock.release() return def read_req_from_cpu( @@ -140,7 +175,7 @@ def read_req_from_cpu( cpu_page_indices: List[int], page_lengths: List[int], start_token_len: int, - ) -> None: + ) -> Optional[List[int]]: """恢复已选定的请求前缀页面，分配资源并登记完整大页引用。 start_token_len 是首个 CPU 页的逻辑起点，可能早于已有 GPU 命中长度。 @@ -193,13 +228,21 @@ def read_req_from_cpu( if big_page_slot_ids is not None: req.hybrid_len_to_big_page_id.update(zip(complete_page_lengths, big_page_slot_ids)) req.cur_kv_len = page_lengths[-1] + return big_page_slot_ids def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> List[InferReq]: """ 将满足cpu kv cache 卸载条件的请求进行处理, 并返回真的满足退出条件的请求list。 """ - # 如果开启了cpu cache，将达到finished状态的请求开启将gpu kv cache 卸载到 cpu cache中的操作。 - # 当 kv cache 卸载完成后，才会进行请求的真实退出操作。 + # 异步路径将源数据打包到独立暂存区后即可退出；FA3 保留等待卸载完成的保护。 + if self.async_cache: + self.store_completed_prefill_pages(finished_reqs, g_infer_context.get_overlap_stream()) + for req in finished_reqs: + self.cache_sessions[req.req_id].closing = True + # 请求不再持有写出任务；CPU 页是否 READY 由独立任务的完成事件决定。 + req.cpu_cache_task_status = InferReq._CpuCacheTaskStatus.FINISHED + return finished_reqs + true_finished_reqs = [] cpu_stream = g_infer_context.get_cpu_kv_cache_stream() for req in finished_reqs: @@ -250,6 +293,213 @@ def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> L return true_finished_reqs + def store_completed_prefill_pages(self, reqs: List[InferReq], producer_stream: torch.cuda.Stream): + """按已完成的 prompt 页增量写出；任务持有独立暂存区，不依赖请求的退出时机。""" + resources = self.backend.model.att_state_manager + page_tokens = self.args.cpu_cache_token_page_size + requests = [] + for req in reqs: + session = self.cache_sessions[req.req_id] + if session.closing or req.shm_req.group_req_id != req.shm_req.request_id: + continue + offload_limit = self.args.linear_att_hash_page_size if self.backend.is_hybrid_att_model else page_tokens + if req.shm_req.input_len <= offload_limit: + continue + page_lens = req.shm_req.token_hash_page_len_list.get_all() + target = bisect.bisect_right(page_lens, req.cur_kv_len) + target = self._handle_hybrid_att_last_page(req, target, page_lens) + if target > session.page_count: + requests.append((req, session, page_lens, target)) + if not requests: + return + + # 一轮内批量分配并广播；只根据各 rank 一致的请求进度决定是否进入 collective。 + if self.backend.is_master_in_dp: + allocations = [] + self.cpu_cache_client.lock.acquire_sleep1ms() + try: + for req, session, _, target in requests: + hashes = req.shm_req.token_hash_list.get_all() + allocation = self.cpu_cache_client.allocate_pages_for_write( + hashes[session.page_count : target], disk_offload_enable=False + ) + if -1 in allocation.page_indices: + count = allocation.page_indices.index(-1) + allocation.page_indices = allocation.page_indices[:count] + allocation.ready_flags = allocation.ready_flags[:count] + allocation.writer_flags = allocation.writer_flags[:count] + session.disk_offload = CacheTier.DISK in req.cache_tiers or ( + self.incremental_offload and self.args.enable_disk_cache + ) + if not session.disk_offload: + # 已有页无需本请求写入；其 writer 或 READY 状态拥有独立生命周期。 + self.cpu_cache_client.deref_pages( + [ + page + for page, writer in zip(allocation.page_indices, allocation.writer_flags) + if not writer + ] + ) + allocations.append(allocation) + finally: + self.cpu_cache_client.lock.release() + data = [allocations] + else: + data = [None] + dist.broadcast_object_list(data, group=self.gloo_group, group_src=0) + for (req, session, page_lens, _), allocation in zip(requests, data[0]): + pages, owners = allocation.page_indices, allocation.writer_flags + start = session.page_count + session.page_count += len(pages) + session.disk_offload = CacheTier.DISK in req.cache_tiers or ( + self.incremental_offload and self.args.enable_disk_cache + ) + if session.disk_offload: + session.disk_pages.extend(pages) + if pages: + session.token_num = page_lens[session.page_count - 1] + shared_slots = [] + if self.backend.is_hybrid_att_model and req.shared_kv_node is not None: + shared_slots = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) + for offset in range(0, len(pages), self.cpu_cache_operator.STAGING_PAGES): + selected = [ + (index, page) + for index, (page, owner) in enumerate( + zip( + pages[offset : offset + self.cpu_cache_operator.STAGING_PAGES], + owners[offset : offset + self.cpu_cache_operator.STAGING_PAGES], + ), + start + offset, + ) + if owner + ] + if not selected: + continue + slot = next((slot for slot in self.staging_slots if not slot.in_use), None) + if slot is None: + # 暂存区满才回收/施加背压，不对每个请求或每组页面执行 object collective。 + self.store_tasks[0].event.synchronize() + self._poll_async_tasks() + slot = next(slot for slot in self.staging_slots if not slot.in_use) + big_slots, small_slots, token_indices = [], [], [] + with torch.cuda.stream(producer_stream): + for index, _ in selected: + end = page_lens[index] + # 最后一页可短于 CPU 页；保持真实长度，不向 transfer kernel 传负索引。 + token_indices.append( + self.backend.model.req_manager.req_to_token_indexs[req.req_idx, index * page_tokens : end] + ) + if self.backend.is_hybrid_att_model: + if end % page_tokens: + small_slots.append(req.tail_small_page_buffer_id) + else: + big_slots.append( + shared_slots[index] + if index < len(shared_slots) + else req.hybrid_len_to_big_page_id[end] + ) + leases = [] + for allocator, slots in ( + (resources.bounded_att_big_page_allocator, big_slots), + (resources.bounded_att_small_page_allocator, small_slots), + ): + if slots: + allocator.retain(slots) + leases.append((allocator, slots)) + staging = slot.buffer[: len(selected)] + self.cpu_cache_operator.write_req( + token_indices=torch.cat(token_indices), + cpu_page_buffers=staging, + cpu_page_indices=list(range(len(selected))), + layouts=self.cpu_cache_operator.get_page_layouts(), + big_page_slot_ids=big_slots, + small_page_slot_ids=small_slots, + ) + packed = torch.cuda.Event() + packed.record() + if leases: + self.checkpoint_tasks.append(CheckpointTask(packed, leases)) + # 源 KV 在计算流上复用；pinned 快照引用只保留到 pack 完成。 + torch.cuda.current_stream().wait_event(packed) + with torch.cuda.stream(g_infer_context.get_cpu_kv_cache_stream()): + torch.cuda.current_stream().wait_event(packed) + self.cpu_cache_operator.scatter_staging_pages( + staging, self.cpu_cache_client.cpu_kv_cache_tensor, [page for _, page in selected] + ) + if self.backend.dp_world_size > 1: + dist.all_reduce(self.offload_sync_tensor, group=self.offload_sync_group) + event = torch.cuda.Event() + event.record() + slot.in_use = True + session.pending_tasks += 1 + self.store_tasks.append(StoreTask([page for _, page in selected], session, slot, event)) + + def close_cache_sessions(self, reqs: List[InferReq]): + # GPU-only、自适应未选中 CPU、diverse 从请求和 abort 也必须结束已有会话。 + for req in reqs: + self.cache_sessions[req.req_id].closing = True + + def _poll_async_tasks(self): + # 只对各 rank 一致的队列/会话状态短路，不能根据本地 event.query() 跳过 collective。 + if not (self.checkpoint_tasks or self.load_tasks or self.store_tasks) and not any( + session.closing for session in self.cache_sessions.values() + ): + return + if self.backend.is_master_in_dp: + completed = [] + for tasks in (self.checkpoint_tasks, self.load_tasks, self.store_tasks): + count = 0 + for task in tasks: + if not task.event.query(): + break + count += 1 + completed.append(count) + self.cpu_cache_client.lock.acquire_sleep1ms() + try: + for task in islice(self.load_tasks, completed[1]): + self.cpu_cache_client.deref_pages(task.pages) + pending = {} + for task in islice(self.store_tasks, completed[2]): + self.cpu_cache_client.update_pages_status_to_ready(task.pages, deref=not task.session.disk_offload) + pending[task.session.req_id] = pending.get(task.session.req_id, 0) + 1 + closed = [] + for req_id, session in self.cache_sessions.items(): + if not session.closing or session.pending_tasks != pending.get(req_id, 0): + continue + # disk 链单独持有引用；不能将其他 writer 尚未完成的页发布到 disk。 + if session.disk_offload: + if not self.cpu_cache_client.check_allpages_ready(session.disk_pages): + continue + self.cpu_cache_client.update_pages_status_to_ready( + session.disk_pages, + deref=True, + disk_offload_enable=True, + token_num_in_page_list=session.token_num, + ) + closed.append(req_id) + finally: + self.cpu_cache_client.lock.release() + data = [completed, closed] + else: + data = [None, None] + # 完成计数和会话关闭合并为一次广播；CPU 页 READY 的事件包含全部 TP 分片。 + dist.broadcast_object_list(data, group=self.filter_group, group_src=0) + completed, closed = data + for _ in range(completed[0]): + task = self.checkpoint_tasks.popleft() + # pack 没有 TP collective；master 的完成不能代替本 rank 的 pinned 源读取完成。 + task.event.synchronize() + for allocator, slots in task.leases: + allocator.free(slots) + for _ in range(completed[1]): + self.load_tasks.popleft() + for _ in range(completed[2]): + task = self.store_tasks.popleft() + task.slot.in_use = False + task.session.pending_tasks -= 1 + for req_id in closed: + del self.cache_sessions[req_id] + def write_req_to_cpu(self, req: InferReq, cpu_kv_cache_stream: torch.cuda.Stream) -> Optional["TransTask"]: """选择请求可缓存的前缀，分配 CPU 页并提交写出，返回待确认完成的任务。""" assert CacheTier.CPU in req.cache_tiers @@ -385,6 +635,9 @@ def _handle_hybrid_att_last_page(self, req: InferReq, move_block_size: int, page return move_block_size def update_cpu_cache_task_states(self): + if self.async_cache: + self._poll_async_tasks() + return if self.backend.is_master_in_dp: trans_ok_tasks = [] while len(self.cpu_cache_handle_queue) != 0: @@ -428,3 +681,40 @@ class TransTask: page_readies: torch.Tensor req_obj: InferReq sync_event: torch.cuda.Event + + +@dataclasses.dataclass +class CacheSession: + req_id: int + page_count: int = 0 + disk_pages: List[int] = dataclasses.field(default_factory=list) + token_num: int = 0 + pending_tasks: int = 0 + closing: bool = False + disk_offload: bool = False + + +@dataclasses.dataclass +class StagingSlot: + buffer: torch.Tensor + in_use: bool = False + + +@dataclasses.dataclass +class CheckpointTask: + event: torch.cuda.Event + leases: list + + +@dataclasses.dataclass +class LoadTask: + pages: List[int] + event: torch.cuda.Event + + +@dataclasses.dataclass +class StoreTask: + pages: List[int] + session: CacheSession + slot: StagingSlot + event: torch.cuda.Event diff --git a/unit_tests/common/att_resources/test_async_cpu_cache.py b/unit_tests/common/att_resources/test_async_cpu_cache.py new file mode 100644 index 0000000000..74dca7fdf3 --- /dev/null +++ b/unit_tests/common/att_resources/test_async_cpu_cache.py @@ -0,0 +1,554 @@ +"""Asynchronous cache copies, writer ownership and request-independent lifetimes.""" + +from collections import deque +from types import SimpleNamespace as NS + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import StatePageAllocator +from lightllm.common.attention_state.full_att import GQAFullAttentionState, GQAAttentionConfig +from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as impl +from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import ( + MultiLevelKvCacheModule, + CacheSession, + StagingSlot, + LoadTask, + StoreTask, + CheckpointTask, +) +from lightllm.server.multi_level_kv_cache import CacheTier +from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient +from unit_tests.common.att_resources.test_state_identity import allocated_state, KINDS +from unit_tests.common.att_resources.test_dsv4_state import make_full +from unit_tests.common.att_resources.test_model_resources import make_resources +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + + +def test_allocator_lifetime(): + a = StatePageAllocator(2) + slots = a.alloc(2) + a.retain(slots) + a.retain([slots[0]]) + a.free(slots) # request exits + assert a.free_count == 0 + a.free(slots) # first copy finishes + assert a.free_count == 1 and a.alloc(1) == [slots[1]] + a.free([slots[0]]) + assert a.free_count == 1 and a.alloc(1) == [slots[0]] + with pytest.raises(AssertionError): + a.retain([2]) + + +@pytest.mark.parametrize("kind", [*KINDS, "c4", "c128"]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_staged_full_pages_match_direct(kind, cpu_args): + cpu_args.cpu_cache_token_page_size = 256 if kind.startswith("c") else 8 + n = cpu_args.cpu_cache_token_page_size + if kind.startswith("c"): + state = make_full(4 if kind == "c4" else 128, device="cuda") + state.allocate_storage(3 * n) + else: + state = allocated_state(kind, device="cuda") + for b in ({"storage": state.storage} if kind.startswith("c") else full_state_test_buffers(state)).values(): + b.view(torch.uint8).random_(0, 256) + owner = AttentionStateManager(request_capacity=2) + owner.register_state("main", state) + owner.finish_state_registration() + op = owner.get_cpu_cache_operator() + size = op.get_page_bytes() + direct = torch.full((3, size), 91, dtype=torch.uint8, pin_memory=True) + actual = torch.full_like(direct, 91, pin_memory=True) + stage = torch.full((2, size), 91, dtype=torch.uint8, device="cuda") + # Non-contiguous CPU page order and a partial final page. + ids = torch.arange(n + n // 2, dtype=torch.int64, device="cuda") + if not kind.startswith("c"): + ids = ids % 8 + op.write_req(ids, direct, [2, 0], op.get_page_layouts()) + op.write_req(ids, stage, [0, 1], op.get_page_layouts()) + op.scatter_staging_pages(stage, actual, [2, 0]) + torch.cuda.synchronize() + # Unused tail bytes are not part of the cached prefix. + for region in op.get_page_layouts()[state.state_name].regions: + for p, length in ((2, n), (0, n // 2)): + a = actual[p, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + b = direct[p, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + slices = [slice(None)] * len(region.axes) + slices[region.axes.index("token")] = slice(0, length) + assert torch.equal(a[tuple(slices)].view(torch.uint8), b[tuple(slices)].view(torch.uint8)) + assert (actual[1] == 91).all() + + +@pytest.mark.parametrize("replicated", [False, True]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_scatter_does_not_overwrite_other_tp_shards(replicated, cpu_args): + cpu_args.cpu_cache_token_page_size = 8 + owners = [] + for rank in range(4): + owner = AttentionStateManager(request_capacity=1) + state = GQAFullAttentionState( + GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + tp_rank=rank, + tp_world_size=4, + replicated=replicated, + ), + page_size=4, + device="cuda", + ) + state.allocate_storage(8) + state.k_storage.fill_(7 + rank) + state.v_storage.fill_(17 + rank) + owner.register_state("main", state) + owner.finish_state_registration() + owners.append(owner) + size = owners[0].get_cpu_cache_operator().get_page_bytes() + actual = torch.full((1, size), 91, dtype=torch.uint8, pin_memory=True) + direct = torch.full_like(actual, 91, pin_memory=True) + for owner in owners: + op = owner.get_cpu_cache_operator() + stage = torch.full((1, size), 203, dtype=torch.uint8, device="cuda") + op.write_req(torch.arange(8, device="cuda"), direct, [0], op.get_page_layouts()) + op.write_req(torch.arange(8, device="cuda"), stage, [0], op.get_page_layouts()) + op.scatter_staging_pages(stage, actual, [0]) + torch.cuda.synchronize() + assert torch.equal(actual, direct) + + +@pytest.mark.parametrize("glm", [False, True]) +@pytest.mark.parametrize("tp", [1, 2, 4]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_staged_bounded_states_and_main_draft(glm, tp, cpu_args): + cpu_args.cpu_cache_token_page_size = 8 + direct = actual = None + for rank in range(tp): + owner, _ = make_resources(glm=glm, tp=tp, rank=rank, mtp=2, draft_layers=2) + for state in owner.get_full_attention_states(): + for b in full_state_test_buffers(state).values(): + b.view(torch.uint8).random_(0, 256) + for state in owner.get_bounded_attention_states(): + if state.get_cpu_transfer_regions(): + pool = state.big_pages_pool + for attr in ("buffer", "conv_state", "ssm_state"): + b = getattr(pool, attr, None) + if b is not None: + b.view(torch.uint8).random_(0, 256) + op = owner.get_cpu_cache_operator() + if direct is None: + direct = torch.full((1, op.get_page_bytes()), 91, dtype=torch.uint8, pin_memory=True) + actual = torch.full_like(direct, 91, pin_memory=True) + stage = torch.full((1, op.get_page_bytes()), 203, dtype=torch.uint8, device="cuda") + kw = dict( + token_indices=torch.arange(8, device="cuda"), + cpu_page_indices=[0], + layouts=op.get_page_layouts(), + big_page_slot_ids=[0], + ) + op.write_req(cpu_page_buffers=direct, **kw) + op.write_req(cpu_page_buffers=stage, **kw) + op.scatter_staging_pages(stage, actual, [0]) + torch.cuda.synchronize() + for layout in op.get_page_layouts().values(): + for region in layout.regions: + assert torch.equal( + actual[:, region.byte_start : region.byte_end], direct[:, region.byte_start : region.byte_end] + ) + + +class FakeEvent: + def __init__(self, ready=False): + self.ready = ready + + def query(self): + return self.ready + + def synchronize(self): + self.ready = True + + +class Lock: + def acquire_sleep1ms(self): + pass + + def release(self): + pass + + +class Client: + def __init__(self): + self.lock = Lock() + self.refs = {} + self.hashes = {} + self.ready = set() + self.disks = [] + self.capacity = 64 + + def allocate_one_page(self, page_items, hash_key, disk_offload_enable): + if hash_key not in self.hashes: + if len(self.hashes) >= self.capacity: + return None, False + self.hashes[hash_key] = len(self.hashes) + p = self.hashes[hash_key] + self.refs[p] = self.refs.get(p, 0) + 1 + return p, p in self.ready + + allocate_pages = CpuKvCacheClient.allocate_pages + allocate_pages_for_write = CpuKvCacheClient.allocate_pages_for_write + + def allocate_pages_with_ownership(self, hashes, disk_offload_enable): + result = self.allocate_pages_for_write(hashes, disk_offload_enable) + return result.page_indices, result.writer_flags + + def deref_pages(self, pages): + for p in pages: + assert self.refs[p] > 0 + self.refs[p] -= 1 + + def update_pages_status_to_ready(self, pages, deref=True, disk_offload_enable=False, token_num_in_page_list=None): + self.ready.update(pages) + if disk_offload_enable: + self.disks.append((list(pages), token_num_in_page_list)) + if deref: + self.deref_pages(pages) + + def check_allpages_ready(self, pages): + return set(pages) <= self.ready + + @property + def page_hash_dict(self): + return self.hashes + + @property + def page_items(self): + return NS(linked_items=None) + + +def module(monkeypatch): + m = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + m.backend = NS(is_master_in_dp=True, dp_world_size=1, is_hybrid_att_model=False) + m.cpu_cache_client = Client() + m.cache_sessions = {} + m.load_tasks = deque() + m.store_tasks = deque() + m.checkpoint_tasks = deque() + m.filter_group = object() + monkeypatch.setattr(impl.dist, "broadcast_object_list", lambda *a, **kw: None) + return m + + +def test_duplicate_loading_page_has_single_writer(): + c = Client() + assert c.allocate_pages_with_ownership([10, 20], False) == ([0, 1], [True, True]) + assert c.allocate_pages_with_ownership([10, 20], False) == ([0, 1], [False, False]) + assert c.allocate_pages([10], False) == ([0], [False]) + c.update_pages_status_to_ready([0], deref=False) + assert c.allocate_pages([10], False) == ([0], [True]) + c.capacity = 2 + assert c.allocate_pages_with_ownership([20, 30, 40], False) == ([1, -1, -1], [False, False, False]) + + +@pytest.mark.parametrize("disk", [False, True]) +def test_abort_detaches_request_but_preserves_pending_copy(monkeypatch, disk): + m = module(monkeypatch) + pages, _ = m.cpu_cache_client.allocate_pages_with_ownership([10, 20], False) + session = CacheSession(1, disk_pages=pages, page_count=len(pages), token_num=16, pending_tasks=1, disk_offload=disk) + m.cache_sessions[1] = session + a = StatePageAllocator(1) + slots = a.alloc(1) + a.retain(slots) + a.free(slots) + event = FakeEvent() + slot = StagingSlot(None, True) + m.store_tasks.append(StoreTask(pages, session, slot, event)) + m.checkpoint_tasks.append(CheckpointTask(event, [(a, slots)])) + m.close_cache_sessions([NS(req_id=1)]) + m._poll_async_tasks() + assert m.cpu_cache_client.refs == {0: 1, 1: 1} and not m.cpu_cache_client.ready + assert a.alloc(1) is None and slot.in_use + event.ready = True + m._poll_async_tasks() + assert m.cpu_cache_client.ready == {0, 1} and m.cpu_cache_client.refs == {0: 0, 1: 0} + assert a.alloc(1) == slots and not slot.in_use and not m.cache_sessions + assert bool(m.cpu_cache_client.disks) == disk + m._poll_async_tasks() # no double free + + +def test_disk_waits_for_other_writer(monkeypatch): + m = module(monkeypatch) + pages, _ = m.cpu_cache_client.allocate_pages_with_ownership([10, 20], False) + m.cache_sessions[1] = CacheSession( + 1, disk_pages=pages, page_count=len(pages), token_num=16, closing=True, disk_offload=True + ) + m.cpu_cache_client.ready.add(0) + m._poll_async_tasks() + assert 1 in m.cache_sessions and not m.cpu_cache_client.disks + m.cpu_cache_client.ready.add(1) + m._poll_async_tasks() + assert m.cpu_cache_client.disks == [([0, 1], 16)] and not m.cache_sessions + + +def test_restore_pages_survive_request_exit(monkeypatch): + m = module(monkeypatch) + pages, _ = m.cpu_cache_client.allocate_pages_with_ownership([10], False) + a = StatePageAllocator(1) + slots = a.alloc(1) + a.retain(slots) + a.free(slots) + e = FakeEvent() + m.load_tasks.append(LoadTask(pages, e)) + m.checkpoint_tasks.append(CheckpointTask(e, [(a, slots)])) + m.cache_sessions[1] = CacheSession(1, closing=True) + m._poll_async_tasks() + assert not m.cache_sessions and m.cpu_cache_client.refs[0] == 1 and a.alloc(1) is None + e.ready = True + m._poll_async_tasks() + assert m.cpu_cache_client.refs[0] == 0 and a.alloc(1) == slots + + +@pytest.mark.parametrize("pages_num", [1, 2, 4, 9]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_real_offload_allows_source_reuse_and_limits_staging(monkeypatch, cpu_args, pages_num): + n = 8 + cpu_args.cpu_cache_token_page_size = n + m = module(monkeypatch) + m.args = NS(cpu_cache_token_page_size=n, enable_disk_cache=True) + m.async_cache = True + m.incremental_offload = False + m.gloo_group = object() + owner = AttentionStateManager(request_capacity=1) + state = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=1, head_dim=16, dtype=torch.bfloat16), + page_size=4, + device="cuda", + ) + state.allocate_storage(n * pages_num) + state.k_storage.fill_(42) + state.v_storage.fill_(67) + owner.register_state("main", state) + owner.finish_state_registration() + m.backend.model = NS( + att_state_manager=owner, + req_manager=NS(req_to_token_indexs=torch.arange(n * pages_num, device="cuda").unsqueeze(0)), + ) + m.cpu_cache_operator = owner.get_cpu_cache_operator() + m.cpu_cache_client.cpu_kv_cache_tensor = torch.empty( + (64, m.cpu_cache_operator.get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) + m.staging_slots = [ + StagingSlot(torch.empty((4, m.cpu_cache_operator.get_page_bytes()), dtype=torch.uint8, device="cuda")) + for _ in range(2) + ] + producer = torch.cuda.Stream() + cpu = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + monkeypatch.setattr(impl.g_infer_context, "get_overlap_stream", lambda: producer) + monkeypatch.setattr(impl.g_infer_context, "get_cpu_kv_cache_stream", lambda: cpu) + monkeypatch.setattr(impl.g_infer_context, "is_hybrid_att_model", False) + req = NS( + req_id=123, + req_idx=0, + cur_kv_len=n * pages_num, + shared_kv_node=None, + cache_tiers=(CacheTier.CPU, CacheTier.DISK), + shm_req=NS( + group_req_id=123, + request_id=123, + input_len=n * pages_num + 1, + token_hash_page_len_list=NS(get_all=lambda: list(range(n, n * (pages_num + 1), n))), + token_hash_list=NS(get_all=lambda: list(range(pages_num))), + ), + ) + m.cache_sessions[123] = CacheSession(123) + with torch.cuda.stream(cpu): + torch.cuda._sleep(20_000_000) + assert m.offload_finished_reqs_to_cpu_cache([req]) == [req] + assert len(m.store_tasks) <= 2 and len(m.staging_slots) == 2 + # Simulate request token pages immediately being reused on the control stream. + state.k_storage.fill_(0) + state.v_storage.fill_(0) + torch.cuda.synchronize() + m._poll_async_tasks() + assert not m.store_tasks and not m.cache_sessions + assert all(count == 0 for count in m.cpu_cache_client.refs.values()) + assert len(m.cpu_cache_client.disks) == 1 + op = m.cpu_cache_operator + op.read_req( + m.cpu_cache_client.cpu_kv_cache_tensor, + list(range(pages_num)), + torch.arange(n * pages_num, device="cuda"), + op.get_page_layouts(), + ) + torch.cuda.synchronize() + assert (state.k_storage[:, : n * pages_num] == 42).all() + assert (state.v_storage[:, : n * pages_num] == 67).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_real_restore_defers_cpu_release_without_host_sync(monkeypatch, cpu_args): + n = 256 + cpu_args.cpu_cache_token_page_size = n + m = module(monkeypatch) + m.args = NS(cpu_cache_token_page_size=n, page_size=4) + m.async_cache = True + m.init_sync_group = object() + owner = AttentionStateManager(request_capacity=1) + state = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=1, head_dim=16, dtype=torch.bfloat16), + page_size=4, + device="cuda", + ) + state.allocate_storage(n) + state.k_storage.fill_(42) + state.v_storage.fill_(67) + owner.register_state("main", state) + owner.finish_state_registration() + op = owner.get_cpu_cache_operator() + cpu = torch.empty((1, op.get_page_bytes()), dtype=torch.uint8, pin_memory=True) + op.write_req(torch.arange(n, device="cuda"), cpu, [0], op.get_page_layouts()) + torch.cuda.synchronize() + state.k_storage.zero_() + state.v_storage.zero_() + m.cpu_cache_client.cpu_kv_cache_tensor = cpu + m.cpu_cache_client.allocate_pages([10], False) + req = NS( + req_id=1, + req_idx=0, + cur_kv_len=0, + hold_kv_len=0, + mm_input_spans=[], + sampling_param=NS(shm_param=NS(prompt_logprobs=-1)), + shm_req=NS( + input_len=n + 1, + disk_prompt_cache_len=0, + cpu_cache_match_page_indexes=NS(get_all=lambda: [0]), + token_hash_page_len_list=NS(get_all=lambda: [n]), + ), + ) + table = torch.zeros((1, n), device="cuda", dtype=torch.int32) + + def alloc(req, size): + ids = torch.arange(size, dtype=torch.int32) + table[req.req_idx, :size] = ids.cuda() + req.hold_kv_len = size + return ids + + m.backend.model = NS(att_state_manager=owner, req_manager=NS(req_to_token_indexs=table)) + m.backend._alloc_req_kv_mem = alloc + m.need_sync_compute_stream = lambda: False + monkeypatch.setattr(impl.g_infer_context, "get_can_alloc_token_num", lambda: n) + monkeypatch.setattr(impl.dist, "barrier", lambda **kw: None) + with monkeypatch.context() as patch: + patch.setattr(torch.cuda.Stream, "synchronize", lambda self: pytest.fail("per-request host sync")) + m.load_cpu_cache_to_reqs([req]) + assert req.cur_kv_len == n and m.cpu_cache_client.refs[0] == 1 and len(m.load_tasks) == 1 + torch.cuda.synchronize() + m._poll_async_tasks() + assert (state.k_storage[:, :n] == 42).all() and (state.v_storage[:, :n] == 67).all() + assert m.cpu_cache_client.refs[0] == 0 and not m.load_tasks + + +@pytest.mark.parametrize("page_tokens,end", [(256, 512), (512, 768)]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_dsv4_incremental_full_and_tail_checkpoint_lifetimes(monkeypatch, cpu_args, page_tokens, end): + from unit_tests.common.att_resources.test_dsv4_state import make_atomic_state, runtime_buffer, gather_swa_reference + + cpu_args.cpu_cache_token_page_size = page_tokens + cpu_args.linear_att_hash_page_size = 256 + cpu_args.linear_att_page_block_num = page_tokens // 256 + owner = AttentionStateManager(3) + full = [make_full(ratio) for ratio in (4, 128)] + bounded = [make_atomic_state(kind, mtp_step=3, device="cuda") for kind in ("swa", "kv_w8", "indexer_w8", "kv_w128")] + for s in (*full, *bounded): + owner.register_state("main", s) + owner.finish_state_registration() + owner.allocate_storage(2048, 256, page_tokens, small_page_capacity=2, device="cuda") + for s in full: + s.storage.random_(0, 256) + for s in bounded: + runtime_buffer(s).random_(0, 127) + big = owner.bounded_att_big_page_allocator.alloc(1)[0] + owner.save_bounded_att_runtime_to_pages([0], [page_tokens], "big", [big]) + m = module(monkeypatch) + m.backend.is_hybrid_att_model = True + m.args = NS( + cpu_cache_token_page_size=page_tokens, + linear_att_hash_page_size=256, + disable_linear_att_small_page_cpu_cache=False, + enable_disk_cache=False, + ) + m.incremental_offload = True + m.async_cache = True + m.gloo_group = object() + m.backend.model = NS( + att_state_manager=owner, req_manager=NS(req_to_token_indexs=torch.arange(end, device="cuda").unsqueeze(0)) + ) + m.cpu_cache_operator = owner.get_cpu_cache_operator() + size = m.cpu_cache_operator.get_page_bytes() + m.cpu_cache_client.cpu_kv_cache_tensor = torch.empty((64, size), dtype=torch.uint8, pin_memory=True) + m.staging_slots = [StagingSlot(torch.empty((4, size), dtype=torch.uint8, device="cuda")) for _ in range(2)] + producer = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + cpu = torch.cuda.Stream() + monkeypatch.setattr(impl.g_infer_context, "get_overlap_stream", lambda: producer) + monkeypatch.setattr(impl.g_infer_context, "get_cpu_kv_cache_stream", lambda: cpu) + monkeypatch.setattr(impl.g_infer_context, "is_hybrid_att_model", True) + req = NS( + req_id=1, + req_idx=0, + cur_kv_len=page_tokens, + shared_kv_node=None, + tail_small_page_buffer_id=None, + hybrid_len_to_big_page_id={page_tokens: big}, + cache_tiers=(CacheTier.GPU,), + shm_req=NS( + group_req_id=1, + request_id=1, + input_len=end + 1, + token_hash_page_len_list=NS(get_all=lambda: [page_tokens, end]), + token_hash_list=NS(get_all=lambda: [10, 20]), + ), + ) + m.cache_sessions[1] = CacheSession(1) + m.store_completed_prefill_pages([req], producer) + assert m.cache_sessions[1].page_count == 1 + if end % page_tokens: + small = owner.bounded_att_small_page_allocator.alloc(1)[0] + req.tail_small_page_buffer_id = small + with torch.cuda.stream(producer): + owner.save_bounded_att_runtime_to_pages([0], [end], "small", [small]) + else: + second = owner.bounded_att_big_page_allocator.alloc(1)[0] + req.hybrid_len_to_big_page_id[end] = second + with torch.cuda.stream(producer): + owner.save_bounded_att_runtime_to_pages([0], [end], "big", [second]) + req.cur_kv_len = end + m.store_completed_prefill_pages([req], producer) + m.close_cache_sessions([req]) + # Request can release the original pinned checkpoints while offload references remain. + owner.bounded_att_big_page_allocator.free(list(req.hybrid_len_to_big_page_id.values())) + if req.tail_small_page_buffer_id is not None: + owner.bounded_att_small_page_allocator.free([req.tail_small_page_buffer_id]) + torch.cuda.synchronize() + expected_swa = gather_swa_reference(bounded[0], 0, end) + expected = [s.storage[:, :end].clone() for s in full] + m._poll_async_tasks() + assert not m.cache_sessions and m.cpu_cache_client.ready == {0, 1} + slots = owner.bounded_att_big_page_allocator.alloc(end // page_tokens) + op = m.cpu_cache_operator + op.read_req( + m.cpu_cache_client.cpu_kv_cache_tensor, + [0, 1], + torch.arange(1024, 1024 + end, device="cuda"), + op.get_page_layouts(), + big_page_slot_ids=slots, + req_idx=1, + state_token_len=end, + ) + torch.cuda.synchronize() + for s, b in zip(full, expected): + assert torch.equal(s.storage[:, 1024 : 1024 + end], b) + assert torch.equal(expected_swa, gather_swa_reference(bounded[0], 1, end)) diff --git a/unit_tests/common/att_resources/test_async_cpu_cache_pressure.py b/unit_tests/common/att_resources/test_async_cpu_cache_pressure.py new file mode 100644 index 0000000000..700edfaba3 --- /dev/null +++ b/unit_tests/common/att_resources/test_async_cpu_cache_pressure.py @@ -0,0 +1,304 @@ +"""Tail-page, checkpoint pressure and CPU-only cache regressions.""" + +from types import SimpleNamespace as NS +import pytest +import torch +from unit_tests.common.att_resources.test_async_cpu_cache import module, FakeEvent, Client +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import StatePageAllocator +from lightllm.common.attention_state.full_att import GQAFullAttentionState, GQAAttentionConfig +from lightllm.server.multi_level_kv_cache import CacheTier +from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as impl +from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import ( + CacheSession, + CheckpointTask, + StoreTask, + StagingSlot, +) +from lightllm.server.router.model_infer.infer_batch import InferenceContext +from unit_tests.common.att_resources.test_dsv4_state import ( + make_full, + make_atomic_state, + runtime_buffer, + gather_swa_reference, +) +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + + +def test_named_allocation_keeps_ready_distinct_from_writer(): + c = Client() + a = c.allocate_pages_for_write([10, 20], False) + assert a.page_indices == [0, 1] and a.ready_flags == [False, False] and a.writer_flags == [True, True] + c.update_pages_status_to_ready([0], deref=False) + a = c.allocate_pages_for_write([10, 20], False) + assert a.ready_flags == [True, False] and a.writer_flags == [False, False] + assert c.allocate_pages([10, 20], False) == ([0, 1], [True, False]) + + +@pytest.mark.parametrize("load", [False, True]) +def test_pressure_reclaims_checkpoint_before_actual_snapshot(monkeypatch, load): + m = module(monkeypatch) + a = StatePageAllocator(4) + ids = a.alloc(4) + a.retain([ids[0]]) + a.free([ids[0]]) # KV/radix owner goes away, asynchronous checkpoint reference remains. + packed = FakeEvent() + m.checkpoint_tasks.append(CheckpointTask(packed, [(a, [ids[0]])])) + a.on_allocation_pressure = m._reclaim_checkpoint_slots + cpu_done = FakeEvent() + if not load: + m.store_tasks.append(StoreTask([], CacheSession(9), StagingSlot(None, True), cpu_done)) + saved = [] + context = InferenceContext() + context.is_hybrid_att_model = True + context.args = NS(linear_att_hash_page_size=256, linear_att_page_block_num=1, disable_chunked_prefill=False) + context.att_state_manager = NS( + bounded_att_big_page_allocator=a, save_bounded_att_runtime_to_pages=lambda **kw: saved.append(kw) + ) + req = NS( + req_idx=0, + cur_kv_len=0, + hybrid_cache_len=512, + hybrid_len_to_big_page_id={}, + get_chuncked_input_token_len=lambda: 256, + ) + context.save_hybrid_state_to_cache(torch.tensor([0]), [req]) + assert req.hybrid_len_to_big_page_id == {256: ids[0]} and packed.ready + assert saved[0]["page_slot_ids"] == [ids[0]] + assert not cpu_done.ready, "allocation must not wait for staging -> CPU completion" + assert not m.checkpoint_tasks + assert a.alloc(1) is None # no invented spare capacity + + +def test_reclaim_waits_only_required_fifo_prefix(monkeypatch): + m = module(monkeypatch) + a = StatePageAllocator(2) + ids = a.alloc(2) + events = [FakeEvent(), FakeEvent()] + for slot, event in zip(ids, events): + a.retain([slot]) + a.free([slot]) + m.checkpoint_tasks.append(CheckpointTask(event, [(a, [slot])])) + a.on_allocation_pressure = m._reclaim_checkpoint_slots + assert a.alloc(1) == [ids[0]] + assert events[0].ready and not events[1].ready and len(m.checkpoint_tasks) == 1 + + +@pytest.mark.parametrize("layout", ["gqa", "packed"]) +@pytest.mark.parametrize("ends,skip", [([256], 0), ([2048, 2304], 0), ([2048, 4096, 4352], 2)]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_real_store_tail_never_passes_negative_slots(monkeypatch, cpu_args, layout, ends, skip): + page_tokens = 2048 + cpu_args.cpu_cache_token_page_size = page_tokens + cpu_args.linear_att_hash_page_size = 256 + cpu_args.linear_att_page_block_num = 8 + owner = AttentionStateManager(3) + if layout == "packed": + full = [make_full(ratio) for ratio in (4, 128)] + else: + full = [ + GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=1, head_dim=16, dtype=torch.bfloat16), + page_size=256, + device="cuda", + ) + ] + swa = make_atomic_state("swa", mtp_step=3, device="cuda") + for state in (*full, swa): + owner.register_state("main", state) + owner.finish_state_registration() + owner.allocate_storage(16384, 256, page_tokens, small_page_capacity=4, device="cuda") + for state in full: + for buffer in ({"packed": state.storage} if layout == "packed" else full_state_test_buffers(state)).values(): + buffer.fill_(42) + runtime_buffer(swa).random_(0, 127) + big_slots = owner.bounded_att_big_page_allocator.alloc(len(ends) - 1) + for end, slot in zip(ends[:-1], big_slots): + owner.save_bounded_att_runtime_to_pages([0], [end], "big", [slot]) + small = owner.bounded_att_small_page_allocator.alloc(1)[0] + owner.save_bounded_att_runtime_to_pages([0], [ends[-1]], "small", [small]) + m = module(monkeypatch) + m.backend.is_hybrid_att_model = True + m.args = NS( + cpu_cache_token_page_size=page_tokens, + linear_att_hash_page_size=256, + disable_linear_att_small_page_cpu_cache=False, + enable_disk_cache=False, + ) + m.async_cache = True + m.incremental_offload = True + m.gloo_group = object() + table = torch.arange(ends[-1], device="cuda", dtype=torch.int32).unsqueeze(0) + m.backend.model = NS(att_state_manager=owner, req_manager=NS(req_to_token_indexs=table)) + m.cpu_cache_operator = owner.get_cpu_cache_operator() + size = m.cpu_cache_operator.get_page_bytes() + m.cpu_cache_client.cpu_kv_cache_tensor = torch.empty((64, size), dtype=torch.uint8, pin_memory=True) + m.staging_slots = [StagingSlot(torch.empty((4, size), device="cuda", dtype=torch.uint8)) for _ in range(2)] + producer = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + cpu = torch.cuda.Stream() + monkeypatch.setattr(impl.g_infer_context, "get_cpu_kv_cache_stream", lambda: cpu) + monkeypatch.setattr(impl.g_infer_context, "is_hybrid_att_model", True) + hashes = list(range(10, 10 + len(ends))) + if skip: + pages, _ = m.cpu_cache_client.allocate_pages(hashes[:skip], False) + m.cpu_cache_client.update_pages_status_to_ready(pages, deref=True) + captured = [] + write = m.cpu_cache_operator.write_req + + def record(**kw): + captured.append(kw["token_indices"].clone()) + return write(**kw) + + monkeypatch.setattr(m.cpu_cache_operator, "write_req", record) + req = NS( + req_id=1, + req_idx=0, + cur_kv_len=ends[-1], + shared_kv_node=None, + tail_small_page_buffer_id=small, + hybrid_len_to_big_page_id=dict(zip(ends[:-1], big_slots)), + cache_tiers=(CacheTier.GPU,), + shm_req=NS( + group_req_id=1, + request_id=1, + input_len=ends[-1] + 1, + token_hash_page_len_list=NS(get_all=lambda: ends), + token_hash_list=NS(get_all=lambda: hashes), + ), + ) + m.cache_sessions[1] = CacheSession(1) + m.store_completed_prefill_pages([req], producer) + torch.cuda.synchronize() + m._poll_async_tasks() + expected_start = skip * page_tokens + assert torch.equal(torch.cat(captured), table[0, expected_start : ends[-1]]) + assert (torch.cat(captured) >= 0).all() + assert all(ref == 0 for ref in m.cpu_cache_client.refs.values()), "CPU-only pages must be evictable during decode" + assert m.cache_sessions[1].page_count == len(ends) and not m.cache_sessions[1].disk_pages + selected_pages = [m.cpu_cache_client.hashes[h] for h in hashes[skip:]] + new_ids = torch.arange(8192, 8192 + ends[-1] - expected_start, device="cuda", dtype=torch.int32) + state_expected = [ + { + name: value[:, expected_start : ends[-1]].clone() + for name, value in ({"packed": s.storage} if layout == "packed" else full_state_test_buffers(s)).items() + } + for s in full + ] + expected_swa = gather_swa_reference(swa, 0, ends[-1]) + restored_big = owner.bounded_att_big_page_allocator.alloc(len(selected_pages) - 1) + m.cpu_cache_operator.read_req( + m.cpu_cache_client.cpu_kv_cache_tensor, + selected_pages, + new_ids, + m.cpu_cache_operator.get_page_layouts(), + big_page_slot_ids=restored_big, + req_idx=1, + state_token_len=ends[-1], + ) + torch.cuda.synchronize() + for state, expected in zip(full, state_expected): + for name, value in ( + {"packed": state.storage} if layout == "packed" else full_state_test_buffers(state) + ).items(): + assert torch.equal(value[:, 8192 : 8192 + new_ids.numel()], expected[name]) + assert torch.equal(gather_swa_reference(swa, 1, ends[-1]), expected_swa) + + +def test_poll_empty_and_batch_allocation_collectives(monkeypatch): + m = module(monkeypatch) + calls = [] + monkeypatch.setattr(impl.dist, "broadcast_object_list", lambda *a, **kw: calls.append(a)) + m._poll_async_tasks() + assert not calls + m.cache_sessions[1] = CacheSession(1, closing=True) + m._poll_async_tasks() + assert len(calls) == 1 and not m.cache_sessions + + +def test_staging_budget_uses_actual_main_and_draft_capabilities(): + owner = AttentionStateManager(1) + op = owner.get_cpu_cache_operator() + # Layout size is independent from the capability selection. + op.get_page_bytes = lambda: 128 + op.configure_backends([NS(requires_cpu_cache_sync=False), None]) + assert op.async_cache and op.get_staging_bytes() == 8 * 128 + op.configure_backends([NS(requires_cpu_cache_sync=False), NS(requires_cpu_cache_sync=True)]) + assert not op.async_cache and op.get_staging_bytes() == 0 + + +def test_fa3_capability_is_inherited_by_all_variants(): + from lightllm.common.basemodel.attention.base_att import BaseAttBackend + from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend + from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend + from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend + from lightllm.common.basemodel.attention.fa3.neo import NeoFa3AttBackend + + assert not BaseAttBackend.requires_cpu_cache_sync + assert all(c.requires_cpu_cache_sync for c in [Fa3AttBackend, Fp8Fa3AttBackend, MlaFa3AttBackend, NeoFa3AttBackend]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_multiple_requests_allocate_once_and_do_not_pin_cpu_prefixes(monkeypatch, cpu_args): + cpu_args.cpu_cache_token_page_size = 8 + owner = AttentionStateManager(2) + state = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=1, head_dim=16, dtype=torch.bfloat16), + page_size=4, + device="cuda", + ) + state.allocate_storage(16) + state.k_storage.fill_(42) + state.v_storage.fill_(67) + owner.register_state("main", state) + owner.finish_state_registration() + m = module(monkeypatch) + m.args = NS(cpu_cache_token_page_size=8, enable_disk_cache=False) + m.async_cache = True + m.incremental_offload = True + m.gloo_group = object() + m.backend.model = NS( + att_state_manager=owner, req_manager=NS(req_to_token_indexs=torch.arange(16, device="cuda").repeat(2, 1)) + ) + m.cpu_cache_operator = owner.get_cpu_cache_operator() + size = m.cpu_cache_operator.get_page_bytes() + m.cpu_cache_client.cpu_kv_cache_tensor = torch.empty((64, size), dtype=torch.uint8, pin_memory=True) + m.staging_slots = [StagingSlot(torch.empty((4, size), device="cuda", dtype=torch.uint8)) for _ in range(2)] + producer = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + cpu = torch.cuda.Stream() + monkeypatch.setattr(impl.g_infer_context, "get_cpu_kv_cache_stream", lambda: cpu) + monkeypatch.setattr(impl.g_infer_context, "is_hybrid_att_model", False) + reqs = [] + for rid, hashes in [(1, [10, 20]), (2, [10, 30])]: + req = NS( + req_id=rid, + req_idx=rid - 1, + cur_kv_len=16, + shared_kv_node=None, + cache_tiers=(CacheTier.GPU,), + shm_req=NS( + group_req_id=rid, + request_id=rid, + input_len=17, + token_hash_page_len_list=NS(get_all=lambda: [8, 16]), + token_hash_list=NS(get_all=lambda hashes=hashes: hashes), + ), + ) + reqs.append(req) + m.cache_sessions[rid] = CacheSession(rid) + calls = [] + monkeypatch.setattr(impl.dist, "broadcast_object_list", lambda *a, **kw: calls.append(a)) + m.store_completed_prefill_pages(reqs, producer) + assert len(calls) == 1 and len(m.store_tasks) == 2 + assert m.cpu_cache_client.refs == {0: 1, 1: 1, 2: 1}, "shared LOADING ref should be released immediately" + torch.cuda.synchronize() + m._poll_async_tasks() + assert len(calls) == 2 and all(v == 0 for v in m.cpu_cache_client.refs.values()) + assert len(m.cache_sessions) == 2, "requests may keep decoding without pinning CPU pages" + m.store_completed_prefill_pages(reqs, producer) + m._poll_async_tasks() + assert len(calls) == 2, "idle replicated queues/cursors must skip collectives" + m.close_cache_sessions(reqs) + m._poll_async_tasks() + assert not m.cache_sessions and len(calls) == 3 diff --git a/unit_tests/common/att_resources/test_async_cpu_cache_tp.py b/unit_tests/common/att_resources/test_async_cpu_cache_tp.py new file mode 100644 index 0000000000..670f04d177 --- /dev/null +++ b/unit_tests/common/att_resources/test_async_cpu_cache_tp.py @@ -0,0 +1,178 @@ +"""Real TP2 completion and checkpoint reuse; requires two CUDA devices.""" + +import os +import json +import socket +from datetime import timedelta +from dataclasses import asdict +from collections import deque +from types import SimpleNamespace as NS +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def worker(rank, port): + torch.cuda.set_device(rank) + os.environ["LIGHTLLM_CURRENT_DEVICE_ID"] = str(rank) + os.environ["LIGHTLLM_CURRENT_RANK_IN_NODE"] = str(rank) + from lightllm.server.core.objs.start_args_type import StartArgs + + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(asdict(StartArgs(cpu_cache_token_page_size=8))) + from unit_tests.common.att_resources.test_async_cpu_cache import Client + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.full_att import GQAFullAttentionState, GQAAttentionConfig + from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import ( + MultiLevelKvCacheModule, + CacheSession, + StagingSlot, + ) + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as impl + from lightllm.server.multi_level_kv_cache import CacheTier + + dist.init_process_group( + "nccl", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=2, timeout=timedelta(seconds=60) + ) + m = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + m.backend = NS(is_master_in_dp=rank == 0, dp_world_size=2, is_hybrid_att_model=False) + m.args = NS(cpu_cache_token_page_size=8, enable_disk_cache=False) + m.async_cache = True + m.incremental_offload = False + m.gloo_group = dist.new_group(backend="gloo") + m.filter_group = dist.new_group(backend="gloo") + m.offload_sync_group = dist.new_group(backend="nccl") + m.offload_sync_tensor = torch.zeros(1, device="cuda") + m.load_tasks = deque() + m.store_tasks = deque() + m.checkpoint_tasks = deque() + m.cache_sessions = {} + owner = AttentionStateManager(1) + state = GQAFullAttentionState( + GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + tp_rank=rank, + tp_world_size=2, + ), + page_size=4, + device=f"cuda:{rank}", + ) + state.allocate_storage(16) + state.k_storage.fill_(42 + rank) + state.v_storage.fill_(67 + rank) + owner.register_state("main", state) + owner.finish_state_registration() + m.backend.model = NS( + att_state_manager=owner, req_manager=NS(req_to_token_indexs=torch.arange(16, device="cuda").unsqueeze(0)) + ) + m.cpu_cache_operator = owner.get_cpu_cache_operator() + size = m.cpu_cache_operator.get_page_bytes() + m.cpu_cache_client = Client() + m.cpu_cache_client.cpu_kv_cache_tensor = torch.full((64, size), 91, dtype=torch.uint8, pin_memory=True) + m.staging_slots = [StagingSlot(torch.empty((4, size), dtype=torch.uint8, device="cuda")) for _ in range(2)] + producer = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + cpu = torch.cuda.Stream() + impl.g_infer_context.get_overlap_stream = lambda: producer + impl.g_infer_context.get_cpu_kv_cache_stream = lambda: cpu + impl.g_infer_context.is_hybrid_att_model = False + req = NS( + req_id=123, + req_idx=0, + cur_kv_len=16, + shared_kv_node=None, + cache_tiers=(CacheTier.CPU,), + shm_req=NS( + group_req_id=123, + request_id=123, + input_len=17, + token_hash_page_len_list=NS(get_all=lambda: [8, 16]), + token_hash_list=NS(get_all=lambda: [10, 20]), + ), + ) + m.cache_sessions[123] = CacheSession(123) + # Warm kernel compilation and first NCCL communicator before timing the delayed shard. + stage = m.staging_slots[0].buffer[:2] + m.cpu_cache_operator.write_req( + torch.arange(16, device="cuda"), stage, [0, 1], m.cpu_cache_operator.get_page_layouts() + ) + m.cpu_cache_operator.scatter_staging_pages(stage, m.cpu_cache_client.cpu_kv_cache_tensor, [0, 1]) + dist.all_reduce(m.offload_sync_tensor, group=m.offload_sync_group) + torch.cuda.synchronize() + m.cpu_cache_client.cpu_kv_cache_tensor.fill_(91) + dist.barrier(group=m.gloo_group) + if rank == 1: + with torch.cuda.stream(cpu): + torch.cuda._sleep(3_000_000_000) + import time + + start = time.perf_counter() + m.offload_finished_reqs_to_cpu_cache([req]) + print(rank, "submit_seconds", time.perf_counter() - start, "done", m.store_tasks[0].event.query(), flush=True) + m._poll_async_tasks() + if rank == 0: + assert not m.cpu_cache_client.ready, "published before other TP shard completed" + torch.cuda.synchronize() + m._poll_async_tasks() + assert not m.store_tasks and not m.cache_sessions + direct = torch.full((2, size), 91, dtype=torch.uint8, pin_memory=True) + m.cpu_cache_operator.write_req( + torch.arange(16, device="cuda"), direct, [0, 1], m.cpu_cache_operator.get_page_layouts() + ) + torch.cuda.synchronize() + assert torch.equal(direct, m.cpu_cache_client.cpu_kv_cache_tensor[:2]) + if rank == 0: + assert m.cpu_cache_client.ready == {0, 1} and m.cpu_cache_client.refs == {0: 0, 1: 0} + print("TP2 NCCL completion/lifetime/scatter checks passed", flush=True) + # A master pack event alone is not proof that the peer's pinned source was consumed. + from lightllm.common.attention_state.allocators import StatePageAllocator + from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import CheckpointTask + + allocator = StatePageAllocator(1) + ids = allocator.alloc(1) + allocator.retain(ids) + allocator.free(ids) + with torch.cuda.stream(producer): + if rank == 1: + torch.cuda._sleep(1_000_000_000) + packed = torch.cuda.Event() + packed.record() + if rank == 0: + packed.synchronize() + m.checkpoint_tasks.append(CheckpointTask(packed, [(allocator, ids)])) + dist.barrier(group=m.gloo_group) + m._poll_async_tasks() + assert packed.query() and not m.checkpoint_tasks + assert allocator.alloc(1) == ids + # Capacity-pressure reclaim must also keep both ranks' slot IDs/order identical. + ids = allocator.alloc(0) # no-op contract + occupied = [0] + allocator.retain(occupied) + allocator.free(occupied) + with torch.cuda.stream(producer): + if rank == 1: + torch.cuda._sleep(100_000_000) + pending = torch.cuda.Event() + pending.record() + m.checkpoint_tasks.append(CheckpointTask(pending, [(allocator, occupied)])) + allocator.on_allocation_pressure = m._reclaim_checkpoint_slots + assert allocator.alloc(1) == occupied and pending.query() and not m.checkpoint_tasks + if rank == 0: + print("TP2 checkpoint release and capacity-pressure checks passed", flush=True) + torch.cuda.synchronize() + dist.barrier(group=m.gloo_group) + dist.destroy_process_group(m.offload_sync_group) + dist.destroy_process_group(m.filter_group) + dist.destroy_process_group(m.gloo_group) + dist.destroy_process_group() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two CUDA devices required") +def test_tp_completion_and_checkpoint_lifetimes(): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + mp.spawn(worker, args=(port,), nprocs=2, join=True) diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index e3a09adb8a..a660273b9b 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -43,7 +43,7 @@ def test_backend_rejects_more_than_36_drafts_before_loading(): @pytest.mark.parametrize("draft_full", [False, True]) @pytest.mark.parametrize("draft_count", [1, 2]) @pytest.mark.parametrize("encoding", ["None", "fp8kv_sph", "fp8kv_spt"]) -@pytest.mark.parametrize("via_backend", [False, True]) +@pytest.mark.parametrize("via_backend,cpu_sync", [(False, None), (True, None), (True, False), (True, True)]) def test_backend_allocates_all_roles_before_warmup( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -52,9 +52,11 @@ def test_backend_allocates_all_roles_before_warmup( draft_count: int, encoding: str, via_backend: bool, + cpu_sync: bool | None, ): args = StartArgs( model_dir=str(tmp_path), + enable_cpu_cache=cpu_sync is not None, data_type="bfloat16", llm_kv_type=encoding, page_size=4, @@ -162,7 +164,9 @@ def mark_backends(model: TpPartBaseModel): ) is state ) - model.prefill_att_backend = model.decode_att_backend = object() + model.prefill_att_backend = model.decode_att_backend = SimpleNamespace( + requires_cpu_cache_sync=bool(cpu_sync and model.is_mtp_draft_model) + ) model.prefill_att_backend1 = model.decode_att_backend1 = None model.backends_bound = True @@ -221,7 +225,14 @@ def memory_info(device: str | torch.device | int): assert manager.draft0_gdn_conv_bf16_ssm_fp32_state.big_pages_pool is None assert manager.full_att_allocator is None assert not warmed - return manager.get_full_att_token_bytes() * (32 + args.page_size), 0 + staging_bytes = 0 + if args.enable_cpu_cache: + assert all(item.backends_bound for item in models) + operator = manager.get_cpu_cache_operator() + assert operator.async_cache is (not cpu_sync) + staging_bytes = operator.get_staging_bytes() + assert (staging_bytes == 0) is cpu_sync + return manager.get_full_att_token_bytes() * (32 + args.page_size) + staging_bytes, 0 monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) monkeypatch.setattr("lightllm.utils.envs_utils.set_model_init_status", ready.append)