diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index f1247e0ef4..bc0fb63a90 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -546,6 +546,7 @@ def _prefill( infer_state.init_some_extra_state(self) infer_state.init_att_state() model_output = self._context_forward(infer_state=infer_state) + infer_state.finish_forward() model_output = self._create_unpad_prefill_model_output( padded_model_output=model_output, @@ -616,11 +617,11 @@ def _decode( else: model_output = self._token_forward(infer_state) + infer_state.finish_forward() return self._create_unpad_decode_model_output(model_output, origin_batch_size=origin_batch_size) @final def _context_forward(self, infer_state: InferStateInfo): - input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" @@ -789,6 +790,8 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input prefill_mem_indexes_ready_event.record() model_output0, model_output1 = self._overlap_tpsp_context_forward(infer_state0, infer_state1=infer_state1) + infer_state0.finish_forward() + infer_state1.finish_forward() model_output0 = self._create_unpad_prefill_model_output( padded_model_output=model_output0, @@ -907,6 +910,8 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 model_output0 = self._create_unpad_decode_model_output(model_output0, origin_batch_size=origin_batch_size0) model_output1 = self._create_unpad_decode_model_output(model_output1, origin_batch_size=origin_batch_size1) + infer_state0.finish_forward() + infer_state1.finish_forward() return model_output0, model_output1 @final diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 29648aa78e..b166e00098 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -140,6 +140,14 @@ def init_att_state(self): if self.decode_att_state1 is not None: self.decode_att_state1.init_state() + def finish_forward(self): + """Update runtime state after one prefill/decode forward, outside graph capture. + + All microbatches have finished their layers before this hook runs. + This precedes prefix checkpoints, not speculative-token acceptance. + """ + return + def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"): for attr_name, attr_value in vars(new_infer_state).items(): if isinstance(attr_value, torch.Tensor): diff --git a/lightllm/common/basemodel/triton_kernel/sliding_window_cpu_cache_copy.py b/lightllm/common/basemodel/triton_kernel/sliding_window_cpu_cache_copy.py new file mode 100644 index 0000000000..cec6227d35 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_cpu_cache_copy.py @@ -0,0 +1,236 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _copy_sliding_window_cpu_cache( + mem_indexes, + page_indexes, + page_readies, + big_page_buffer_ids, + gpu_full_att_kv_state, + cpu_kv_sliding_state, + cpu_cache, + page_num, + full_stride_l, + full_stride_t, + cpu_stride_p, + tp_rank, + FULL_LAYER_NUM: tl.constexpr, + FULL_TOKEN_LAYER_SIZE: tl.constexpr, + FULL_RANK_SIZE: tl.constexpr, + FULL_TOTAL_SIZE: tl.constexpr, + WINDOW_RANK_SIZE: tl.constexpr, + BIG_PAGE_TOKEN_NUM: tl.constexpr, + OFFLOAD: tl.constexpr, + BLOCK: tl.constexpr, +): + # All sizes/offsets below are uint64 elements, not bytes. CPU page offsets + # must remain int64: a shared CPU cache can be much larger than 2 GiB. + full_stride_l = tl.cast(full_stride_l, tl.int64) + full_stride_t = tl.cast(full_stride_t, tl.int64) + cpu_stride_p = tl.cast(cpu_stride_p, tl.int64) + tp_rank = tl.cast(tp_rank, tl.int64) + block_start = tl.program_id(0) + block_count = tl.num_programs(0) + for page in range(page_num): + cpu_page = tl.load(page_indexes + page).to(tl.int64) + copy_page = cpu_page != -1 + if OFFLOAD: + copy_page = copy_page & ~tl.load(page_readies + page).to(tl.int1) + if copy_page: + cpu_page_start = cpu_page * cpu_stride_p + for block in range(block_start, tl.cdiv(FULL_RANK_SIZE, BLOCK), block_count): + offsets = tl.cast(block, tl.int64) * BLOCK + tl.arange(0, BLOCK) + valid = offsets < FULL_RANK_SIZE + token = offsets // (FULL_LAYER_NUM * FULL_TOKEN_LAYER_SIZE) + layer = (offsets // FULL_TOKEN_LAYER_SIZE) % FULL_LAYER_NUM + dim = offsets % FULL_TOKEN_LAYER_SIZE + mem_index = tl.load(mem_indexes + page * BIG_PAGE_TOKEN_NUM + token, valid, other=-1).to(tl.int64) + valid = valid & (mem_index != -1) + gpu_ptr = gpu_full_att_kv_state + layer * full_stride_l + mem_index * full_stride_t + dim + cpu_ptr = cpu_cache + cpu_page_start + tp_rank * FULL_RANK_SIZE + offsets + if OFFLOAD: + value = tl.load(gpu_ptr, valid, other=0) + tl.store(cpu_ptr, value, valid) + else: + value = tl.load(cpu_ptr, valid, other=0) + tl.store(gpu_ptr, value, valid) + + big_page = tl.load(big_page_buffer_ids + page).to(tl.int64) + for block in range(block_start, tl.cdiv(WINDOW_RANK_SIZE, BLOCK), block_count): + offsets = tl.cast(block, tl.int64) * BLOCK + tl.arange(0, BLOCK) + valid = offsets < WINDOW_RANK_SIZE + state_ptr = cpu_kv_sliding_state + big_page * WINDOW_RANK_SIZE + offsets + cpu_ptr = cpu_cache + cpu_page_start + FULL_TOTAL_SIZE + tp_rank * WINDOW_RANK_SIZE + offsets + if OFFLOAD: + value = tl.load(state_ptr, valid, other=0) + tl.store(cpu_ptr, value, valid) + else: + value = tl.load(cpu_ptr, valid, other=0) + tl.store(state_ptr, value, valid) + + +def _copy_state_cache( + mem_indexes, + page_indexes, + page_readies, + big_page_buffer_ids, + gpu_full_att_kv_state, + cpu_kv_sliding_state, + cpu_cache_tensor, + tp_rank, + tp_world_size, + big_page_token_num, + offload, + grid_num, +): + page_num = len(page_indexes) + assert len(big_page_buffer_ids) == page_num + assert len(mem_indexes) == page_num * big_page_token_num + assert not offload or len(page_readies) == page_num + if page_num == 0: + return + + assert gpu_full_att_kv_state.is_contiguous() and cpu_kv_sliding_state.is_contiguous() + assert cpu_cache_tensor.is_contiguous() + + # Packing preserves the original bit patterns, including BF16/FP16 NaNs. + # Storage dimensions, not a model config, define each rank's payload. + # view(uint64) also checks the required element-size alignment. + full_state = gpu_full_att_kv_state.flatten(2).view(torch.uint64) + window_state = cpu_kv_sliding_state.flatten(1).view(torch.uint64) + cpu_cache = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint64) + full_rank_size = big_page_token_num * full_state.shape[0] * full_state.shape[2] + window_rank_size = window_state.shape[1] + page_size = (full_rank_size + window_rank_size) * tp_world_size + assert ( + cpu_cache.shape[1] == triton.cdiv(page_size, 2) * 2 + ), "CPU byte-page layout does not match GPU/checkpoint storage" + + _copy_sliding_window_cpu_cache[(grid_num,)]( + mem_indexes, + page_indexes, + page_readies, + big_page_buffer_ids, + full_state, + window_state, + cpu_cache, + page_num, + full_state.stride(0), + full_state.stride(1), + cpu_cache.stride(0), + tp_rank, + FULL_LAYER_NUM=full_state.shape[0], + FULL_TOKEN_LAYER_SIZE=full_state.shape[2], + FULL_RANK_SIZE=full_rank_size, + FULL_TOTAL_SIZE=full_rank_size * tp_world_size, + WINDOW_RANK_SIZE=window_rank_size, + BIG_PAGE_TOKEN_NUM=big_page_token_num, + OFFLOAD=offload, + BLOCK=4096, + ) + + +def copy_kv_buffer_to_cpu_cache( + mem_indexes: torch.Tensor, + page_indexes: torch.Tensor, + page_readies: torch.Tensor, + big_page_buffer_ids: torch.Tensor, + gpu_full_att_kv_state: torch.Tensor, + cpu_kv_sliding_state: torch.Tensor, + cpu_cache_tensor: torch.Tensor, + tp_rank: int, + tp_world_size: int, + big_page_token_num: int, + grid_num: int = 12, +): + """Pack GPU full KV and pinned CPU checkpoints into this TP rank's CPU page slices.""" + _copy_state_cache( + mem_indexes, + page_indexes, + page_readies, + big_page_buffer_ids, + gpu_full_att_kv_state, + cpu_kv_sliding_state, + cpu_cache_tensor, + tp_rank, + tp_world_size, + big_page_token_num, + offload=True, + grid_num=grid_num, + ) + + +def copy_cpu_cache_to_kv_buffer( + mem_indexes: torch.Tensor, + page_indexes: torch.Tensor, + big_page_buffer_ids: torch.Tensor, + gpu_full_att_kv_state: torch.Tensor, + cpu_kv_sliding_state: torch.Tensor, + cpu_cache_tensor: torch.Tensor, + tp_rank: int, + tp_world_size: int, + big_page_token_num: int, + grid_num: int = 12, +): + """Restore GPU full KV and pinned CPU checkpoints from a packed CPU page.""" + _copy_state_cache( + mem_indexes, + page_indexes, + None, + big_page_buffer_ids, + gpu_full_att_kv_state, + cpu_kv_sliding_state, + cpu_cache_tensor, + tp_rank, + tp_world_size, + big_page_token_num, + offload=False, + grid_num=grid_num, + ) + + +@triton.jit +def _copy_sliding_window_state( + Src, + Dst, + src_layer_stride, + dst_layer_stride, + LAYER_BYTES: tl.constexpr, + TOTAL_BYTES: tl.constexpr, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + for block in range(pid, tl.cdiv(TOTAL_BYTES, BLOCK), tl.num_programs(0)): + offsets = tl.cast(block, tl.int64) * BLOCK + tl.arange(0, BLOCK) + layer = offsets // LAYER_BYTES + within_layer = offsets % LAYER_BYTES + src = Src + layer * tl.cast(src_layer_stride, tl.int64) + within_layer + dst = Dst + layer * tl.cast(dst_layer_stride, tl.int64) + within_layer + values = tl.load(src, mask=offsets < TOTAL_BYTES, other=0) + tl.store(dst, values, mask=offsets < TOTAL_BYTES) + + +@torch.no_grad() +def copy_sliding_window_state(src_state: torch.Tensor, dst_state: torch.Tensor): + """Copy [layer, ...payload] on the CUDA stream, including pinned CPU staging.""" + assert src_state.shape == dst_state.shape + assert src_state.dtype == dst_state.dtype + if not src_state.numel(): + return + assert src_state[0].is_contiguous() and dst_state[0].is_contiguous() + # Only the layer stride may contain gaps in a request's GPU runtime view. + # Reinterpret bytes without flattening/copying that noncontiguous view. + src_bytes, dst_bytes = src_state.view(torch.uint8), dst_state.view(torch.uint8) + total_bytes = src_bytes.numel() + _copy_sliding_window_state[(16,)]( + src_bytes, + dst_bytes, + src_bytes.stride(0), + dst_bytes.stride(0), + LAYER_BYTES=total_bytes // src_state.shape[0], + TOTAL_BYTES=total_bytes, + BLOCK=4096, + ) diff --git a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py new file mode 100644 index 0000000000..4305f96802 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py @@ -0,0 +1,75 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _copy_sliding_window_checkpoint( + gpu_sliding_kv_ptr, # uint8 view: [layer, token_capacity, ...] + req_to_sliding_window, # [req_num, max_seq_len], absolute token position -> GPU mem_index + cache_len, + cpu_kv_sliding_ptr, # uint8 view of one checkpoint: [layer, window, ...] + req_idx, + req_table_stride, + gpu_sliding_layer_stride_bytes, + gpu_sliding_token_stride_bytes, + SLIDING_WINDOW: tl.constexpr, + KV_TOKEN_BYTES: tl.constexpr, + CPU_STATE_BYTES: tl.constexpr, + RESTORE: tl.constexpr, + BLOCK_BYTES: tl.constexpr, +): + for block_index in range(tl.program_id(0), tl.cdiv(CPU_STATE_BYTES, BLOCK_BYTES), tl.num_programs(0)): + state_byte_offsets = tl.cast(block_index, tl.int64) * BLOCK_BYTES + tl.arange(0, BLOCK_BYTES) + layer_index = state_byte_offsets // (SLIDING_WINDOW * KV_TOKEN_BYTES) + window_offset = state_byte_offsets // KV_TOKEN_BYTES % SLIDING_WINDOW + # CPU checkpoints retain absolute token position % W ordering, independent of GPU slot allocation. + token_position = cache_len - 1 - (cache_len - 1 - window_offset + SLIDING_WINDOW) % SLIDING_WINDOW + valid_token = (state_byte_offsets < CPU_STATE_BYTES) & (token_position >= 0) + mem_index = tl.load( + req_to_sliding_window + tl.cast(req_idx, tl.int64) * req_table_stride + token_position, valid_token, other=0 + ).to(tl.int64) + gpu_kv_ptr = ( + gpu_sliding_kv_ptr + + layer_index * gpu_sliding_layer_stride_bytes + + mem_index * gpu_sliding_token_stride_bytes + + state_byte_offsets % KV_TOKEN_BYTES + ) + if RESTORE: + kv_data = tl.load(cpu_kv_sliding_ptr + state_byte_offsets, valid_token, other=0) + tl.store(gpu_kv_ptr, kv_data, valid_token) + else: + kv_data = tl.load(gpu_kv_ptr, valid_token, other=0) + tl.store(cpu_kv_sliding_ptr + state_byte_offsets, kv_data, state_byte_offsets < CPU_STATE_BYTES) + + +def copy_sliding_window_checkpoint( + gpu_sliding_kv_buffer: torch.Tensor, + req_to_sliding_window: torch.Tensor, + cache_len: int, + req_idx: int, + cpu_kv_sliding_state: torch.Tensor, + restore: bool = False, +): + """GPU [layer, token_capacity, ...] 与 CPU pinned [layer, window, ...] checkpoint 的逐字节拷贝。 + + restore=True: CPU checkpoint → GPU KV;否则 GPU KV → CPU checkpoint。 + cache_len 是窗口的绝对 token 右边界;KV_TOKEN_BYTES 是每层每个 token 的 KV 字节数。 + """ + gpu_sliding_kv_bytes = gpu_sliding_kv_buffer.view(torch.uint8) + cpu_kv_sliding_bytes = cpu_kv_sliding_state.view(torch.uint8) + _copy_sliding_window_checkpoint[(16,)]( + gpu_sliding_kv_ptr=gpu_sliding_kv_bytes, + req_to_sliding_window=req_to_sliding_window, + cache_len=cache_len, + cpu_kv_sliding_ptr=cpu_kv_sliding_bytes, + req_idx=req_idx, + req_table_stride=req_to_sliding_window.stride(0), + gpu_sliding_layer_stride_bytes=gpu_sliding_kv_bytes.stride(0), + gpu_sliding_token_stride_bytes=gpu_sliding_kv_bytes.stride(1), + SLIDING_WINDOW=cpu_kv_sliding_state.shape[1], + KV_TOKEN_BYTES=gpu_sliding_kv_bytes.stride(1), + CPU_STATE_BYTES=cpu_kv_sliding_bytes.numel(), + RESTORE=restore, + BLOCK_BYTES=4096, + ) diff --git a/lightllm/common/kv_cache_mem_manager/allocator.py b/lightllm/common/kv_cache_mem_manager/allocator.py index 850c158778..58360f0647 100644 --- a/lightllm/common/kv_cache_mem_manager/allocator.py +++ b/lightllm/common/kv_cache_mem_manager/allocator.py @@ -9,7 +9,7 @@ class KvCacheAllocator: - def __init__(self, size: int) -> None: + def __init__(self, size: int, publish_usage: bool = True) -> None: self.size = size self.mem_state = torch.arange( 0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True @@ -24,12 +24,14 @@ def __init__(self, size: int) -> None: self.can_use_mem_size = self.size - rank_in_node = get_current_rank_in_node() - # 用共享内存进行共享,router 模块读取进行精确的调度估计, nccl port 作为一个单机中单实列的标记。防止冲突。 - self.shared_can_use_token_num = SharedInt( - f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank_in_node}" - ) - self.shared_can_use_token_num.set_value(self.can_use_mem_size) + # Only the full KV pool publishes scheduler capacity; attention-state pools are private. + self.shared_can_use_token_num = None + if publish_usage: + rank_in_node = get_current_rank_in_node() + self.shared_can_use_token_num = SharedInt( + f"{get_unique_server_name()}_mem_manger_can_use_token_num_{rank_in_node}" + ) + self.shared_can_use_token_num.set_value(self.can_use_mem_size) return def alloc(self, need_size) -> torch.Tensor: @@ -42,7 +44,8 @@ def alloc(self, need_size) -> torch.Tensor: self.mark_start += need_size self.can_use_mem_size -= need_size - self.shared_can_use_token_num.set_value(self.can_use_mem_size) + if self.shared_can_use_token_num is not None: + self.shared_can_use_token_num.set_value(self.can_use_mem_size) # 利用缓冲区返回,避免异步情况下的内存竞争 if self._return_start + need_size > self._mem_state_return.shape[0]: @@ -72,7 +75,8 @@ def free(self, free_index: Union[torch.Tensor, List[int]]): self.mark_start -= len(free_index) self.can_use_mem_size += len(free_index) - self.shared_can_use_token_num.set_value(self.can_use_mem_size) + if self.shared_can_use_token_num is not None: + self.shared_can_use_token_num.set_value(self.can_use_mem_size) if self.can_use_mem_size == len(self.mem_state): logger.debug(f"freed all gpu mem size {self.can_use_mem_size}") @@ -83,7 +87,8 @@ def free_all(self): self.mark_start = 0 self.mark_end = len(self.mem_state) self.can_use_mem_size = len(self.mem_state) - self.shared_can_use_token_num.set_value(self.can_use_mem_size) + if self.shared_can_use_token_num is not None: + self.shared_can_use_token_num.set_value(self.can_use_mem_size) return def resize(self, new_size: int) -> None: @@ -103,4 +108,5 @@ def resize(self, new_size: int) -> None: self._return_start = 0 self.can_use_mem_size = self.size - self.shared_can_use_token_num.set_value(self.can_use_mem_size) + if self.shared_can_use_token_num is not None: + self.shared_can_use_token_num.set_value(self.can_use_mem_size) diff --git a/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py new file mode 100644 index 0000000000..1ecb752166 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -0,0 +1,62 @@ +import torch +import triton + +from lightllm.common.state_cache_manager import SlidingWindowStateCacheManager +from lightllm.utils.envs_utils import get_env_start_args + +from .mem_manager import MemoryManager +from .operator.hybrid_sliding import HybridSlidingMemOperator + + +class HybridSlidingMemoryManager(MemoryManager): + """管理 token 粒度的 full KV 和大页 checkpoint,向 attention 提供窗口运行池的引用。""" + + operator_class = HybridSlidingMemOperator + # Bound by the model to req_manager's runtime pool; no allocation or request-slot ownership here. + sliding_kv_buffer: torch.Tensor + + def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): + args = get_env_start_args() + self.sliding_config = sliding_config + self.big_page_token_num = args.linear_att_page_block_num * args.linear_att_hash_page_size + super().__init__( + size=size, + dtype=sliding_config.dtype, + head_num=sliding_config.full_head_num, + head_dim=sliding_config.full_head_dim, + layer_num=sliding_config.full_layer_num, + always_copy=always_copy, + mem_fraction=mem_fraction, + ) + + def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): + super()._init_buffers(size, dtype, head_num, head_dim, layer_num) + # Match linear attention: CPU checkpoints plus two reserved tail-transfer slots. + self.big_page_buffers = SlidingWindowStateCacheManager( + size=triton.cdiv(size, self.big_page_token_num) + 2, + sliding_config=self.sliding_config, + keep_num=2, + ) + self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 + + def write_to_shm(self, req_manager): + # As in Qwen3NextMemManager, keep pickling from replacing pinned CPU + # checkpoints with ordinary shared storage inaccessible to Triton. + big_page_buffers = self.big_page_buffers + self.big_page_buffers = None + try: + return super().write_to_shm(req_manager) + finally: + self.big_page_buffers = big_page_buffers + + def get_att_input_params(self, layer_index: int): + if layer_index in self.sliding_config.sliding_layer_to_cache_index: + layer_buffer = self.sliding_kv_buffer[self.sliding_config.sliding_layer_to_cache_index[layer_index]] + head_num = self.sliding_config.sliding_head_num + return layer_buffer[:, :head_num], layer_buffer[:, head_num:] + return super().get_att_input_params(self.sliding_config.full_layer_to_cache_index[layer_index]) + + def _free_buffers(self): + super()._free_buffers() + self.big_page_buffers = None diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index d217e05c78..5c29b92cb6 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -30,7 +30,9 @@ class MemoryManager: operator_class = NormalMemOperator - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): + def __init__( + self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9, *, publish_usage=True + ): self.size = size self.head_num = head_num self.head_dim = head_dim @@ -40,7 +42,7 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False # profile the max total token num if the size is None self.profile_size(mem_fraction) - self.allocator = KvCacheAllocator(self.size) + self.allocator = KvCacheAllocator(self.size, publish_usage=publish_usage) self._init_buffers( self.size, diff --git a/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py new file mode 100644 index 0000000000..0c9bce6d27 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py @@ -0,0 +1,113 @@ +import torch +import triton + +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size + +from .normal import NormalMemOperator + + +class HybridSlidingMemOperator(NormalMemOperator): + """Full-KV operations and CPU transfers of hybrid sliding checkpoints.""" + + def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): + layer_index = self.mem_manager.sliding_config.full_layer_to_cache_index[layer_index] + return super().copy_kv_to_mem_manager(layer_index, mem_index, kv) + + def load_cpu_cache_to_gpu(self, mem_indexes, page_indexes, cpu_cache_client, req): + from lightllm.common.basemodel.triton_kernel.sliding_window_cpu_cache_copy import ( + copy_cpu_cache_to_kv_buffer, + ) + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + if not len(page_indexes): + return + + mem_manager = self.mem_manager + page_size = mem_manager.big_page_token_num + big_page_num = len(mem_indexes) // page_size + max_kv_len = (req.cur_kv_len // page_size) * page_size + big_page_ids = [] + for _ in range(big_page_num): + page_id = mem_manager.big_page_buffers.alloc_one_state_cache() + assert page_id is not None + req.hybrid_len_to_big_page_id[max_kv_len] = page_id + big_page_ids.append(page_id) + max_kv_len -= page_size + big_page_ids.reverse() + + if len(mem_indexes) % page_size: + padded_token_num = triton.cdiv(len(mem_indexes), page_size) * page_size - len(mem_indexes) + mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), value=-1) + # The CPU tail carries a checkpoint before a big-page boundary. + # Restore through a reserved slot; it must not become a radix big page. + big_page_ids.append(mem_manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) + + big_page_ids_gpu = torch.tensor(big_page_ids, dtype=torch.int64, device="cpu").cuda(non_blocking=True) + copy_cpu_cache_to_kv_buffer( + mem_indexes=mem_indexes, + page_indexes=page_indexes, + big_page_buffer_ids=big_page_ids_gpu, + gpu_full_att_kv_state=mem_manager.kv_buffer, + cpu_kv_sliding_state=mem_manager.big_page_buffers.state_cache, + cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, + tp_rank=get_current_rank_in_dp(), + tp_world_size=get_dp_world_size(), + big_page_token_num=page_size, + ) + # Loads and this restore use the inference stream. The next load may + # reuse its reserved slot only after this copy has been queued. + g_infer_context.req_manager.restore_big_page_state(big_page_buffer_idx=big_page_ids[-1], req=req) + + def offload_gpu_kv_to_cpu_cache(self, mem_indexes, page_indexes, page_readies, cpu_cache_client, req): + from lightllm.common.basemodel.triton_kernel.sliding_window_cpu_cache_copy import ( + copy_kv_buffer_to_cpu_cache, + copy_sliding_window_state, + ) + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + if not len(page_indexes): + return + + mem_manager = self.mem_manager + page_size = mem_manager.big_page_token_num + radix_cache = g_infer_context.radix_cache + big_page_ids = radix_cache.get_big_page_ids_by_node(req.shared_kv_node) + max_kv_len = (len(mem_indexes) // page_size) * page_size + start_kv_len = (len(big_page_ids) + 1) * page_size + for seq_len in range(start_kv_len, max_kv_len + 1, page_size): + big_page_ids.append(req.hybrid_len_to_big_page_id[seq_len]) + + if len(mem_indexes) % page_size: + padded_token_num = triton.cdiv(len(mem_indexes), page_size) * page_size - len(mem_indexes) + mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), value=-1) + assert req.tail_small_page_buffer_id is not None + temp_id = mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID + src_state = radix_cache.small_page_buffers.get_state_cache(req.tail_small_page_buffer_id) + copy_sliding_window_state(src_state, mem_manager.big_page_buffers.get_state_cache(temp_id)) + big_page_ids.append(temp_id) + + big_page_ids_gpu = torch.tensor(big_page_ids, dtype=torch.int64, device="cpu").cuda(non_blocking=True) + # Both staging and the transfer run on the CPU-cache offload stream. + # Serial stream order protects this slot across requests; load uses a + # different reserved slot and cannot overwrite an in-flight offload. + copy_kv_buffer_to_cpu_cache( + mem_indexes=mem_indexes, + page_indexes=page_indexes, + page_readies=page_readies, + big_page_buffer_ids=big_page_ids_gpu, + gpu_full_att_kv_state=mem_manager.kv_buffer, + cpu_kv_sliding_state=mem_manager.big_page_buffers.state_cache, + cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, + tp_rank=get_current_rank_in_dp(), + tp_world_size=get_dp_world_size(), + big_page_token_num=page_size, + ) + + def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + copy_kv_buffer_to_kv_buffer( + src_mem_index.cuda(non_blocking=True), + dst_mem_index.cuda(non_blocking=True), + self.mem_manager.kv_buffer, + ) diff --git a/lightllm/common/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index abcd4f491e..75d7bccb8e 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -1,6 +1,13 @@ from .base import ReqManager -from .linear_att import ReqManagerForMamba from .hybrid_base import HybridAttentionReqManager +from .linear_att import ReqManagerForMamba from .req_sampling_params import ReqSamplingParamsManager +from .sliding_window import ReqManagerForSlidingWindow -__all__ = ["ReqManager", "HybridAttentionReqManager", "ReqManagerForMamba", "ReqSamplingParamsManager"] +__all__ = [ + "ReqManager", + "HybridAttentionReqManager", + "ReqManagerForMamba", + "ReqManagerForSlidingWindow", + "ReqSamplingParamsManager", +] diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index 372ca9deb3..d0223064ca 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -72,10 +72,7 @@ def alloc(self): def free(self, free_req_indexes: List[int], free_token_index): for req_index in free_req_indexes: - self.req_list.free(req_index) - - if self.req_list.is_all_free(): - logger.debug(f"freed all request size {self.req_list.can_alloc_size}") + self.free_req(req_index) self.mem_manager.free(free_token_index) def free_req(self, free_req_index: int): diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py new file mode 100644 index 0000000000..48782e6f0c --- /dev/null +++ b/lightllm/common/req_manager/sliding_window.py @@ -0,0 +1,141 @@ +from typing import TYPE_CHECKING, Optional + +import torch + +from lightllm.common.basemodel.triton_kernel.sliding_window_state import copy_sliding_window_checkpoint +from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager +from lightllm.common.state_cache_manager import SlidingWindowStateCacheManager +from lightllm.utils.dist_utils import get_dp_world_size +from lightllm.utils.envs_utils import get_env_start_args + +from .hybrid_base import HybridAttentionReqManager + + +if TYPE_CHECKING: + from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager + from lightllm.common.state_cache_manager import SlidingWindowCacheConfig + from lightllm.server.router.model_infer.infer_batch import InferReq + + +class ReqManagerForSlidingWindow(HybridAttentionReqManager): + """管理请求的窗口索引与 checkpoint;私有 MemoryManager 负责 GPU KV 和物理槽位。""" + + def __init__( + self, + max_request_num: int, + max_sequence_length: int, + mem_manager: Optional["HybridSlidingMemoryManager"], + sliding_config: "SlidingWindowCacheConfig", + ): + super().__init__(max_request_num, max_sequence_length, mem_manager) + self.sliding_config = sliding_config + self.sliding_window = sliding_config.sliding_window + # Like linear attention, reserve runtime state before mem_manager profiles full KV capacity. + self._init_runtime_buffer() + + def _init_runtime_buffer(self): + args = get_env_start_args() + # Both microbatches share batch_max_tokens; allow TP/dummy padding for each. + max_forward_tokens = max(args.batch_max_tokens, self.max_request_num) + 2 * get_dp_world_size() + self.sliding_mem_manager = MemoryManager( + size=self.max_request_num * self.sliding_window + max_forward_tokens, + dtype=self.sliding_config.dtype, + head_num=self.sliding_config.sliding_head_num, + head_dim=self.sliding_config.sliding_head_dim, + layer_num=self.sliding_config.sliding_layer_num, + publish_usage=False, + ) + # Absolute-token addressing matches full attention, using a separate physical pool. + self.req_to_sliding_window = torch.zeros_like(self.req_to_token_indexs) + self.req_to_sliding_window[self.HOLD_REQUEST_ID].fill_(self.sliding_mem_manager.HOLD_TOKEN_MEMINDEX) + # Each request owns W slots, addressed by absolute position % W on the CPU. + # Prefill may exchange their physical indices; decode reuses them in place. + self._sliding_req_indexes = torch.empty( + (self.max_request_num, self.sliding_window), dtype=torch.int32, device="cpu" + ) + self._sliding_seq_lens = [0] * self.max_request_num + + def alloc(self): + req_idx = super().alloc() + if req_idx is not None: + self._sliding_req_indexes[req_idx].copy_(self.sliding_mem_manager.alloc(self.sliding_window)) + return req_idx + + def init_hybrid_attention_state(self, req: "InferReq"): + # The request already owns its window; a cache miss has no valid history. + self._sliding_seq_lens[req.req_idx] = 0 + + def alloc_sliding_window_indexes(self, req_idx: int, token_num: int): + """先复用请求窗口的空闲槽,再借用本轮额外槽位;返回 CPU 索引片段供 batch 合并。""" + seq_len = self._sliding_seq_lens[req_idx] + # The first query needs at most W-1 history tokens, so at least one slot is reusable. + reuse_num = min(token_num, self.sliding_window - min(seq_len, self.sliding_window - 1)) + ring_start = seq_len % self.sliding_window + indexes = [self._sliding_req_indexes[req_idx, ring_start : ring_start + reuse_num]] + if token_num > reuse_num: + indexes.append(self.sliding_mem_manager.alloc(token_num - reuse_num)) + return indexes + + def update_sliding_window(self, req_idx: int, seq_len: int, new_indexes: torch.Tensor): + """所有层读取后将借用的尾部槽纳入请求窗口,归还被替换和过期的槽;不移动 KV。""" + token_num = new_indexes.numel() + old_seq_len = seq_len - token_num + reuse_num = min(token_num, self.sliding_window - min(old_seq_len, self.sliding_window - 1)) + if token_num > reuse_num: + # Retain only borrowed tokens in the final W positions, replacing their old ring slots. + retain_start = max(reuse_num, token_num - self.sliding_window) + ring_positions = torch.arange(old_seq_len + retain_start, seq_len, device="cpu") % self.sliding_window + ring = self._sliding_req_indexes[req_idx] + expired_indexes = torch.cat((ring[ring_positions], new_indexes[reuse_num:retain_start])) + self.sliding_mem_manager.free(expired_indexes) + ring[ring_positions] = new_indexes[retain_start:] + # Decode only advances the position: no allocator call or window-index copy. + self._sliding_seq_lens[req_idx] = seq_len + + def free_req(self, free_req_index: int): + self.sliding_mem_manager.free(self._sliding_req_indexes[free_req_index]) + self._sliding_seq_lens[free_req_index] = 0 + super().free_req(free_req_index) + + def free_all(self): + self.sliding_mem_manager.free_all() + self._sliding_seq_lens = [0] * self.max_request_num + self.req_to_sliding_window.zero_() + self.req_to_sliding_window[self.HOLD_REQUEST_ID].fill_(self.sliding_mem_manager.HOLD_TOKEN_MEMINDEX) + super().free_all() + + def create_small_page_cache_manager(self, size: int): + self.small_page_buffers = SlidingWindowStateCacheManager(size=size, sliding_config=self.sliding_config) + return self.small_page_buffers + + def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): + cache_len = req.cur_kv_len + if req.shared_kv_node is not None: + # GPU small-page matching restores before updating cur_kv_len; + # a subsequent CPU-cache load can extend beyond this shared node. + cache_len = max(cache_len, req.shared_kv_node.node_prefix_total_len) + window_len = min(cache_len, self.sliding_window) + # Restore into the window reserved by alloc(), ordered by absolute token position. + ring_positions = torch.arange(cache_len - window_len, cache_len, device="cpu") % self.sliding_window + indexes = torch.empty(window_len, dtype=torch.int32, device="cpu", pin_memory=True) + indexes.copy_(self._sliding_req_indexes[req.req_idx, ring_positions]) + self._sliding_seq_lens[req.req_idx] = cache_len + self.req_to_sliding_window[req.req_idx, cache_len - window_len : cache_len].copy_(indexes, non_blocking=True) + copy_sliding_window_checkpoint( + gpu_sliding_kv_buffer=self.sliding_mem_manager.kv_buffer, + req_to_sliding_window=self.req_to_sliding_window, + cache_len=cache_len, + req_idx=req.req_idx, + cpu_kv_sliding_state=state_cache_manager.get_state_cache(buffer_idx), + restore=True, + ) + + def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: SlidingWindowStateCacheManager): + copy_sliding_window_checkpoint( + gpu_sliding_kv_buffer=self.sliding_mem_manager.kv_buffer, + req_to_sliding_window=self.req_to_sliding_window, + cache_len=self._sliding_seq_lens[req_idx], + req_idx=req_idx, + cpu_kv_sliding_state=state_cache_manager.get_state_cache(buffer_idx), + restore=False, + ) diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py index af0479a158..cbaf8fbe3b 100644 --- a/lightllm/common/state_cache_manager/__init__.py +++ b/lightllm/common/state_cache_manager/__init__.py @@ -1,13 +1,17 @@ from .base import StateCacheManager from .layer_cache import LayerCache from .linear_att import LinearAttCacheConfig, LinearAttCacheManager +from .sliding_window import SlidingWindowCacheConfig, SlidingWindowStateCacheManager def get_hybrid_cache_config(): """Return the model-specific layout used by hybrid CPU/disk cache pages.""" - from lightllm.utils.config_utils import is_linear_att_mixed_model + from lightllm.utils.config_utils import is_linear_att_mixed_model, is_sliding_att_mixed_model from lightllm.utils.envs_utils import get_env_start_args - if is_linear_att_mixed_model(get_env_start_args().model_dir): + model_dir = get_env_start_args().model_dir + if is_linear_att_mixed_model(model_dir): return LinearAttCacheConfig.load_from_args() + if is_sliding_att_mixed_model(model_dir): + return SlidingWindowCacheConfig.load_from_args() raise ValueError("No hybrid state-cache layout registered for this model") diff --git a/lightllm/common/state_cache_manager/sliding_window.py b/lightllm/common/state_cache_manager/sliding_window.py new file mode 100644 index 0000000000..e4566ac71a --- /dev/null +++ b/lightllm/common/state_cache_manager/sliding_window.py @@ -0,0 +1,101 @@ +import dataclasses +import math +from typing import Dict + +import torch + +from .base import StateCacheManager + + +@dataclasses.dataclass +class SlidingWindowCacheConfig: + """Physical cache layout for a full + sliding-window transformer.""" + + sliding_layer_to_cache_index: Dict[int, int] + full_layer_to_cache_index: Dict[int, int] + sliding_window: int + sliding_head_num: int + sliding_head_dim: int + full_head_num: int + full_head_dim: int + dtype: torch.dtype + + def __post_init__(self): + assert self.sliding_window > 0 + assert self.sliding_layer_to_cache_index and self.full_layer_to_cache_index + assert not self.sliding_layer_to_cache_index.keys() & self.full_layer_to_cache_index.keys() + self.sliding_layer_num = len(set(self.sliding_layer_to_cache_index.values())) + self.full_layer_num = len(set(self.full_layer_to_cache_index.values())) + assert set(self.sliding_layer_to_cache_index.values()) == set(range(self.sliding_layer_num)) + assert set(self.full_layer_to_cache_index.values()) == set(range(self.full_layer_num)) + + def get_state_shape(self): + return ( + self.sliding_layer_num, + self.sliding_window, + 2 * self.sliding_head_num, + self.sliding_head_dim, + ) + + def get_state_nbytes(self): + return math.prod(self.get_state_shape()) * self.dtype.itemsize + + def get_cpu_cache_full_att_bytes(self, big_page_token_num: int, tp_world_size: int): + return ( + big_page_token_num + * self.full_layer_num + * 2 + * self.full_head_num + * self.full_head_dim + * self.dtype.itemsize + * tp_world_size + ) + + def get_cpu_cache_state_bytes(self, tp_world_size: int): + return self.get_state_nbytes() * tp_world_size + + def get_cpu_cache_big_page_bytes(self, big_page_token_num: int = None, tp_world_size: int = None): + if big_page_token_num is None or tp_world_size is None: + from lightllm.utils.envs_utils import get_env_start_args + + args = get_env_start_args() + if big_page_token_num is None: + big_page_token_num = args.linear_att_hash_page_size * args.linear_att_page_block_num + assert args.cpu_cache_token_page_size == big_page_token_num + if tp_world_size is None: + tp_world_size = args.tp // args.dp + # One CPU page contains all TP shards: full KV, window state, padding. + payload_bytes = self.get_cpu_cache_full_att_bytes(big_page_token_num, tp_world_size) + payload_bytes += self.get_cpu_cache_state_bytes(tp_world_size) + return (payload_bytes + 15) // 16 * 16 + + @classmethod + def load_from_args(cls): + from lightllm.models.gemma4.kv_layout import build_sliding_cache_config + from lightllm.utils.config_utils import get_config_json + from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type + + args = get_env_start_args() + model_config = get_config_json(args.model_dir) + text_config = model_config.get("text_config", model_config) + return build_sliding_cache_config(text_config, args.tp // args.dp, get_llm_data_type()) + + +class SlidingWindowStateCacheManager(StateCacheManager): + """CPU pinned 窗口 checkpoint,布局为 [slot, layer, window, 2 * heads, dim]。""" + + def __init__(self, size: int, sliding_config: SlidingWindowCacheConfig, keep_num: int = 0): + super().__init__(size, keep_num) + self.state_cache = torch.zeros( + (size, *sliding_config.get_state_shape()), + dtype=sliding_config.dtype, + device="cpu", + pin_memory=True, + ) + + def get_state_cache(self, buffer_idx: int): + return self.state_cache[buffer_idx] + + def clear_to_init_state(self): + self.state_cache.zero_() + super().clear_to_init_state() diff --git a/lightllm/models/gemma4/infer_struct.py b/lightllm/models/gemma4/infer_struct.py index 89ad34acbd..3cc21916f7 100644 --- a/lightllm/models/gemma4/infer_struct.py +++ b/lightllm/models/gemma4/infer_struct.py @@ -1,5 +1,10 @@ +import copy +from types import SimpleNamespace + import torch from lightllm.common.basemodel import InferStateInfo +from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import copy_kv_index_to_req +from lightllm.common.infer_utils import init_req_to_token_indexes from lightllm.models.gemma4.triton_kernel.build_b_image_token_end import build_b_image_token_end @@ -21,6 +26,10 @@ def __init__(self): # 则对应的 b_image_token_end 为 [0, 0, 4, 4, 0], # image token 可以看到自己当前这个token以及后面的 image token。 self.b_image_token_end = None + self.has_image_tokens = False + self.sliding_window_mem_index = None + self.sliding_window_mem_index_cpu = None + self.sliding_requests = None def init_some_extra_state(self, model): super().init_some_extra_state(model) @@ -38,10 +47,69 @@ def init_some_extra_state(self, model): position_ids.shape[0], -1 ) if self.is_prefill: - self.max_seq_len = self.max_kv_seq_len self._build_b_image_token_end() + sliding_mem_manager = self.req_manager.sliding_mem_manager + index_chunks = [] + token_num = 0 + for req_idx, _, q_len in self.sliding_requests: + if req_idx != self.req_manager.HOLD_REQUEST_ID: + index_chunks.extend(self.req_manager.alloc_sliding_window_indexes(req_idx, q_len)) + else: + index_chunks.append( + torch.full((q_len,), sliding_mem_manager.HOLD_TOKEN_MEMINDEX, dtype=torch.int32, device="cpu") + ) + token_num += q_len + padding_token_num = self.input_ids.shape[0] - token_num + if padding_token_num > 0: + index_chunks.append( + torch.full( + (padding_token_num,), sliding_mem_manager.HOLD_TOKEN_MEMINDEX, dtype=torch.int32, device="cpu" + ) + ) + # Combine request-window and allocator views into owned pinned storage for asynchronous H2D. + self.sliding_window_mem_index_cpu = torch.empty( + (self.input_ids.shape[0],), dtype=torch.int32, device="cpu", pin_memory=True + ) + if index_chunks: + torch.cat(index_chunks, out=self.sliding_window_mem_index_cpu) + self.sliding_window_mem_index = self.sliding_window_mem_index_cpu.cuda(non_blocking=True) + if self.is_prefill: + init_req_to_token_indexes( + self.req_manager.req_to_sliding_window, + self.b_req_idx, + self.b_seq_len, + self.b_ready_cache_len, + self.b_q_start_loc, + self.sliding_window_mem_index, + self.max_q_seq_len, + ) + else: + copy_kv_index_to_req( + self.req_manager.req_to_sliding_window, + self.b_req_idx, + self.b_seq_len, + self.sliding_window_mem_index, + ) return + def init_att_state(self): + # Share batch tensors, but bind sliding attention to its own token table. + sliding_state = copy.copy(self) + sliding_state.req_manager = SimpleNamespace(req_to_token_indexs=self.req_manager.req_to_sliding_window) + att_state = self.prefill_att_state1 if self.is_prefill else self.decode_att_state1 + att_state.infer_state = sliding_state + super().init_att_state() + + def finish_forward(self): + # Keep the latest W token slots after every KV-sharing reader has finished. + start = 0 + for req_idx, seq_len, q_len in self.sliding_requests: + if req_idx != self.req_manager.HOLD_REQUEST_ID: + self.req_manager.update_sliding_window( + req_idx, seq_len, self.sliding_window_mem_index_cpu[start : start + q_len] + ) + start += q_len + def _build_b_image_token_end(self): device = self.position_ids.device self.b_image_token_end = torch.zeros(self.position_ids.shape[0], dtype=torch.int32, device=device) @@ -66,6 +134,7 @@ def _build_b_image_token_end(self): if image_start_num == 0: return + self.has_image_tokens = True build_b_image_token_end( b_image_start_idx=torch.tensor(b_image_start_idx, dtype=torch.int32).cuda(non_blocking=True), b_image_len=torch.tensor(b_image_len, dtype=torch.int32).cuda(non_blocking=True), diff --git a/lightllm/models/gemma4/kv_layout.py b/lightllm/models/gemma4/kv_layout.py new file mode 100644 index 0000000000..87dd516553 --- /dev/null +++ b/lightllm/models/gemma4/kv_layout.py @@ -0,0 +1,40 @@ +from lightllm.common.state_cache_manager import SlidingWindowCacheConfig + + +def get_kv_cache_layout(config): + """Map Gemma's shared tail layers to their physical KV owners.""" + layer_types = config["layer_types"] + cutoff = len(layer_types) - (config.get("num_kv_shared_layers") or 0) + assert 0 < cutoff <= len(layer_types) + layer_maps = {"sliding_attention": {}, "full_attention": {}} + last_owner = {} + owners = [] + for layer_index, layer_type in enumerate(layer_types): + cache_map = layer_maps[layer_type] + if layer_index < cutoff: + last_owner[layer_type] = layer_index + cache_map[layer_index] = len(cache_map) + else: + cache_map[layer_index] = cache_map[last_owner[layer_type]] + owner = last_owner[layer_type] + owners.append(owner) + return layer_maps, owners + + +def build_sliding_cache_config(config, tp_world_size, dtype): + """Use the same physical owner layout in model and CPU-cache processes.""" + num_sliding_kv = config["num_key_value_heads"] + num_full_kv = config.get("num_global_key_value_heads") or num_sliding_kv + assert num_sliding_kv % tp_world_size == 0, "sliding KV heads must be divisible by TP size" + assert num_full_kv % tp_world_size == 0, "full KV heads must be divisible by TP size" + layer_maps, _ = get_kv_cache_layout(config) + return SlidingWindowCacheConfig( + sliding_layer_to_cache_index=layer_maps["sliding_attention"], + full_layer_to_cache_index=layer_maps["full_attention"], + sliding_window=config["sliding_window"], + sliding_head_num=num_sliding_kv // tp_world_size, + sliding_head_dim=config["head_dim"], + full_head_num=num_full_kv // tp_world_size, + full_head_dim=config["global_head_dim"], + dtype=dtype, + ) diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 2f0c01dbf6..3a681a6a88 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -4,8 +4,10 @@ from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward from lightllm.models.gemma4.layer_weights.transformer_layer_weight import Gemma4TransformerLayerWeight +from lightllm.models.gemma4.kv_layout import get_kv_cache_layout from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import ( context_attention_fwd_gemma4_mm, ) @@ -16,12 +18,8 @@ class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): """ - Gemma-4 decoder block. Per-layer heterogeneity (sliding vs full attention) - is handled by switching shape / RoPE table / sliding-window flag at init - time. The KV cache layout is uniform (sliding shape: num_kv_heads=16, - head_dim=256); full-attention layers pack their (4, 512) tensor into the - first 8 heads of the 16-head slot at cache-write time, then reshape on - read. See Gemma4TpPartModel._init_mem_manager for context. + Gemma-4 decoder block. Full-attention KV stays token granular, while + sliding attention indexes a bounded KV pool through its own token table. """ def __init__(self, layer_num, network_config): @@ -53,25 +51,12 @@ def __init__(self, layer_num, network_config): # TP shard counts for this layer self.tp_q_head_num_ = network_config["num_attention_heads"] // self.tp_world_size_ - self.tp_k_head_num_ = max(total_kv_heads // self.tp_world_size_, 1) + self.tp_k_head_num_ = total_kv_heads // self.tp_world_size_ self.tp_v_head_num_ = self.tp_k_head_num_ self.tp_o_head_num_ = self.tp_q_head_num_ - self.kv_cache_slot_dim_ = network_config["head_dim"] - sliding_total = network_config["num_key_value_heads"] * network_config["head_dim"] - full_total = num_global_kv * network_config["global_head_dim"] - per_token_k_width = max(sliding_total, full_total) - assert ( - per_token_k_width % self.kv_cache_slot_dim_ == 0 - ), f"per-token K width {per_token_k_width} not aligned to kv_cache_slot_dim {self.kv_cache_slot_dim_}" - self.kv_cache_slot_num_ = (per_token_k_width // self.kv_cache_slot_dim_) // self.tp_world_size_ - - # Sliding window (None on full-attn layers) - if self.is_sliding: - sw = network_config.get("sliding_window", 0) - self.sliding_window_ = int(sw) if sw else 0 - else: - self.sliding_window_ = 0 + # Sliding window (unused on full-attention layers). + self.sliding_window_ = network_config["sliding_window"] if self.is_sliding else 0 # E-series Per-Layer Embeddings gate (HF: config.hidden_size_per_layer_input, # absent or 0 on 31B). @@ -81,20 +66,11 @@ def __init__(self, layer_num, network_config): # HF: config.num_kv_shared_layers (may be missing or null on non-E # checkpoints — treat as 0). - kv_shared_count = network_config.get("num_kv_shared_layers") or 0 - total_layers = network_config["num_hidden_layers"] - self.is_kv_shared_ = kv_shared_count > 0 and layer_num >= total_layers - kv_shared_count - self.kv_share_target_layer_ = None - if self.is_kv_shared_: - cutoff = total_layers - kv_shared_count - for j in range(cutoff - 1, -1, -1): - if network_config["layer_types"][j] == layer_type: - self.kv_share_target_layer_ = j - break - assert self.kv_share_target_layer_ is not None, ( - f"layer {layer_num} ({layer_type}) is KV-shared but no earlier non-shared " - f"layer of the same type found below cutoff={cutoff}" - ) + cache_maps, kv_owners = get_kv_cache_layout(network_config) + kv_owner = kv_owners[layer_num] + self.is_kv_shared_ = kv_owner != layer_num + self.kv_cache_layer_index_ = kv_owner + self.sliding_cache_index_ = cache_maps["sliding_attention"].get(layer_num) # Always 1.0: NoPE dims for full-attn layers are zero-padded into # cos/sin (cos=1, sin=0 → identity), so the kernel walks the whole @@ -130,7 +106,7 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight: Gemma4Trans q = infer_state._all_to_all_unbalance_get(data=q) return q, None - # ---- non-shared: full K/V path ---- + # ---- non-shared: project the owner's K/V ---- k = layer_weight.k_proj.mm(input).view(-1, kv_heads, head_dim) if self.k_eq_v: # Full-attn k_eq_v variant (e.g. 31B): K weights serve as V. @@ -155,23 +131,7 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight: Gemma4Trans # kernel's division cancels out, yielding scores = Q @ K^T. q = q * math.sqrt(head_dim) - # Pack into the uniform KV-cache layout (N, 2*slot_num, slot_dim). - # K occupies slots [0, used_slots); V occupies - # [slot_num, slot_num + used_slots). If this layer's K/V width is - # smaller than the allocated cache slot width, pad with zeros. - cache_slot_num = self.kv_cache_slot_num_ - cache_slot_dim = self.kv_cache_slot_dim_ - N = k.shape[0] - k_packed = k.reshape(N, -1, cache_slot_dim) - v_packed = v.reshape(N, -1, cache_slot_dim) - used_cache_slots = k_packed.shape[1] - if used_cache_slots == cache_slot_num: - cache_kv = torch.cat([k_packed, v_packed], dim=1) - else: - cache_kv = self.alloc_tensor((N, 2 * cache_slot_num, cache_slot_dim), dtype=k.dtype) - cache_kv.zero_() - cache_kv[:, :used_cache_slots, :] = k_packed - cache_kv[:, cache_slot_num : cache_slot_num + used_cache_slots, :] = v_packed + cache_kv = torch.cat([k, v], dim=1) if infer_state.need_dp_prefill_balance: q = infer_state._all_to_all_unbalance_get(data=q) @@ -180,39 +140,22 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight: Gemma4Trans return q, cache_kv def _post_cache_kv(self, cache_kv, infer_state, layer_weight): - if self.is_kv_shared_ or cache_kv is None: + if self.is_kv_shared_: + return + if self.is_sliding: + # Prefill packs KV inside the attention callback, using its current batch metadata. + if infer_state.is_prefill: + return + destindex_copy_kv( + cache_kv, + infer_state.sliding_window_mem_index, + infer_state.mem_manager.sliding_kv_buffer[self.sliding_cache_index_], + ) return - return super()._post_cache_kv(cache_kv, infer_state, layer_weight) + super()._post_cache_kv(cache_kv, infer_state, layer_weight) # ----- Attention kernels (sliding window + per-layer KV reshape) --- - def _att_control(self): - if self.is_sliding and self.sliding_window_ > 0: - w = self.sliding_window_ - 1 - return AttControl(use_sliding_window=True, sliding_window=(w, 0)) - return AttControl(use_sliding_window=False, sliding_window=(-1, -1)) - - def _get_layer_kv(self, infer_state: InferStateInfo): - # KV-shared layers read from the target layer's cache slot. - layer_idx = self.kv_share_target_layer_ if self.is_kv_shared_ else self.layer_num_ - _k_raw, _v_raw = infer_state.mem_manager.get_att_input_params(layer_index=layer_idx) - # _k_raw / _v_raw shape (S, cache_slot_num, cache_slot_dim). Use .view - # (not .reshape) so any non-contiguous layout from a future mem_manager - # backend fails loudly instead of silently copying — slice + view is - # O(1) on the standard MemoryManager layout (inner (kv_heads, head_dim) - # span is contiguous). - kv_heads = self.tp_k_head_num_ - head_dim = self.head_dim_ - cache_slot_dim = self.kv_cache_slot_dim_ - used_cache_slots = kv_heads * head_dim // cache_slot_dim - if used_cache_slots == _k_raw.shape[1]: - # Layout already matches this layer's natural shape. - return _k_raw.view(-1, kv_heads, head_dim), _v_raw.view(-1, kv_heads, head_dim) - # Otherwise the K/V live in the first used_cache_slots; the rest is zero pad. - _k = _k_raw[:, :used_cache_slots, :].view(-1, kv_heads, head_dim) - _v = _v_raw[:, :used_cache_slots, :].view(-1, kv_heads, head_dim) - return _k, _v - def _context_attention_kernel( self, q: torch.Tensor, @@ -221,34 +164,40 @@ def _context_attention_kernel( layer_weight: Gemma4TransformerLayerWeight, out=None, ) -> torch.Tensor: - _k, _v = self._get_layer_kv(infer_state) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) + _k, _v = infer_state.mem_manager.get_att_input_params(self.kv_cache_layer_index_) if self.is_sliding: - # Sliding layers always go through the gemma4_mm Triton kernel: it - # handles SWA + image bidirectional masking in one pass. - o_tensor = self.alloc_tensor(_q.shape, q.dtype) - sw = (self.sliding_window_ - 1, 0) if self.sliding_window_ > 0 else (-1, -1) - context_attention_fwd_gemma4_mm( - _q, - _k, - _v, - o_tensor, - infer_state.b_req_idx, - infer_state.b_q_start_loc, - infer_state.b_seq_len, - infer_state.b_ready_cache_len, - infer_state.max_q_seq_len, - infer_state.req_manager.req_to_token_indexs, - infer_state.b_image_token_end, - sliding_window=sw, - ) - return o_tensor.view(q.shape) - - # Full-attn layers: head_dim=512, no SWA, no image bidi — standard - # triton via backend1. - o_tensor = infer_state.prefill_att_state1.prefill_att( - q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor - ) + if not self.is_kv_shared_: + # Use the callback's live indices on prefill graph replay. + destindex_copy_kv( + kv, + infer_state.sliding_window_mem_index, + infer_state.mem_manager.sliding_kv_buffer[self.sliding_cache_index_], + ) + if infer_state.has_image_tokens: + # Image tokens need Gemma's bidirectional mask in addition to SWA. + o_tensor = self.alloc_tensor(_q.shape, q.dtype) + context_attention_fwd_gemma4_mm( + _q, + _k, + _v, + o_tensor, + infer_state.b_req_idx, + infer_state.b_q_start_loc, + infer_state.b_seq_len, + infer_state.b_ready_cache_len, + infer_state.max_q_seq_len, + infer_state.req_manager.req_to_sliding_window, + infer_state.b_image_token_end, + sliding_window=(self.sliding_window_ - 1, 0), + ) + return o_tensor.view(q.shape) + att_state = infer_state.prefill_att_state1 + att_control = AttControl(use_sliding_window=True, sliding_window=(self.sliding_window_ - 1, 0)) + else: + att_state = infer_state.prefill_att_state + att_control = AttControl() + o_tensor = att_state.prefill_att(q=_q, k=_k, v=_v, att_control=att_control, alloc_func=self.alloc_tensor) return o_tensor.view(q.shape) def _token_attention_kernel( @@ -258,10 +207,15 @@ def _token_attention_kernel( layer_weight: Gemma4TransformerLayerWeight, out=None, ) -> torch.Tensor: - _k, _v = self._get_layer_kv(infer_state) + _k, _v = infer_state.mem_manager.get_att_input_params(self.kv_cache_layer_index_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - att_state = infer_state.decode_att_state if self.is_sliding else infer_state.decode_att_state1 - o_tensor = att_state.decode_att(q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor) + if self.is_sliding: + att_state = infer_state.decode_att_state1 + att_control = AttControl(use_sliding_window=True, sliding_window=(self.sliding_window_ - 1, 0)) + else: + att_state = infer_state.decode_att_state + att_control = AttControl() + o_tensor = att_state.decode_att(q=_q, k=_k, v=_v, att_control=att_control, alloc_func=self.alloc_tensor) return o_tensor.view(q.shape) # ----- FFN (Gemma gelu-tanh, fused gate_up + down) ----------------- diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 061c135b4c..447cf7d035 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -3,16 +3,18 @@ import torch from lightllm.models.registry import ModelRegistry from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class +from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager +from lightllm.common.req_manager import ReqManagerForSlidingWindow from lightllm.common.build_utils import repair_config from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo +from lightllm.models.gemma4.kv_layout import build_sliding_cache_config from lightllm.models.gemma4.layer_infer.pre_layer_infer import Gemma4PreLayerInfer from lightllm.models.gemma4.layer_infer.post_layer_infer import Gemma4PostLayerInfer from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer from lightllm.models.gemma4.layer_weights.pre_and_post_layer_weight import Gemma4PreAndPostLayerWeight from lightllm.models.gemma4.layer_weights.transformer_layer_weight import Gemma4TransformerLayerWeight -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger from lightllm.distributed.communication_op import dist_group_manager @@ -65,65 +67,77 @@ def _init_config(self): return def _verify_params(self): + args = get_env_start_args() assert self.load_way == "HF", "Gemma-4 only supports HF format." - assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 - assert self.config["num_key_value_heads"] % self.tp_world_size_ == 0 - # Use `or` rather than the dict.get default: E4B-style configs ship - # `num_global_key_value_heads: null`, which the default form would - # leave as None. - num_global_kv = self.config.get("num_global_key_value_heads") or self.config["num_key_value_heads"] - assert ( - num_global_kv % self.tp_world_size_ == 0 - ), f"num_global_key_value_heads={num_global_kv} must be divisible by tp={self.tp_world_size_}" - kv_shared = self.config.get("num_kv_shared_layers") or 0 - assert 0 <= kv_shared < self.config["num_hidden_layers"], ( - f"num_kv_shared_layers={kv_shared} out of range for " - f"num_hidden_layers={self.config['num_hidden_layers']}" - ) + self.sliding_cache_config = build_sliding_cache_config(self.config, self.tp_world_size_, self.data_type) + if self.config.get("hidden_size_per_layer_input"): + # PLE uses one static buffer, not independent microbatch storage. + assert not ( + args.enable_prefill_microbatch_overlap or args.enable_decode_microbatch_overlap + ), "Gemma-4 PLE does not support microbatch overlap yet" + assert args.mtp_step == 0, "Gemma-4 hybrid sliding-window cache does not support MTP yet" + if args.enable_cpu_cache: + assert not args.disable_dynamic_prompt_cache, "Gemma-4 CPU cache requires GPU prefix cache" + assert not args.disable_chunked_prefill, "Gemma-4 hybrid sliding-window cache requires chunked prefill" + assert args.run_mode == "normal", "Gemma-4 hybrid sliding-window cache does not support PD mode yet" + assert args.llm_kv_type == "None", "Gemma-4 hybrid sliding-window cache does not support quantized KV yet" + assert not args.enable_dp_prompt_cache_fetch, "Gemma-4 sliding-window state does not support DP cache fetch yet" + assert not args.diverse_mode, "Gemma-4 sliding-window state does not support diverse mode yet" return + def _init_req_manager(self): + self.req_manager = ReqManagerForSlidingWindow( + max_request_num=self.max_req_num, + max_sequence_length=max(self.batch_max_tokens, self.max_seq_length), + mem_manager=None, + sliding_config=self.sliding_cache_config, + ) + def _init_mem_manager(self): - # Uniform per-layer KV cache layout. The per-layer cache slot must fit - # whichever layer type has the largest per-token K/V width: sliding - # (num_key_value_heads * head_dim) or full - # (num_global_kv * global_head_dim). Keep cache_slot_dim = head_dim - # and pick cache_slot_num = max-width / head_dim. For 31B this - # collapses to num_key_value_heads; for E4B the full-attn shape wins - # (2*512 > 2*256), so it uses 4 storage slots of 256 dims. - # Gemma4TransformerLayerInfer.__init__ computes the same value and - # uses it to pack/unpack K/V at write/read time. - head_dim = self.config["head_dim"] - num_global_kv = self.config.get("num_global_key_value_heads") or self.config["num_key_value_heads"] - sliding_total = self.config["num_key_value_heads"] * self.config["head_dim"] - full_total = num_global_kv * self.config["global_head_dim"] - per_token_k_width = max(sliding_total, full_total) - head_num_per_rank = (per_token_k_width // head_dim) // self.tp_world_size_ - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=head_num_per_rank, - head_dim=head_dim, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), + self.mem_manager = HybridSlidingMemoryManager( + size=self.max_total_token_num, + sliding_config=self.sliding_cache_config, mem_fraction=self.mem_fraction, ) + self.mem_manager.sliding_kv_buffer = self.req_manager.sliding_mem_manager.kv_buffer return + def _prepare_sliding_requests(self, *model_inputs): + # Capture Gemma-only metadata before H2D; padding's copy.copy preserves it. + # Normal scheduling supplies CPU tensors; only synthetic GPU warmups need D2H. + for model_input in model_inputs: + req_indexes = model_input.b_req_idx.cpu() + seq_lens = model_input.b_seq_len.cpu() + if model_input.is_prefill: + q_lens = seq_lens - model_input.b_ready_cache_len.cpu() + else: + q_lens = torch.ones_like(seq_lens) + model_input.sliding_requests = list(zip(req_indexes.tolist(), seq_lens.tolist(), q_lens.tolist())) + + def forward(self, model_input): + self._prepare_sliding_requests(model_input) + return super().forward(model_input) + + def microbatch_overlap_prefill(self, model_input0, model_input1): + self._prepare_sliding_requests(model_input0, model_input1) + return super().microbatch_overlap_prefill(model_input0, model_input1) + + def microbatch_overlap_decode(self, model_input0, model_input1): + self._prepare_sliding_requests(model_input0, model_input1) + return super().microbatch_overlap_decode(model_input0, model_input1) + + def _create_inferstate(self, model_input, microbatch_index=0): + infer_state = super()._create_inferstate(model_input, microbatch_index) + # CPU rows exclude padding added inside the model; dummy slots are filled separately. + infer_state.sliding_requests = model_input.sliding_requests + return infer_state + def _init_att_backend(self): - # Gemma-4 has per-layer heterogeneous attention: sliding layers use - # (head_dim=256, kv_heads=16); full-attn layers use (head_dim=512, - # kv_heads=4, k_eq_v). FA3 caps head_dim at 256 and flashinfer plans - # once per infer_state on a single shape — both unworkable for the - # heterogeneous layout. Both layer kinds go through triton. - # - # Primary backend = sliding layers. Sliding prefill bypasses the - # backend and calls gemma4_mm directly (SWA + image bidi in one - # pass); the prefill_att_state created here is unused but the - # framework requires prefill_att_backend to be non-None. + # Full-attention head_dim can be 512, beyond FA3's supported limit. self.prefill_att_backend = TritonAttBackend(model=self) self.decode_att_backend = TritonAttBackend(model=self) def _init_att_backend1(self): - # Secondary backend = full-attn layers (head_dim=512, plain causal). self.prefill_att_backend1 = TritonAttBackend(model=self) self.decode_att_backend1 = TritonAttBackend(model=self) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 18a8f6c04f..dd37dace9e 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -392,7 +392,7 @@ def get_can_alloc_token_num(self): def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): """Snapshot request-level attention state at big/small-page boundaries.""" - if not self.is_hybrid_att_model: + if not self.is_hybrid_att_model or self.radix_cache is None: return # Request-state snapshot at a big-page boundary. diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 4bd78d887f..d930bd6374 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -468,9 +468,14 @@ def is_linear_att_mixed_model(model_path: str) -> bool: return False +@lru_cache(maxsize=None) +def is_sliding_att_mixed_model(model_path: str) -> bool: + return get_model_type(model_path) in {"gemma4", "gemma4_text"} + + def is_hybrid_att_model(model_path: str) -> bool: """Models whose non-full attention state follows hybrid checkpoint pages.""" - return is_linear_att_mixed_model(model_path) + return is_linear_att_mixed_model(model_path) or is_sliding_att_mixed_model(model_path) def get_model_type(model_path: str) -> Optional[str]: diff --git a/test/kernel/test_sliding_window_cpu_cache_copy.py b/test/kernel/test_sliding_window_cpu_cache_copy.py new file mode 100644 index 0000000000..81720c66ef --- /dev/null +++ b/test/kernel/test_sliding_window_cpu_cache_copy.py @@ -0,0 +1,202 @@ +import math + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.sliding_window_cpu_cache_copy import ( + copy_cpu_cache_to_kv_buffer, + copy_kv_buffer_to_cpu_cache, + copy_sliding_window_state, +) +from lightllm.common.state_cache_manager import SlidingWindowCacheConfig + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _random_bits(shape, dtype, generator): + data = torch.randint( + 256, + (math.prod(shape) * dtype.itemsize,), + dtype=torch.uint8, + generator=generator, + ) + return data.view(dtype).reshape(shape) + + +def _assert_same_bits(actual, expected): + torch.testing.assert_close(actual.cpu().view(torch.uint8), expected.view(torch.uint8), atol=0, rtol=0) + + +@pytest.mark.parametrize("tp_world_size", [1, 2, 4]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + "layout", + [ + (2, 3, 7, 1, 64, 2, 32), + (4, 1, 5, 2, 32, 1, 128), + # Both full KV and window state span multiple blocks per program. + (9, 2, 257, 4, 256, 1, 128), + # Deliberately leaves alignment padding in the CPU page for TP=1. + (1, 1, 2, 1, 6, 1, 6), + ], +) +def test_multi_page_round_trip_preserves_tp_slices_tail_and_ring(tp_world_size, dtype, layout): + ( + full_layers, + sliding_layers, + window, + full_heads, + full_dim, + sliding_heads, + sliding_dim, + ) = layout + big_page_tokens, page_num, cpu_page_num, slots = 5, 4, 5, 5 + token_num = page_num * big_page_tokens + 7 + # Duplicate logical layers model shared read-only owners. Only distinct + # physical owners must be copied into the CPU page. + sliding_map = {index: index for index in range(sliding_layers)} + sliding_map[sliding_layers + full_layers] = sliding_layers - 1 + full_map = {sliding_layers + index: index for index in range(full_layers)} + config = SlidingWindowCacheConfig( + sliding_map, + full_map, + window, + sliding_heads, + sliding_dim, + full_heads, + full_dim, + dtype, + ) + page_bytes = config.get_cpu_cache_big_page_bytes(big_page_tokens, tp_world_size) + full_rank_bytes = config.get_cpu_cache_full_att_bytes(big_page_tokens, tp_world_size) // tp_world_size + window_rank_bytes = config.get_cpu_cache_state_bytes(tp_world_size) // tp_world_size + full_total_bytes = full_rank_bytes * tp_world_size + token_bytes = full_rank_bytes // big_page_tokens + cpu_cache = torch.full((cpu_page_num, 1, 1, 1, page_bytes), 0xAB, dtype=torch.uint8, pin_memory=True) + expected_cache = cpu_cache.view(cpu_page_num, page_bytes).clone() + generator = torch.Generator().manual_seed(47) + + mem_indexes = torch.randperm(token_num, generator=generator)[: page_num * big_page_tokens].reshape(page_num, -1) + mem_indexes[-1, -2:] = -1 + page_indexes = torch.tensor([3, 0, -1, 2], dtype=torch.int32) + page_readies = torch.tensor([False, True, False, False]) + # A skipped page must not read its checkpoint ID or dereference slot -1. + big_page_ids = torch.tensor([4, -1, -1, 2], dtype=torch.int64) + sources = [] + for rank in range(tp_world_size): + full_cpu = _random_bits((full_layers, token_num, 2 * full_heads, full_dim), dtype, generator) + window_cpu = _random_bits((slots, *config.get_state_shape()), dtype, generator).pin_memory() + sources.append((full_cpu, window_cpu)) + for page in [0, 3]: + cpu_page = page_indexes[page].item() + for offset, token in enumerate(mem_indexes[page].tolist()): + if token != -1: + start = rank * full_rank_bytes + offset * token_bytes + expected_cache[cpu_page, start : start + token_bytes].copy_( + full_cpu[:, token].contiguous().view(torch.uint8).flatten() + ) + start = full_total_bytes + rank * window_rank_bytes + expected_cache[cpu_page, start : start + window_rank_bytes].copy_( + window_cpu[big_page_ids[page]].view(torch.uint8).flatten() + ) + copy_kv_buffer_to_cpu_cache( + mem_indexes=mem_indexes.flatten().cuda(), + page_indexes=page_indexes.cuda(), + page_readies=page_readies.cuda(), + big_page_buffer_ids=big_page_ids.cuda(), + gpu_full_att_kv_state=full_cpu.cuda(), + cpu_kv_sliding_state=window_cpu, + cpu_cache_tensor=cpu_cache, + tp_rank=rank, + tp_world_size=tp_world_size, + big_page_token_num=big_page_tokens, + grid_num=3, + ) + torch.cuda.synchronize() + # This also checks ready/invalid pages, invalid tail tokens, other TP + # ranks, untouched CPU pages and final alignment padding. + _assert_same_bits(cpu_cache.view(cpu_page_num, page_bytes), expected_cache) + + load_indexes = torch.randperm(token_num, generator=generator)[: page_num * big_page_tokens].reshape(page_num, -1) + load_indexes[-1, -2:] = -1 + load_pages = torch.tensor([3, -1, -1, 2], dtype=torch.int32, device="cuda") + load_slots = torch.tensor([1, -1, -1, 4], dtype=torch.int64) + for rank, (full_cpu, window_cpu) in enumerate(sources): + expected_full = torch.full_like(full_cpu.view(torch.uint8), 0xCD).view(dtype) + expected_window = torch.full_like(window_cpu.view(torch.uint8), 0xCD).view(dtype) + full_gpu, window_pinned = expected_full.cuda(), expected_window.pin_memory() + for page in [0, 3]: + for offset, target in enumerate(load_indexes[page].tolist()): + if target != -1: + expected_full[:, target].copy_(full_cpu[:, mem_indexes[page, offset]]) + # The ring's physical order must be unchanged, including when the + # token page length differs from the sliding window length. + expected_window[load_slots[page]].copy_(window_cpu[big_page_ids[page]]) + copy_cpu_cache_to_kv_buffer( + mem_indexes=load_indexes.flatten().cuda(), + page_indexes=load_pages, + big_page_buffer_ids=load_slots.cuda(), + gpu_full_att_kv_state=full_gpu, + cpu_kv_sliding_state=window_pinned, + cpu_cache_tensor=cpu_cache, + tp_rank=rank, + tp_world_size=tp_world_size, + big_page_token_num=big_page_tokens, + grid_num=3, + ) + torch.cuda.synchronize() + _assert_same_bits(full_gpu, expected_full) + _assert_same_bits(window_pinned, expected_window) + _assert_same_bits(cpu_cache.view(cpu_page_num, page_bytes), expected_cache) + + +def test_empty_copy_is_a_noop(): + indexes = torch.empty(0, dtype=torch.int64, device="cuda") + kwargs = dict( + mem_indexes=indexes, + page_indexes=indexes, + big_page_buffer_ids=indexes, + gpu_full_att_kv_state=None, + cpu_kv_sliding_state=None, + cpu_cache_tensor=None, + tp_rank=0, + tp_world_size=1, + big_page_token_num=16, + ) + copy_kv_buffer_to_cpu_cache(page_readies=indexes, **kwargs) + copy_cpu_cache_to_kv_buffer(**kwargs) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("layers,window,heads,dim", [(3, 17, 2, 6), (5, 129, 4, 64)]) +def test_state_copy_preserves_layer_strides_and_cpu_staging_stream_order(dtype, layers, window, heads, dim): + generator = torch.Generator().manual_seed(71) + source = _random_bits((layers, 5 * window + 7, heads, dim), dtype, generator) + source_gpu = source.cuda() + source_requests = source_gpu[:, : 5 * window].view(layers, 5, window, heads, dim) + expected = torch.full((layers, 6 * window + 11, heads, dim), -3, dtype=dtype) + restored = expected.cuda() + restored_requests = restored[:, : 6 * window].view(layers, 6, window, heads, dim) + # Checkpoints are size-first; runtime requests are layer-first views of a + # larger pool, so copying one request must preserve both layer strides. + checkpoint_pool = torch.zeros((2, layers, window, heads, dim), dtype=dtype, pin_memory=True) + checkpoint = checkpoint_pool[1] + staging = torch.empty_like(checkpoint, pin_memory=True) + staging.zero_() + assert checkpoint.is_contiguous() and not source_requests[:, 1].is_contiguous() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + # Keep preceding GPU work pending: a host-side CPU copy_ would bypass it. + torch.cuda._sleep(1_000_000) + for src_req, dst_req in [(1, 0), (3, 4)]: + copy_sliding_window_state(source_requests[:, src_req], checkpoint) + copy_sliding_window_state(checkpoint, staging) + copy_sliding_window_state(staging, restored_requests[:, dst_req]) + stream.synchronize() + for src_start, dst_start in [(window, 0), (3 * window, 4 * window)]: + expected[:, dst_start : dst_start + window] = source[:, src_start : src_start + window] + _assert_same_bits(restored, expected) + _assert_same_bits(checkpoint, source[:, 3 * window : 4 * window]) + _assert_same_bits(staging, source[:, 3 * window : 4 * window]) + assert torch.count_nonzero(checkpoint_pool[0]).item() == 0 diff --git a/test/kernel/test_sliding_window_decode.py b/test/kernel/test_sliding_window_decode.py new file mode 100644 index 0000000000..ba2eda2e4d --- /dev/null +++ b/test/kernel/test_sliding_window_decode.py @@ -0,0 +1,64 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import ( + gqa_token_decode_attention_flash_decoding, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("window", [512, 1024]) +@pytest.mark.parametrize("q_heads,kv_heads,head_dim", [(4, 1, 64), (8, 2, 256), (32, 2, 128)]) +def test_main_gqa_reads_scattered_window_slots(dtype, window, q_heads, kv_heads, head_dim): + torch.manual_seed(42) + req_ids = [6, 0, 4, 2, 7, 1] + lengths = [1, 2, window - 1, window, window + 1, 3 * window + 7] + full = torch.randn((sum(lengths), 2 * kv_heads, head_dim), device="cuda", dtype=dtype) + pool = torch.full((len(req_ids) * window + 1, 2 * kv_heads, head_dim), float("nan"), device="cuda", dtype=dtype) + full_table = torch.zeros((8, max(lengths)), dtype=torch.int32, device="cuda") + # Unmapped positions point to NaN KV; slot 0 holds the first real request's KV. + window_table = torch.full_like(full_table, pool.shape[0] - 1) + shuffled_slots = torch.arange(pool.shape[0] - 1, device="cuda") + shuffled_slots[1:] = torch.randperm(pool.shape[0] - 2, device="cuda") + 1 + full_offset, window_offset = 0, 0 + for req, length in zip(req_ids, lengths): + indexes = torch.arange(full_offset, full_offset + length, device="cuda", dtype=torch.int32) + full_table[req, :length] = indexes + retained = min(length, window) + slots = shuffled_slots[window_offset : window_offset + retained] + pool[slots] = full[indexes[-retained:].long()] + window_table[req, length - retained : length] = slots.int() + full_offset += length + window_offset += retained + q = torch.randn((len(req_ids), q_heads, head_dim), device="cuda", dtype=dtype) + state = SimpleNamespace( + batch_size=len(req_ids), + b_req_idx=torch.tensor(req_ids, dtype=torch.int32, device="cuda"), + b_seq_len=torch.tensor(lengths, dtype=torch.int32, device="cuda"), + max_kv_seq_len=max(lengths), + req_manager=SimpleNamespace(req_to_token_indexs=full_table), + ) + expected = gqa_token_decode_attention_flash_decoding( + q, + state, + full[:, :kv_heads], + full[:, kv_heads:], + max_len_in_batch=state.max_kv_seq_len, + out=torch.empty_like(q), + sliding_window=(window - 1, 0), + ) + state.req_manager.req_to_token_indexs = window_table + actual = gqa_token_decode_attention_flash_decoding( + q, + state, + pool[:, :kv_heads], + pool[:, kv_heads:], + max_len_in_batch=state.max_kv_seq_len, + out=torch.empty_like(q), + sliding_window=(window - 1, 0), + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) diff --git a/test/kernel/test_sliding_window_prefill.py b/test/kernel/test_sliding_window_prefill.py new file mode 100644 index 0000000000..0389d3abf1 --- /dev/null +++ b/test/kernel/test_sliding_window_prefill.py @@ -0,0 +1,103 @@ +import pytest +import torch + +from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import context_attention_fwd_gemma4_mm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _compare_runtime_and_paged(window, q_len, dtype, head_dim=64, image_span=None): + torch.manual_seed(42) + req_ids, q_lens = [5, 1, 3], [q_len, 7, 173] + histories = [2 * window + 73, 0, window - 19] + lengths = [history + count for history, count in zip(histories, q_lens)] + # Request IDs and physical query offsets are deliberately unrelated. The + # flattened query buffer also has gaps, which must not be read or written. + starts = [11, 47 + q_lens[0], 83 + q_lens[0] + q_lens[1]] + query_num = max(start + count for start, count in zip(starts, q_lens)) + 13 + kv_heads, q_heads = 2, 4 + req_slots = max(req_ids) + 2 + reference = torch.randn((sum(lengths) + 29, 2 * kv_heads, head_dim), device="cuda", dtype=dtype) + mapping = torch.full((req_slots, max(lengths) + 11), -1, device="cuda", dtype=torch.int32) + shuffled_indexes = torch.randperm(reference.shape[0], device="cuda") + offset = 0 + for req_id, length in zip(req_ids, lengths): + mapping[req_id, :length] = shuffled_indexes[offset : offset + length].to(torch.int32) + offset += length + + runtime = torch.full( + (req_slots * window + query_num, 2 * kv_heads, head_dim), float("nan"), device="cuda", dtype=dtype + ) + # Unmapped positions point to NaN KV, not a zero-filled fallback slot. + runtime_mapping = torch.full_like(mapping, runtime.shape[0] - 1) + runtime_slots = torch.randperm(runtime.shape[0] - 1, device="cuda") + offset = 0 + for req_id, history, length in zip(req_ids, histories, lengths): + positions = torch.arange(max(0, history - window), length, device="cuda") + slots = runtime_slots[offset : offset + positions.numel()] + offset += positions.numel() + runtime_mapping[req_id, positions] = slots.int() + runtime[slots] = reference[mapping[req_id, positions].long()] + + image_ends = torch.zeros(query_num, device="cuda", dtype=torch.int32) + if image_span is not None: + image_start, image_end = image_span + image_ends[starts[0] + max(0, image_start) : starts[0] + image_end] = histories[0] + image_end + q = torch.randn((query_num, q_heads, head_dim), device="cuda", dtype=dtype) + expected, actual = torch.full_like(q, -11), torch.full_like(q, -11) + int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + kwargs = dict( + q=q, + b_req_idx=int_tensor(req_ids), + b_start_loc=int_tensor(starts), + b_seq_len=int_tensor(lengths), + b_prompt_cache_len=int_tensor(histories), + max_input_len=max(q_lens), + b_image_token_end=image_ends, + sliding_window=(window - 1, 0), + ) + context_attention_fwd_gemma4_mm( + k=reference[:, :kv_heads], + v=reference[:, kv_heads:], + o=expected, + req_to_token_indexs=mapping, + **kwargs, + ) + context_attention_fwd_gemma4_mm( + k=runtime[:, :kv_heads], + v=runtime[:, kv_heads:], + o=actual, + req_to_token_indexs=runtime_mapping, + **kwargs, + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + if image_span is not None: + # Verify that this case actually exercises bidirectional attention, + # rather than comparing two paths where the image mask is a no-op. + causal = torch.full_like(q, -11) + kwargs["b_image_token_end"] = torch.zeros_like(image_ends) + context_attention_fwd_gemma4_mm( + k=reference[:, :kv_heads], + v=reference[:, kv_heads:], + o=causal, + req_to_token_indexs=mapping, + **kwargs, + ) + assert not torch.equal(expected, causal) + + +@pytest.mark.parametrize("window", [512, 1024]) +@pytest.mark.parametrize("q_len", [1, 31, 4096, 8192]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_runtime_prefill_matches_paged_history_and_current_kv(window, q_len, dtype): + _compare_runtime_and_paged(window, q_len, dtype) + + +@pytest.mark.parametrize("window", [512, 1024]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("image_span", [(17, 43), (31, 197), (-23, 177)]) +def test_runtime_prefill_preserves_image_bidirectional_mask(window, dtype, image_span): + # Production sliding head dimension uses 64-token query tiles. The cases + # cover an image inside one tile, multiple tiles, and the cached boundary. + _compare_runtime_and_paged(window, 384, dtype, head_dim=256, image_span=image_span) diff --git a/test/kernel/test_sliding_window_state.py b/test/kernel/test_sliding_window_state.py new file mode 100644 index 0000000000..ff793750bf --- /dev/null +++ b/test/kernel/test_sliding_window_state.py @@ -0,0 +1,33 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.sliding_window_state import copy_sliding_window_checkpoint + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.mark.parametrize("dtype,payload", [(torch.bfloat16, (4, 32)), (torch.float16, (2, 64)), (torch.uint8, (584,))]) +@pytest.mark.parametrize("seq_len", [0, 1, 17, 32, 97]) +def test_checkpoint_preserves_ring_byte_layout_with_scattered_runtime(dtype, payload, seq_len): + window, capacity = 32, 200 + table = torch.zeros((2, 128), device="cuda", dtype=torch.int32) + pool = torch.randint(0, 100, (3, capacity, *payload), device="cuda", dtype=dtype) + checkpoint = torch.empty((3, window, *payload), dtype=dtype, pin_memory=True) + window_len = min(window, seq_len) + # Slot 0 is ordinary KV; force the checkpoint to include it. + free = torch.arange(capacity - 1, device="cuda", dtype=torch.int32) + free[1:] = torch.randperm(capacity - 2, device="cuda", dtype=torch.int32) + 1 + slots, restored = free[:window_len], free[window_len : 2 * window_len] + table[1, seq_len - window_len : seq_len] = slots + table[0, seq_len - window_len : seq_len] = restored + + copy_sliding_window_checkpoint(pool, table, seq_len, 1, checkpoint) + positions = torch.arange(seq_len - window_len, seq_len, device="cuda") + expected = torch.zeros_like(checkpoint, device="cuda") + expected[:, positions % window] = pool[:, slots.long()] + torch.cuda.synchronize() + torch.testing.assert_close(checkpoint.cuda(), expected, atol=0, rtol=0) + + # Restore the same checkpoint to different physical slots. + copy_sliding_window_checkpoint(pool, table, seq_len, 0, checkpoint, restore=True) + torch.testing.assert_close(pool[:, restored.long()], pool[:, slots.long()], atol=0, rtol=0) diff --git a/unit_tests/models/gemma4/test_hybrid_cache_config.py b/unit_tests/models/gemma4/test_hybrid_cache_config.py new file mode 100644 index 0000000000..c9b3cb8e80 --- /dev/null +++ b/unit_tests/models/gemma4/test_hybrid_cache_config.py @@ -0,0 +1,48 @@ +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.state_cache_manager import SlidingWindowCacheConfig, get_hybrid_cache_config +from lightllm.utils import config_utils, envs_utils + + +@pytest.mark.parametrize("wrapped", [False, True]) +@pytest.mark.parametrize("tp_world_size", [1, 2]) +def test_gemma_hybrid_cache_factory_preserves_shared_owners_and_cpu_page_bytes( + monkeypatch, tmp_path, wrapped, tp_world_size +): + text_config = { + "model_type": "gemma4_text", + "layer_types": ["sliding_attention", "full_attention", "sliding_attention", "full_attention"], + "num_kv_shared_layers": 2, + "num_key_value_heads": 4, + "num_global_key_value_heads": 2, + "sliding_window": 8, + "head_dim": 4, + "global_head_dim": 8, + } + config = {"model_type": "gemma4", "text_config": text_config} if wrapped else text_config + (tmp_path / "config.json").write_text(json.dumps(config)) + args = SimpleNamespace( + model_dir=str(tmp_path), + tp=tp_world_size * 2, + dp=2, + linear_att_hash_page_size=8, + linear_att_page_block_num=2, + cpu_cache_token_page_size=16, + ) + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr(envs_utils, "get_llm_data_type", lambda: torch.bfloat16) + + assert config_utils.is_hybrid_att_model(args.model_dir) + layout = get_hybrid_cache_config() + assert isinstance(layout, SlidingWindowCacheConfig) + assert layout.sliding_layer_to_cache_index == {0: 0, 2: 0} + assert layout.full_layer_to_cache_index == {1: 0, 3: 0} + assert layout.sliding_head_num == 4 // tp_world_size + assert layout.full_head_num == 2 // tp_world_size + # Across all TP ranks: 1024 bytes of full KV plus 512 bytes of window checkpoint. + assert layout.get_cpu_cache_big_page_bytes() == 1536 + assert layout.get_cpu_cache_big_page_bytes(16, tp_world_size) == 1536 diff --git a/unit_tests/models/gemma4/test_sliding_runtime.py b/unit_tests/models/gemma4/test_sliding_runtime.py new file mode 100644 index 0000000000..2eea6daf1b --- /dev/null +++ b/unit_tests/models/gemma4/test_sliding_runtime.py @@ -0,0 +1,167 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.common.req_manager import req_sampling_params, sliding_window +from lightllm.common.state_cache_manager import SlidingWindowCacheConfig +from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo +from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer +from lightllm.models.gemma4.model import Gemma4TpPartModel + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.mark.parametrize("window", [1, 8, 31]) +def test_request_windows_preserve_kv_across_wrap_restore_and_release(monkeypatch, window): + monkeypatch.setattr(req_sampling_params, "ReqSamplingParamsManager", lambda _: None) + monkeypatch.setattr(sliding_window, "get_env_start_args", lambda: SimpleNamespace(batch_max_tokens=128)) + monkeypatch.setattr(sliding_window, "get_dp_world_size", lambda: 1) + layout = SlidingWindowCacheConfig({0: 0}, {1: 0}, window, 1, 4, 1, 4, torch.float32) + manager = sliding_window.ReqManagerForSlidingWindow(2, 512, None, layout) + pool = manager.sliding_mem_manager + req_ids = [manager.alloc(), manager.alloc()] + assert manager.alloc() is None + available = pool.size - 2 * window + assert pool.allocator.can_use_mem_size == available + lengths = {req_idx: 0 for req_idx in req_ids} + alloc = Mock(wraps=pool.alloc) + free = Mock(wraps=pool.free) + monkeypatch.setattr(pool, "alloc", alloc) + monkeypatch.setattr(pool, "free", free) + + def assert_window(req_idx): + length = lengths[req_idx] + positions = torch.arange(max(0, length - window), length, device="cuda") + slots = manager.req_to_sliding_window[req_idx, positions].long() + actual = pool.kv_buffer[0, slots, 0, 0] + torch.testing.assert_close(actual, (positions + req_idx * 1000).float()) + + for req_idx, count in [(req_ids[1], 3), (req_ids[0], 2), (req_ids[0], window * 3 + 5)] + [ + (req_ids[i % 2], 1) for i in range(window * 3 + 5) + ]: + old_len = lengths[req_idx] + alloc.reset_mock() + free.reset_mock() + indexes = torch.cat(manager.alloc_sliding_window_indexes(req_idx, count)) + assert indexes.numel() == count + assert indexes.unique().numel() == count + positions = torch.arange(old_len, old_len + count, device="cuda") + manager.req_to_sliding_window[req_idx, positions] = indexes.cuda() + pool.kv_buffer[0, indexes.long().cuda()] = (positions + req_idx * 1000).float()[:, None, None] + # Writing current KV must retain the W-1 history tokens needed by its first query. + history = torch.arange(max(0, old_len - window + 1), old_len, device="cuda") + old_slots = manager.req_to_sliding_window[req_idx, history].long() + torch.testing.assert_close(pool.kv_buffer[0, old_slots, 0, 0], (history + req_idx * 1000).float()) + lengths[req_idx] += count + manager.update_sliding_window(req_idx, lengths[req_idx], indexes) + if count == 1: + alloc.assert_not_called() + free.assert_not_called() + assert pool.allocator.can_use_mem_size == available + for other_req in req_ids: + assert_window(other_req) + + req_idx = req_ids[0] + cache = manager.create_small_page_cache_manager(1) + manager.save_state(req_idx, 0, cache) + torch.cuda.synchronize() + manager.free_req(req_idx) + assert manager.alloc() == req_idx + # Poison the reserved slots so that restore must actually copy checkpoint data. + reserved = manager._sliding_req_indexes[req_idx].long().cuda() + pool.kv_buffer[:, reserved] = float("nan") + alloc.reset_mock() + free.reset_mock() + req = SimpleNamespace( + req_idx=req_idx, cur_kv_len=0, shared_kv_node=SimpleNamespace(node_prefix_total_len=lengths[req_idx]) + ) + manager.restore_state(req, cache, 0) + torch.cuda.synchronize() + alloc.assert_not_called() + free.assert_not_called() + for other_req in req_ids: + assert_window(other_req) + manager.free_req(other_req) + assert pool.allocator.can_use_mem_size == pool.size + manager.alloc() + manager.free_all() + assert pool.allocator.can_use_mem_size == pool.size + assert manager.req_list.is_all_free() + assert torch.all(manager.req_to_sliding_window[manager.HOLD_REQUEST_ID] == pool.HOLD_TOKEN_MEMINDEX) + + +@pytest.mark.parametrize("is_prefill", [True, False]) +@pytest.mark.parametrize("is_sliding", [True, False]) +def test_gemma_attention_uses_correct_pool_and_supports_full_head_dim_512(is_prefill, is_sliding): + torch.manual_seed(42) + seq_len, window = 29, 8 + q_len = 5 if is_prefill else 1 + head_dim = 64 if is_sliding else 512 + q = torch.randn((q_len, 4, head_dim), device="cuda", dtype=torch.bfloat16) + kv = torch.randn((seq_len, 4, head_dim), device="cuda", dtype=q.dtype) + pool = torch.full((seq_len + 1, 4, head_dim), float("nan"), device="cuda", dtype=q.dtype) + slots = torch.randperm(seq_len, device="cuda") + pool[slots] = kv + valid_table = slots.int()[None, :] + invalid_table = torch.full_like(valid_table, seq_len) + state = Gemma4InferStateInfo() + state.is_prefill = is_prefill + state.b_req_idx = torch.tensor([0], device="cuda", dtype=torch.int32) + state.b_seq_len = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + state.b_q_start_loc = torch.tensor([0], device="cuda", dtype=torch.int32) + state.b_ready_cache_len = torch.tensor([seq_len - q_len], device="cuda", dtype=torch.int32) + state.max_q_seq_len = q_len + state.max_kv_seq_len = seq_len + state.batch_size = 1 + state.total_token_num = seq_len + state.req_manager = SimpleNamespace( + req_to_token_indexs=invalid_table if is_sliding else valid_table, + req_to_sliding_window=valid_table if is_sliding else invalid_table, + ) + state.sliding_window_mem_index = slots[-q_len:].int() + state.mem_manager = SimpleNamespace( + sliding_kv_buffer=pool[None, :], get_att_input_params=lambda _: (pool[:, :2], pool[:, 2:]) + ) + model = object.__new__(Gemma4TpPartModel) + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda _: 0) + model.is_mtp_draft_model = False + model._init_att_backend() + model._init_att_backend1() + if is_prefill: + state.prefill_att_state = model.prefill_att_backend.create_att_prefill_state(state) + state.prefill_att_state1 = model.prefill_att_backend1.create_att_prefill_state(state) + else: + state.decode_att_state = model.decode_att_backend.create_att_decode_state(state) + state.decode_att_state1 = model.decode_att_backend1.create_att_decode_state(state) + state.init_att_state() + layer = object.__new__(Gemma4TransformerLayerInfer) + layer.tp_q_head_num_ = 4 + layer.head_dim_ = head_dim + layer.kv_cache_layer_index_ = 0 + layer.sliding_cache_index_ = 0 + layer.is_sliding = is_sliding + layer.is_kv_shared_ = False + layer.sliding_window_ = window + layer.alloc_tensor = lambda shape, dtype, device="cuda": torch.empty(shape, dtype=dtype, device=device) + if is_prefill: + actual = layer._context_attention_kernel(q, kv[-q_len:], state, None) + else: + actual = layer._token_attention_kernel(q, state, None) + query_positions = torch.arange(seq_len - q_len, seq_len, device="cuda")[:, None] + key_positions = torch.arange(seq_len, device="cuda")[None, :] + mask = key_positions <= query_positions + if is_sliding: + mask &= key_positions > query_positions - window + expected = ( + torch.nn.functional.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0).float(), + kv[:, :2].repeat_interleave(2, dim=1).transpose(0, 1).unsqueeze(0).float(), + kv[:, 2:].repeat_interleave(2, dim=1).transpose(0, 1).unsqueeze(0).float(), + attn_mask=mask, + ) + .squeeze(0) + .transpose(0, 1) + ) + torch.testing.assert_close(actual.float(), expected, atol=2e-2, rtol=2e-2)