From 57e69ccc5c38a0a561199441406684fcd8e0d4fa Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:04:51 +0000 Subject: [PATCH 01/14] feat: support request-level sliding-window cache --- .../common/basemodel/attention/triton/fp.py | 10 +- .../gqa/flash_decoding/gqa_flash_decoding.py | 6 +- .../triton_kernel/sliding_window_state.py | 148 +++++++++++++++ .../hybrid_sliding_mem_manager.py | 61 +++++++ .../kv_cache_mem_manager/operator/__init__.py | 1 + .../operator/hybrid_sliding.py | 16 ++ lightllm/common/req_manager/__init__.py | 11 +- lightllm/common/req_manager/hybrid_att.py | 46 +++++ lightllm/common/req_manager/linear_att.py | 49 ++++- lightllm/common/req_manager/sliding_window.py | 172 ++++++++++++++++++ .../sliding_window_cache_manager/__init__.py | 4 + .../sliding_window_cache_manager/config.py | 73 ++++++++ lightllm/models/gemma4/infer_struct.py | 4 + .../layer_infer/transformer_layer_infer.py | 87 ++++----- lightllm/models/gemma4/model.py | 66 ++++--- lightllm/server/core/objs/req.py | 6 +- .../server/router/model_infer/infer_batch.py | 92 ++++------ .../model_infer/mode_backend/base_backend.py | 18 +- .../mode_backend/chunked_prefill/impl.py | 4 +- .../mode_backend/dp_backend/impl.py | 12 +- lightllm/utils/config_utils.py | 26 +++ 21 files changed, 757 insertions(+), 155 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/sliding_window_state.py create mode 100644 lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py create mode 100644 lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py create mode 100644 lightllm/common/req_manager/hybrid_att.py create mode 100644 lightllm/common/req_manager/sliding_window.py create mode 100644 lightllm/common/sliding_window_cache_manager/__init__.py create mode 100644 lightllm/common/sliding_window_cache_manager/config.py diff --git a/lightllm/common/basemodel/attention/triton/fp.py b/lightllm/common/basemodel/attention/triton/fp.py index e7ce66c774..d5f5cce39b 100644 --- a/lightllm/common/basemodel/attention/triton/fp.py +++ b/lightllm/common/basemodel/attention/triton/fp.py @@ -114,6 +114,7 @@ def decode_att( v: torch.Tensor, att_control: AttControl = AttControl(), alloc_func=torch.empty, + req_to_token_indexs=None, ): if att_control.use_alibi: assert att_control.use_sliding_window is False, "alibi + sliding_window not supported" @@ -133,7 +134,12 @@ def decode_att( return self._normal_decode_flash_decoding_att(q=q, k=k, v=v, alloc_func=alloc_func) elif q_head_num > k_head_num: return self._normal_decode_gqa_flash_decoding_att( - q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func + q=q, + k=k, + v=v, + att_control=att_control, + alloc_func=alloc_func, + req_to_token_indexs=req_to_token_indexs, ) else: raise NotImplementedError("error") @@ -195,6 +201,7 @@ def _normal_decode_gqa_flash_decoding_att( v: torch.Tensor, att_control: AttControl = AttControl(), alloc_func=torch.empty, + req_to_token_indexs=None, ): from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import ( gqa_token_decode_attention_flash_decoding, @@ -215,6 +222,7 @@ def _normal_decode_gqa_flash_decoding_att( out=out, alloc_tensor_func=alloc_func, sliding_window=sliding_window, + req_to_token_indexs=req_to_token_indexs, ) return out diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py index 59a7d4f742..d4a806d6f4 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py @@ -9,6 +9,7 @@ def gqa_token_decode_attention_flash_decoding( out=None, alloc_tensor_func=torch.empty, sliding_window=(-1, -1), + req_to_token_indexs=None, ): batch_size = infer_state.batch_size q_head_num, head_dim = q.shape[1], q.shape[2] @@ -34,11 +35,14 @@ def gqa_token_decode_attention_flash_decoding( mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device="cuda") + if req_to_token_indexs is None: + req_to_token_indexs = infer_state.req_manager.req_to_token_indexs + flash_decode_stage1( q=q.view(calcu_shape1), k=cache_k, v=cache_v, - Req_to_tokens=infer_state.req_manager.req_to_token_indexs, + Req_to_tokens=req_to_token_indexs, B_req_idx=infer_state.b_req_idx, B_Seqlen=infer_state.b_seq_len, max_len_in_batch=infer_state.max_kv_seq_len, 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..37a7693c96 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py @@ -0,0 +1,148 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _prepare_history_indexes( + ReqToSliding, + BReqIdx, + BSeqLen, + BQSeqLen, + stride_req, + stride_seq, + WINDOW: tl.constexpr, + BLOCK: tl.constexpr, +): + batch_idx = tl.program_id(0) + block_idx = tl.program_id(1) + req_idx = tl.load(BReqIdx + batch_idx) + seq_len = tl.load(BSeqLen + batch_idx) + q_len = tl.load(BQSeqLen + batch_idx) + history_end = seq_len - q_len + history_start = tl.maximum(0, history_end - WINDOW) + pos = history_start + block_idx * BLOCK + tl.arange(0, BLOCK) + mask = pos < history_end + physical = req_idx * WINDOW + pos % WINDOW + tl.store(ReqToSliding + req_idx * stride_req + pos * stride_seq, physical, mask=mask) + + +@triton.jit +def _prepare_current_indexes( + ReqToSliding, + BReqIdx, + BSeqLen, + BQSeqLen, + BQStartLoc, + stride_req, + stride_seq, + SCRATCH_START: tl.constexpr, + BLOCK: tl.constexpr, +): + batch_idx = tl.program_id(0) + block_idx = tl.program_id(1) + req_idx = tl.load(BReqIdx + batch_idx) + seq_len = tl.load(BSeqLen + batch_idx) + q_len = tl.load(BQSeqLen + batch_idx) + q_start = tl.load(BQStartLoc + batch_idx) + offset = block_idx * BLOCK + tl.arange(0, BLOCK) + mask = offset < q_len + pos = seq_len - q_len + offset + physical = SCRATCH_START + q_start + offset + tl.store(ReqToSliding + req_idx * stride_req + pos * stride_seq, physical, mask=mask) + + +@torch.no_grad() +def prepare_sliding_window_indexes( + req_to_sliding_window_indexs, + b_req_idx, + b_seq_len, + b_q_seq_len, + b_q_start_loc, + sliding_window, + scratch_start, + max_q_seq_len, +): + block = 256 + _prepare_history_indexes[(b_req_idx.shape[0], triton.cdiv(sliding_window, block))]( + req_to_sliding_window_indexs, + b_req_idx, + b_seq_len, + b_q_seq_len, + *req_to_sliding_window_indexs.stride(), + WINDOW=sliding_window, + BLOCK=block, + ) + _prepare_current_indexes[(b_req_idx.shape[0], triton.cdiv(max_q_seq_len, block))]( + req_to_sliding_window_indexs, + b_req_idx, + b_seq_len, + b_q_seq_len, + b_q_start_loc, + *req_to_sliding_window_indexs.stride(), + SCRATCH_START=scratch_start, + BLOCK=block, + ) + + +@triton.jit +def _commit_sliding_window_state( + LayerBuffer, + BReqIdx, + BSeqLen, + BQSeqLen, + BQStartLoc, + stride_token, + stride_head, + stride_dim, + HEAD_NUM: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW: tl.constexpr, + SCRATCH_START: tl.constexpr, + BLOCK_D: tl.constexpr, +): + batch_idx = tl.program_id(0) + q_offset = tl.program_id(1) + head_idx = tl.program_id(2) + req_idx = tl.load(BReqIdx + batch_idx) + seq_len = tl.load(BSeqLen + batch_idx) + q_len = tl.load(BQSeqLen + batch_idx) + q_start = tl.load(BQStartLoc + batch_idx) + pos = seq_len - q_len + q_offset + mask_token = (q_offset < q_len) & (pos >= seq_len - WINDOW) + src_token = SCRATCH_START + q_start + q_offset + dst_token = req_idx * WINDOW + pos % WINDOW + dims = tl.arange(0, BLOCK_D) + mask = mask_token & (head_idx < HEAD_NUM) & (dims < HEAD_DIM) + src = src_token * stride_token + head_idx * stride_head + dims * stride_dim + dst = dst_token * stride_token + head_idx * stride_head + dims * stride_dim + value = tl.load(LayerBuffer + src, mask=mask, other=0.0) + tl.store(LayerBuffer + dst, value, mask=mask) + + +@torch.no_grad() +def commit_sliding_window_state( + layer_buffer, + b_req_idx, + b_seq_len, + b_q_seq_len, + b_q_start_loc, + sliding_window, + scratch_start, + max_q_seq_len, +): + block_d = triton.next_power_of_2(layer_buffer.shape[-1]) + grid = (b_req_idx.shape[0], max_q_seq_len, layer_buffer.shape[1]) + _commit_sliding_window_state[grid]( + layer_buffer, + b_req_idx, + b_seq_len, + b_q_seq_len, + b_q_start_loc, + *layer_buffer.stride(), + HEAD_NUM=layer_buffer.shape[1], + HEAD_DIM=layer_buffer.shape[2], + WINDOW=sliding_window, + SCRATCH_START=scratch_start, + BLOCK_D=block_d, + ) 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..3caa7dcaab --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -0,0 +1,61 @@ +import triton + +from lightllm.common.req_manager.sliding_window 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-granular full KV plus request-granular sliding-window KV.""" + + operator_class = HybridSlidingMemOperator + + def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): + self.sliding_config = sliding_config + 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) + big_page_token_num = ( + get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size + ) + self.hybrid_att_big_page_buffers = SlidingWindowStateCacheManager( + size=max(1, triton.cdiv(self.size, big_page_token_num)), + sliding_config=self.sliding_config, + ) + # Compatibility with the unchanged big/small-page radix implementation. + self.linear_att_big_page_buffers = self.hybrid_att_big_page_buffers + + def get_att_input_params(self, layer_index: int): + return super().get_att_input_params(self.sliding_config.get_full_layer_index(layer_index)) + + def get_full_cache_layer_index(self, layer_index: int): + return self.sliding_config.get_full_layer_index(layer_index) + + def _free_buffers(self): + super()._free_buffers() + self.hybrid_att_big_page_buffers = None + self.linear_att_big_page_buffers = None + + def write_to_shm(self, req_manager): + # Host-side checkpoints are local to the inference process. Excluding + # them also preserves their pinned allocation during serialization. + big_page_buffers = self.hybrid_att_big_page_buffers + legacy_big_page_buffers = self.linear_att_big_page_buffers + self.hybrid_att_big_page_buffers = None + self.linear_att_big_page_buffers = None + try: + return super().write_to_shm(req_manager) + finally: + self.hybrid_att_big_page_buffers = big_page_buffers + self.linear_att_big_page_buffers = legacy_big_page_buffers diff --git a/lightllm/common/kv_cache_mem_manager/operator/__init__.py b/lightllm/common/kv_cache_mem_manager/operator/__init__.py index 85c37ad39b..26740d567d 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/__init__.py +++ b/lightllm/common/kv_cache_mem_manager/operator/__init__.py @@ -2,6 +2,7 @@ from .normal import NormalMemOperator from .quant import QuantScaleMemOperator, PPLInt4KVMemOperator, PPLInt8KVMemOperator from .linear_att import LinearAttMemOperator +from .hybrid_sliding import HybridSlidingMemOperator from .deepseek import ( Deepseek2MemOperator, Deepseek3_2MemOperator, 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..4dafe300b4 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py @@ -0,0 +1,16 @@ +import torch + +from .normal import NormalMemOperator + + +class HybridSlidingMemOperator(NormalMemOperator): + """GPU KV operations for the token-granular full-attention cache.""" + + 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 078aa8c155..11b92fb15a 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -1,5 +1,14 @@ from .base import ReqManager +from .hybrid_att import HybridAttentionReqManager from .linear_att import ReqManagerForMamba from .req_sampling_params import ReqSamplingParamsManager +from .sliding_window import ReqManagerForSlidingWindow, SlidingWindowStateCacheManager -__all__ = ["ReqManager", "ReqManagerForMamba", "ReqSamplingParamsManager"] +__all__ = [ + "ReqManager", + "HybridAttentionReqManager", + "ReqManagerForMamba", + "ReqManagerForSlidingWindow", + "ReqSamplingParamsManager", + "SlidingWindowStateCacheManager", +] diff --git a/lightllm/common/req_manager/hybrid_att.py b/lightllm/common/req_manager/hybrid_att.py new file mode 100644 index 0000000000..d24b427ea6 --- /dev/null +++ b/lightllm/common/req_manager/hybrid_att.py @@ -0,0 +1,46 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, List, Union + +import torch + +from .base import ReqManager + + +if TYPE_CHECKING: + from lightllm.server.router.model_infer.infer_batch import InferReq + + +class HybridAttentionReqManager(ReqManager, ABC): + """Request manager contract for token/full + request-state attention models. + + The token index table remains the virtual, token-granular address space used + by prefix-cache matching. The non-full attention state is managed through + this interface and may have a different physical granularity. + """ + + is_linear_attention = False + + @abstractmethod + def create_state_cache_manager(self, size: int): + """Create checkpoint storage used by request-state page boundaries.""" + + @abstractmethod + def init_hybrid_attention_state(self, req: "InferReq"): + """Initialize request runtime state when no prefix cache is restored.""" + + @abstractmethod + def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): + """Restore runtime state from a big-page checkpoint.""" + + @abstractmethod + def restore_small_page_state(self, req: "InferReq", small_page_buffers): + """Restore runtime state from a small-page checkpoint.""" + + @abstractmethod + def copy_runtime_state_to_cache( + self, + req_indexes: Union[List[int], torch.Tensor], + buffer_indexes: List[int], + state_cache_manager, + ): + """Copy selected request runtime states into host-side page buffers.""" diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index 967bc9bf7a..e33a9d7d65 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, List, Union import torch @@ -7,14 +7,16 @@ from lightllm.common.linear_att_cache_manager.linear_att_buffer_manager import LinearAttCacheManager from lightllm.utils.envs_utils import get_env_start_args -from .base import ReqManager +from .hybrid_att import HybridAttentionReqManager if TYPE_CHECKING: from lightllm.server.router.model_infer.infer_batch import InferReq -class ReqManagerForMamba(ReqManager): +class ReqManagerForMamba(HybridAttentionReqManager): + is_linear_attention = True + def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig): super().__init__(max_request_num, max_sequence_length, mem_manager) self.mtp_step = get_env_start_args().mtp_step @@ -50,6 +52,47 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_con ) return + def create_state_cache_manager(self, size: int): + return LinearAttCacheManager(size=size, linear_config=self.linear_config) + + def init_hybrid_attention_state(self, req: "InferReq"): + return self.init_linear_att_state(req) + + def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): + return self.copy_big_page_buffer_to_linear_att_state(big_page_buffer_idx=big_page_buffer_idx, req=req) + + def restore_small_page_state(self, req: "InferReq", small_page_buffers): + return self.copy_small_page_buffer_to_linear_att_state( + req=req, + linear_att_small_page_buffers=small_page_buffers, + ) + + def copy_runtime_state_to_cache( + self, + req_indexes: Union[List[int], torch.Tensor], + buffer_indexes: list[int], + state_cache_manager: LinearAttCacheManager, + ): + assert len(req_indexes) == len(buffer_indexes) + if not any(buffer_idx != -1 for buffer_idx in buffer_indexes): + return + + from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer + + if not isinstance(req_indexes, torch.Tensor): + req_indexes = torch.tensor(req_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) + buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) + copy_linear_att_state_to_kv_buffer( + b_req_idx=req_indexes, + big_page_buffer_ids=buffer_indexes, + gpu_conv_state=self.req_to_conv_state.buffer, + gpu_ssm_state=self.req_to_ssm_state.buffer, + cpu_kv_conv_state=state_cache_manager.conv_state_cache.buffer, + cpu_kv_ssm_state=state_cache_manager.ssm_state_cache.buffer, + mtp_step=self.mtp_step, + ) + return + def init_linear_att_state(self, req: "InferReq"): conv_index = req.req_idx ssm_start = req.req_idx * (self.mtp_step + 1) diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py new file mode 100644 index 0000000000..1910017650 --- /dev/null +++ b/lightllm/common/req_manager/sliding_window.py @@ -0,0 +1,172 @@ +import collections +from typing import TYPE_CHECKING, List, Optional, Union + +import torch + +from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( + commit_sliding_window_state, + prepare_sliding_window_indexes, +) + +from .hybrid_att import HybridAttentionReqManager + + +if TYPE_CHECKING: + from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager + from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig + from lightllm.server.router.model_infer.infer_batch import InferReq + + +class SlidingWindowStateCacheManager: + """Pinned host storage for request-level sliding-window checkpoints.""" + + def __init__(self, size: int, sliding_config: "SlidingWindowCacheConfig", keep_num: int = 0): + self.size = size + self.keep_num = keep_num + self.sliding_config = sliding_config + assert 0 <= keep_num <= size + self.state_cache = torch.empty( + (size, *sliding_config.get_state_shape()), + dtype=sliding_config.dtype, + device="cpu", + pin_memory=True, + ) + self.clear_to_init_state() + + def get_state_cache(self, buffer_idx: int): + return self.state_cache[buffer_idx] + + def alloc_one_state_cache(self) -> Optional[int]: + return None if not self.free_list else self.free_list.popleft() + + def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: + if need_size > len(self.free_list): + return None + return [self.free_list.popleft() for _ in range(need_size)] + + def free_state_cache(self, free_indexes: List[int]): + alloc_upper_bound = self.size - self.keep_num + assert all(0 <= idx < alloc_upper_bound for idx in free_indexes) + self.free_list.extend(free_indexes) + assert len(self.free_list) <= alloc_upper_bound + + def get_free_cache_num(self): + return len(self.free_list) + + def get_used_cache_num(self): + return self.size - len(self.free_list) + + def clear_to_init_state(self): + self.state_cache.zero_() + self.free_list = collections.deque(range(self.size - self.keep_num)) + + +class ReqManagerForSlidingWindow(HybridAttentionReqManager): + """Token-granular virtual addresses plus request-granular sliding KV.""" + + def __init__( + self, + max_request_num: int, + max_sequence_length: int, + mem_manager: Optional["HybridSlidingMemoryManager"], + sliding_config: "SlidingWindowCacheConfig", + scratch_token_num: int, + ): + super().__init__(max_request_num, max_sequence_length, mem_manager) + self.sliding_config = sliding_config + self.sliding_window = sliding_config.sliding_window + self.scratch_token_num = scratch_token_num + self.runtime_token_num = (max_request_num + 1) * self.sliding_window + self.scratch_start = self.runtime_token_num + self.req_to_sliding_window = torch.zeros( + ( + sliding_config.sliding_layer_num, + self.runtime_token_num + scratch_token_num, + 2 * sliding_config.sliding_head_num, + sliding_config.sliding_head_dim, + ), + dtype=sliding_config.dtype, + device="cuda", + ) + self.req_to_sliding_window_indexs = torch.zeros( + (max_request_num + 1, max_sequence_length), + dtype=torch.int32, + device="cuda", + ) + + def create_state_cache_manager(self, size: int): + return SlidingWindowStateCacheManager(size=size, sliding_config=self.sliding_config) + + def init_hybrid_attention_state(self, req: "InferReq"): + start = req.req_idx * self.sliding_window + self.req_to_sliding_window[:, start : start + self.sliding_window].zero_() + + def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): + self._restore_state(req.req_idx, self.mem_manager.hybrid_att_big_page_buffers, big_page_buffer_idx) + + def restore_small_page_state(self, req: "InferReq", small_page_buffers): + self._restore_state(req.req_idx, small_page_buffers, req.shared_kv_node.small_page_buffer_idx) + + def _restore_state(self, req_idx: int, state_cache_manager, buffer_idx: int): + start = req_idx * self.sliding_window + self.req_to_sliding_window[:, start : start + self.sliding_window].copy_( + state_cache_manager.get_state_cache(buffer_idx), + non_blocking=True, + ) + + def copy_runtime_state_to_cache( + self, + req_indexes: Union[List[int], torch.Tensor], + buffer_indexes: List[int], + state_cache_manager: SlidingWindowStateCacheManager, + ): + assert len(req_indexes) == len(buffer_indexes) + if isinstance(req_indexes, torch.Tensor): + req_indexes = req_indexes.tolist() + for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): + if buffer_idx == -1: + continue + start = req_idx * self.sliding_window + state_cache_manager.get_state_cache(buffer_idx).copy_( + self.req_to_sliding_window[:, start : start + self.sliding_window], + non_blocking=True, + ) + + def prepare_sliding_window(self, infer_state): + q_token_num = infer_state.input_ids.shape[0] + assert q_token_num <= self.scratch_token_num + prepare_sliding_window_indexes( + req_to_sliding_window_indexs=self.req_to_sliding_window_indexs, + b_req_idx=infer_state.b_req_idx, + b_seq_len=infer_state.b_seq_len, + b_q_seq_len=infer_state.b_q_seq_len, + b_q_start_loc=infer_state.b_q_start_loc, + sliding_window=self.sliding_window, + scratch_start=self.scratch_start, + max_q_seq_len=infer_state.max_q_seq_len, + ) + infer_state.sliding_window_mem_index = torch.arange( + self.scratch_start, + self.scratch_start + q_token_num, + dtype=torch.int64, + device="cuda", + ) + + def get_layer_kv(self, layer_index: int): + local_layer = self.sliding_config.get_sliding_layer_index(layer_index) + layer_buffer = self.req_to_sliding_window[local_layer] + head_num = self.sliding_config.sliding_head_num + return layer_buffer[:, :head_num], layer_buffer[:, head_num:] + + def commit_layer_state(self, layer_index: int, infer_state): + local_layer = self.sliding_config.get_sliding_layer_index(layer_index) + commit_sliding_window_state( + layer_buffer=self.req_to_sliding_window[local_layer], + b_req_idx=infer_state.b_req_idx, + b_seq_len=infer_state.b_seq_len, + b_q_seq_len=infer_state.b_q_seq_len, + b_q_start_loc=infer_state.b_q_start_loc, + sliding_window=self.sliding_window, + scratch_start=self.scratch_start, + max_q_seq_len=infer_state.max_q_seq_len, + ) diff --git a/lightllm/common/sliding_window_cache_manager/__init__.py b/lightllm/common/sliding_window_cache_manager/__init__.py new file mode 100644 index 0000000000..540971d4d8 --- /dev/null +++ b/lightllm/common/sliding_window_cache_manager/__init__.py @@ -0,0 +1,4 @@ +from .config import SlidingWindowCacheConfig + + +__all__ = ["SlidingWindowCacheConfig"] diff --git a/lightllm/common/sliding_window_cache_manager/config.py b/lightllm/common/sliding_window_cache_manager/config.py new file mode 100644 index 0000000000..5f06c33711 --- /dev/null +++ b/lightllm/common/sliding_window_cache_manager/config.py @@ -0,0 +1,73 @@ +import dataclasses +from typing import Dict, List + +import torch + + +@dataclasses.dataclass +class SlidingWindowCacheConfig: + """Physical cache layout for a full + sliding-window transformer.""" + + layer_types: List[str] + num_kv_shared_layers: 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 "sliding_attention" in self.layer_types + assert "full_attention" in self.layer_types + cutoff = len(self.layer_types) - self.num_kv_shared_layers + assert 0 < cutoff <= len(self.layer_types) + + self.sliding_layer_to_cache_index: Dict[int, int] = {} + self.full_layer_to_cache_index: Dict[int, int] = {} + owner_to_index = {"sliding_attention": {}, "full_attention": {}} + next_index = {"sliding_attention": 0, "full_attention": 0} + + for layer_idx, layer_type in enumerate(self.layer_types[:cutoff]): + assert layer_type in owner_to_index, f"unsupported attention layer type: {layer_type}" + owner_to_index[layer_type][layer_idx] = next_index[layer_type] + next_index[layer_type] += 1 + + for layer_idx, layer_type in enumerate(self.layer_types): + if layer_idx < cutoff: + owner = layer_idx + else: + owner = next(idx for idx in range(cutoff - 1, -1, -1) if self.layer_types[idx] == layer_type) + cache_index = owner_to_index[layer_type][owner] + if layer_type == "sliding_attention": + self.sliding_layer_to_cache_index[layer_idx] = cache_index + else: + self.full_layer_to_cache_index[layer_idx] = cache_index + + self.sliding_layer_num = next_index["sliding_attention"] + self.full_layer_num = next_index["full_attention"] + + @property + def all_layer_num(self): + return len(self.layer_types) + + def get_sliding_layer_index(self, layer_index: int) -> int: + return self.sliding_layer_to_cache_index[layer_index] + + def get_full_layer_index(self, layer_index: int) -> int: + return self.full_layer_to_cache_index[layer_index] + + 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): + elements = 1 + for dim in self.get_state_shape(): + elements *= dim + return elements * self.dtype.itemsize diff --git a/lightllm/models/gemma4/infer_struct.py b/lightllm/models/gemma4/infer_struct.py index 89ad34acbd..d18a470dc5 100644 --- a/lightllm/models/gemma4/infer_struct.py +++ b/lightllm/models/gemma4/infer_struct.py @@ -21,6 +21,7 @@ def __init__(self): # 则对应的 b_image_token_end 为 [0, 0, 4, 4, 0], # image token 可以看到自己当前这个token以及后面的 image token。 self.b_image_token_end = None + self.sliding_window_mem_index = None def init_some_extra_state(self, model): super().init_some_extra_state(model) @@ -40,6 +41,9 @@ def init_some_extra_state(self, model): if self.is_prefill: self.max_seq_len = self.max_kv_seq_len self._build_b_image_token_end() + else: + self.b_q_start_loc = self.b1_cu_q_seq_len[:-1] + self.req_manager.prepare_sliding_window(self) return def _build_b_image_token_end(self): diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 2f0c01dbf6..9ebf82230a 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -16,12 +16,9 @@ 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 KV is written to request-window state plus per-forward + scratch storage. """ def __init__(self, layer_num, network_config): @@ -57,15 +54,6 @@ def __init__(self, layer_num, network_config): 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) @@ -155,23 +143,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) @@ -182,7 +154,21 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight: Gemma4Trans def _post_cache_kv(self, cache_kv, infer_state, layer_weight): if self.is_kv_shared_ or cache_kv is None: return - return super()._post_cache_kv(cache_kv, infer_state, layer_weight) + if self.is_sliding: + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + layer_idx = infer_state.req_manager.sliding_config.get_sliding_layer_index(self.layer_num_) + destindex_copy_kv( + cache_kv, + infer_state.sliding_window_mem_index, + infer_state.req_manager.req_to_sliding_window[layer_idx], + ) + return + infer_state.mem_manager.operator.copy_kv_to_mem_manager( + layer_index=infer_state.mem_manager.get_full_cache_layer_index(self.layer_num_), + mem_index=infer_state.mem_index, + kv=cache_kv, + ) # ----- Attention kernels (sliding window + per-layer KV reshape) --- @@ -195,23 +181,9 @@ def _att_control(self): 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 + if self.is_sliding: + return infer_state.req_manager.get_layer_kv(layer_idx) + return infer_state.mem_manager.get_att_input_params(layer_index=layer_idx) def _context_attention_kernel( self, @@ -238,10 +210,12 @@ def _context_attention_kernel( 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.req_manager.req_to_sliding_window_indexs, infer_state.b_image_token_end, sliding_window=sw, ) + if not self.is_kv_shared_: + infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) return o_tensor.view(q.shape) # Full-attn layers: head_dim=512, no SWA, no image bidi — standard @@ -261,7 +235,16 @@ def _token_attention_kernel( _k, _v = self._get_layer_kv(infer_state) _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) + o_tensor = att_state.decode_att( + q=_q, + k=_k, + v=_v, + att_control=self._att_control(), + alloc_func=self.alloc_tensor, + req_to_token_indexs=(infer_state.req_manager.req_to_sliding_window_indexs if self.is_sliding else None), + ) + if self.is_sliding and not self.is_kv_shared_: + infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) 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..907ad10ed6 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -3,7 +3,9 @@ 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.sliding_window_cache_manager import SlidingWindowCacheConfig from lightllm.common.build_utils import repair_config from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo @@ -12,7 +14,7 @@ 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,6 +67,7 @@ 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 @@ -80,30 +83,49 @@ def _verify_params(self): f"num_kv_shared_layers={kv_shared} out of range for " f"num_hidden_layers={self.config['num_hidden_layers']}" ) + assert args.mtp_step == 0, "Gemma-4 hybrid sliding-window cache does not support MTP yet" + assert not args.enable_cpu_cache, "Gemma-4 hybrid sliding-window cache does not support CPU 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" return - 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"] + def _get_sliding_cache_config(self): + if hasattr(self, "sliding_cache_config"): + return self.sliding_cache_config 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, + self.sliding_cache_config = SlidingWindowCacheConfig( + layer_types=self.config["layer_types"], + num_kv_shared_layers=self.config.get("num_kv_shared_layers") or 0, + sliding_window=self.config["sliding_window"], + sliding_head_num=self.config["num_key_value_heads"] // self.tp_world_size_, + sliding_head_dim=self.config["head_dim"], + full_head_num=num_global_kv // self.tp_world_size_, + full_head_dim=self.config["global_head_dim"], 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(), + ) + return self.sliding_cache_config + + def _init_req_manager(self): + args = get_env_start_args() + create_max_seq_len = max(int(self.batch_max_tokens or 0), int(self.max_seq_length or 0)) + scratch_token_num = max( + int(self.batch_max_tokens or 0), + int(self.graph_max_batch_size or 0), + int(args.prefill_cudagraph_max_handle_token or 0) if args.enable_prefill_cudagraph else 0, + ) + self.req_manager = ReqManagerForSlidingWindow( + max_request_num=self.max_req_num, + max_sequence_length=create_max_seq_len, + mem_manager=None, + sliding_config=self._get_sliding_cache_config(), + scratch_token_num=scratch_token_num, + ) + + def _init_mem_manager(self): + self.mem_manager = HybridSlidingMemoryManager( + size=self.max_total_token_num, + sliding_config=self._get_sliding_cache_config(), mem_fraction=self.mem_fraction, ) return diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index 9729a8205c..a0993099eb 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -11,7 +11,7 @@ from lightllm.server.req_id_generator import convert_sub_id_to_group_id from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.config_utils import is_linear_att_mixed_model +from lightllm.utils.config_utils import is_hybrid_att_mixed_model, is_linear_att_mixed_model from lightllm.utils.kv_cache_utils import compute_token_list_hash from typing import Any, Dict, List, Union from lightllm.utils.log_utils import init_logger @@ -216,9 +216,9 @@ def init( self.post_init() args = get_env_start_args() - if is_linear_att_mixed_model(args.model_dir): + if is_hybrid_att_mixed_model(args.model_dir): self._fill_linear_att_token_hash() - if args.enable_cpu_cache: + if args.enable_cpu_cache and is_linear_att_mixed_model(args.model_dir): cpu_cache_hash_list, cpu_cache_page_len_list = self._calcu_linear_att_cpu_cache_page_len_list() self.token_hash_list = TokenHashList() self.token_hash_list.clear() diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 3585c223e2..2da7577064 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -8,7 +8,7 @@ from sortedcontainers import SortedDict from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union -from lightllm.common.req_manager import ReqManager, ReqManagerForMamba +from lightllm.common.req_manager import HybridAttentionReqManager, ReqManager, ReqManagerForMamba from lightllm.utils.infer_utils import mark_start, mark_end from lightllm.server.core.objs import Req, SamplingParams, FinishStatus, ShmReqManager from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode @@ -46,6 +46,7 @@ class InferenceContext: overlap_stream: torch.cuda.Stream = None # 一些情况下推理进程进行异步折叠操作的异步流对象。 cpu_kv_cache_stream: torch.cuda.Stream = None # 用 cpu kv cache 操作的 stream is_linear_att_mixed_model: bool = False # 标记模型是否是full att 混合 linear att 的混合模型。 + is_hybrid_att_mixed_model: bool = False def register( self, @@ -70,6 +71,7 @@ def register( self.vocab_size = vocab_size self.is_linear_att_mixed_model = isinstance(self.req_manager, ReqManagerForMamba) + self.is_hybrid_att_mixed_model = isinstance(self.req_manager, HybridAttentionReqManager) return @@ -133,7 +135,7 @@ def free_a_req_mem(self, free_token_index: List, req: "InferReq"): elif CacheTier.GPU not in req.cache_tiers: self._free_req_mem_without_radix_insert(free_token_index=free_token_index, req=req) else: - if not self.is_linear_att_mixed_model: + if not self.is_hybrid_att_mixed_model: self._full_att_free_req(free_token_index=free_token_index, req=req) else: self._linear_att_free_req(free_token_index=free_token_index, req=req) @@ -146,7 +148,7 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer shared_kv_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][shared_kv_len : req.cur_kv_len]) - if self.is_linear_att_mixed_model: + if self.is_hybrid_att_mixed_model: # 释放请求尾部尚未移交给 radix cache 的 linear attention 小页状态。 if req.tail_linear_att_small_page_buffer_id is not None: self.radix_cache.linear_att_small_page_buffers.free_state_cache( @@ -182,7 +184,7 @@ def _full_att_free_req(self, free_token_index: List, req: "InferReq"): return def _linear_att_free_req(self, free_token_index: List, req: "InferReq"): - assert g_infer_context.is_linear_att_mixed_model is True + assert g_infer_context.is_hybrid_att_mixed_model is True args = get_env_start_args() hash_page_size = args.linear_att_hash_page_size big_page_num = args.linear_att_page_block_num @@ -375,7 +377,7 @@ def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bo if prefill_need_token_num > can_alloc_token_num: break - if g_infer_context.is_linear_att_mixed_model: + if g_infer_context.is_hybrid_att_mixed_model: req._linear_match_radix_cache() else: req._match_radix_cache() @@ -396,14 +398,12 @@ def get_can_alloc_token_num(self): ) return self.req_manager.mem_manager.allocator.can_use_mem_size + radix_cache_unref_token_num - def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): - """ - 该函数用于在线性混合模型prefill后,如果存在大页匹配的情况下，将线性层状态复制到 - """ - if not self.is_linear_att_mixed_model: + def copy_hybrid_att_state_to_cache_buffer(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_mixed_model or self.radix_cache is None: return - # 大页对应的 linear att 的拷贝 + # Request-state snapshot at a big-page boundary. big_page_token_num = self.args.linear_att_hash_page_size * self.args.linear_att_page_block_num big_page_buffer_ids = [] for req in reqs: @@ -418,27 +418,15 @@ def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L big_page_buffer_ids.append(-1) assert len(b_req_idx) == len(big_page_buffer_ids) - if any(buffer_id != -1 for buffer_id in big_page_buffer_ids): - big_page_buffer_ids = torch.tensor( - big_page_buffer_ids, dtype=torch.int32, requires_grad=False, device="cpu" - ) - big_page_buffer_ids = big_page_buffer_ids.cuda(non_blocking=True) - - from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer - - copy_linear_att_state_to_kv_buffer( - b_req_idx=b_req_idx, - big_page_buffer_ids=big_page_buffer_ids, - gpu_conv_state=self.req_manager.req_to_conv_state.buffer, - gpu_ssm_state=self.req_manager.req_to_ssm_state.buffer, - cpu_kv_conv_state=self.radix_cache.linear_att_big_page_buffers.conv_state_cache.buffer, - cpu_kv_ssm_state=self.radix_cache.linear_att_big_page_buffers.ssm_state_cache.buffer, - mtp_step=self.args.mtp_step, - ) + self.req_manager.copy_runtime_state_to_cache( + req_indexes=b_req_idx, + buffer_indexes=big_page_buffer_ids, + state_cache_manager=self.radix_cache.linear_att_big_page_buffers, + ) - assert not self.args.disable_chunked_prefill, "chunked prefill mode must be enabled for linear att mixed model" + assert not self.args.disable_chunked_prefill, "chunked prefill must be enabled for hybrid attention models" - # tail small page 的linear att 状态的存储 + # Request-state snapshot at the final small-page boundary. for req in reqs: # 判断本次prefill 完以后 kv 的长度是否到达linear att 块存储的临界点。 if req.get_chuncked_input_token_len() == req.linear_att_cache_len: @@ -449,23 +437,17 @@ def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L self.radix_cache.linear_att_small_page_buffers.alloc_one_state_cache() ) if req.tail_linear_att_small_page_buffer_id is not None: - conv_src_idx = req.req_idx - ssm_src_idx = req.req_idx * (self.args.mtp_step + 1) - conv_cache_width = self.req_manager.linear_config.get_conv_state_shape()[-1] - gpu_conv_state = self.req_manager.req_to_conv_state.buffer[ - :, conv_src_idx, ..., :conv_cache_width - ] - gpu_ssm_state = self.req_manager.req_to_ssm_state.buffer[:, ssm_src_idx, ...] dst_buffer_idx = req.tail_linear_att_small_page_buffer_id - - dst_conv_state, dst_ssm_state = self.radix_cache.linear_att_small_page_buffers.get_state_cache( - buffer_idx=dst_buffer_idx + self.req_manager.copy_runtime_state_to_cache( + req_indexes=[req.req_idx], + buffer_indexes=[dst_buffer_idx], + state_cache_manager=self.radix_cache.linear_att_small_page_buffers, ) - # TODO 对于非连续对象调用 copy_ 效率并不高 - dst_conv_state.copy_(gpu_conv_state, non_blocking=True) - dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) return + # Compatibility for out-of-tree backends while the hybrid name becomes canonical. + copy_linear_att_state_to_cache_buffer = copy_hybrid_att_state_to_cache_buffer + g_infer_context = InferenceContext() @@ -611,7 +593,7 @@ def __init__( else: self.decode_need_token_num = self._normal_decode_need_token_num - if g_infer_context.is_linear_att_mixed_model: + if g_infer_context.is_hybrid_att_mixed_model: self.get_chuncked_input_token_len = self.get_chuncked_input_token_len_for_linear_att self.get_chuncked_input_token_ids = self.get_chuncked_input_token_ids_for_linear_att @@ -623,7 +605,7 @@ def __init__( self.generator.manual_seed(self.sampling_param.shm_param.seed) if init_prefix_cache: - if g_infer_context.is_linear_att_mixed_model: + if g_infer_context.is_hybrid_att_mixed_model: self._linear_match_radix_cache() else: self._match_radix_cache() @@ -653,7 +635,7 @@ def _init_all_state(self): self.finish_status = FinishStatus() # 申请线性att混合模型使用的缓存资源 - if g_infer_context.is_linear_att_mixed_model: + if g_infer_context.is_hybrid_att_mixed_model: linear_block_num = self.shm_req.linear_att_token_hash_list.size self.linear_att_cache_len = linear_block_num * self.args.linear_att_hash_page_size self.linear_att_len_to_big_page_id = SortedDict() @@ -662,7 +644,7 @@ def _init_all_state(self): def _match_radix_cache(self): assert ( - g_infer_context.is_linear_att_mixed_model is False + g_infer_context.is_hybrid_att_mixed_model is False ), "current _match_radix_cache does not support linear att hybrid model, to do..." enable_prompt_cache = (not self.sampling_param.disable_prompt_cache) and g_infer_context.radix_cache is not None if enable_prompt_cache and self.get_cur_total_len() > 1 and self.cur_kv_len == 0: @@ -683,7 +665,7 @@ def _match_radix_cache(self): def _linear_match_radix_cache(self): assert ( - g_infer_context.is_linear_att_mixed_model is True + g_infer_context.is_hybrid_att_mixed_model is True ), "current _linear_match_radix_cache only support linear att hybrid model, to do..." enable_prompt_cache = (not self.sampling_param.disable_prompt_cache) and g_infer_context.radix_cache is not None linear_hash_list = self.shm_req.linear_att_token_hash_list.get_all() @@ -715,7 +697,7 @@ def _linear_match_radix_cache(self): self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 assert self.tail_linear_att_small_page_buffer_id is None # 恢复linear att 状态 - g_infer_context.req_manager.copy_big_page_buffer_to_linear_att_state( + g_infer_context.req_manager.restore_big_page_state( big_page_buffer_idx=share_node.big_page_buffer_idx, req=self ) else: @@ -730,9 +712,9 @@ def _linear_match_radix_cache(self): self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 assert self.tail_linear_att_small_page_buffer_id is None # 恢复linear att 状态 - g_infer_context.req_manager.copy_small_page_buffer_to_linear_att_state( + g_infer_context.req_manager.restore_small_page_state( req=self, - linear_att_small_page_buffers=g_infer_context.radix_cache.linear_att_small_page_buffers, + small_page_buffers=g_infer_context.radix_cache.linear_att_small_page_buffers, ) else: # 如果 大页本质是被启用的，则需要使用小页的匹配结果, 将小页的kv 复制到的新申请的kv位置，同时释放 @@ -763,9 +745,9 @@ def _linear_match_radix_cache(self): ) self.shared_kv_node = share_node # 只是为了保证 copy_small_page_buffer_to_linear_att_state 正确调用 - g_infer_context.req_manager.copy_small_page_buffer_to_linear_att_state( + g_infer_context.req_manager.restore_small_page_state( req=self, - linear_att_small_page_buffers=g_infer_context.radix_cache.linear_att_small_page_buffers, + small_page_buffers=g_infer_context.radix_cache.linear_att_small_page_buffers, ) self.shared_kv_node = None @@ -789,7 +771,7 @@ def _linear_match_radix_cache(self): self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 assert self.tail_linear_att_small_page_buffer_id is None # 恢复linear att 状态 - g_infer_context.req_manager.copy_big_page_buffer_to_linear_att_state( + g_infer_context.req_manager.restore_big_page_state( big_page_buffer_idx=share_node.big_page_buffer_idx, req=self ) @@ -797,7 +779,7 @@ def _linear_match_radix_cache(self): if self.cur_kv_len == 0: # 说明没有任何命中 - g_infer_context.req_manager.init_linear_att_state(req=self) + g_infer_context.req_manager.init_hybrid_attention_state(req=self) return def is_master_req(self): diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 560c3b6de4..69f117a8d0 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -15,8 +15,7 @@ from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.logprobs_manager import PromptLogprobsCaptureManager from lightllm.common.basemodel.moe_route_info_manager import MoeRouteInfoManager -from lightllm.common.req_manager import ReqManagerForMamba -from lightllm.common.linear_att_cache_manager import LinearAttCacheManager +from lightllm.common.req_manager import HybridAttentionReqManager, ReqManagerForMamba from lightllm.server.router.dynamic_prompt.linear_att_radix_cache import LinearAttPagedRadixCache from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput @@ -152,19 +151,20 @@ def init_model(self, kvargs): self.model: TpPartBaseModel = self.model # for easy typing set_random_seed(2147483647) self.is_linear_att_mixed_model = isinstance(self.model.req_manager, ReqManagerForMamba) + self.is_hybrid_att_mixed_model = isinstance(self.model.req_manager, HybridAttentionReqManager) - if self.is_linear_att_mixed_model: - self.linear_att_cache_manager = LinearAttCacheManager( - size=self.args.linear_att_cache_size, - linear_config=self.model.req_manager.linear_config, + if self.is_hybrid_att_mixed_model: + self.hybrid_att_cache_manager = self.model.req_manager.create_state_cache_manager( + size=self.args.linear_att_cache_size ) else: - self.linear_att_cache_manager = None + self.hybrid_att_cache_manager = None + self.linear_att_cache_manager = self.hybrid_att_cache_manager if self.is_linear_att_mixed_model else None if not self.use_dynamic_prompt_cache: self.radix_cache = None else: - if self.is_linear_att_mixed_model: + if self.is_hybrid_att_mixed_model: self.radix_cache = LinearAttPagedRadixCache( unique_name=get_unique_server_name(), total_token_num=self.model.mem_manager.size, @@ -172,7 +172,7 @@ def init_model(self, kvargs): hash_page_size=self.args.linear_att_hash_page_size, big_page_num=self.args.linear_att_page_block_num, kv_cache_mem_manager=self.model.mem_manager, - linear_att_small_page_buffers=self.linear_att_cache_manager, + linear_att_small_page_buffers=self.hybrid_att_cache_manager, ) else: self.radix_cache = RadixCache( diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 4d09476849..50a0e66e8d 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -119,7 +119,7 @@ def prefill_normal( b_prefill_has_output_cpu=model_input.b_prefill_has_output_cpu, mask_func=self.prefill_mask_func, ) - g_infer_context.copy_linear_att_state_to_cache_buffer( + g_infer_context.copy_hybrid_att_state_to_cache_buffer( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) @@ -215,7 +215,7 @@ def prefill_mtp( target_model_output=model_output, target_next_token_ids=next_token_ids, ) - g_infer_context.copy_linear_att_state_to_cache_buffer( + g_infer_context.copy_hybrid_att_state_to_cache_buffer( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 9a81927bc1..63cbfe74e8 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -195,7 +195,7 @@ def prefill_normal( b_prefill_has_output_cpu=model_input.b_prefill_has_output_cpu, mask_func=None, ) - g_infer_context.copy_linear_att_state_to_cache_buffer( + g_infer_context.copy_hybrid_att_state_to_cache_buffer( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) @@ -316,8 +316,8 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer mask_func=None, ) - if g_infer_context.is_linear_att_mixed_model: - g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + if g_infer_context.is_hybrid_att_mixed_model: + g_infer_context.copy_hybrid_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() @@ -440,7 +440,7 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] target_next_token_ids=next_token_ids, ) if req_num > 0: - g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + g_infer_context.copy_hybrid_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() @@ -695,8 +695,8 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I target_next_token_ids1=target_next_token_ids_gpu1, ) - if req_num > 0 and g_infer_context.is_linear_att_mixed_model: - g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + if req_num > 0 and g_infer_context.is_hybrid_att_mixed_model: + g_infer_context.copy_hybrid_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 6695f0ec44..4dfabe66e9 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -465,6 +465,32 @@ 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: + try: + config_json = get_config_json(model_path) + llm_config = config_json.get("text_config", config_json) + model_type = config_json.get("model_type") or llm_config.get("model_type") + layer_types = set(llm_config.get("layer_types", [])) + # Keep the shared-request ABI opt-in aligned with models that actually + # instantiate ReqManagerForSlidingWindow. Other architectures may use + # the same layer-type strings while retaining token-granular KV. + return ( + model_type in {"gemma4", "gemma4_text"} + and { + "full_attention", + "sliding_attention", + }.issubset(layer_types) + ) + except Exception: + logger.info(f"model path: {model_path} does not have hybrid sliding-window attention") + return False + + +def is_hybrid_att_mixed_model(model_path: str) -> bool: + 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]: """Get model type from config.json""" try: From ffd84b0bd10006c778a8132098ab58e8a4bcb5cc Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Sun, 6 Sep 2026 16:28:22 +0800 Subject: [PATCH 02/14] refactor: resolve decode kv indexes in request manager --- lightllm/common/basemodel/attention/triton/fp.py | 10 +--------- .../gqa/flash_decoding/gqa_flash_decoding.py | 6 +----- lightllm/common/req_manager/base.py | 4 ++++ lightllm/common/req_manager/sliding_window.py | 5 +++++ .../gemma4/layer_infer/transformer_layer_infer.py | 9 +-------- 5 files changed, 12 insertions(+), 22 deletions(-) diff --git a/lightllm/common/basemodel/attention/triton/fp.py b/lightllm/common/basemodel/attention/triton/fp.py index d5f5cce39b..e7ce66c774 100644 --- a/lightllm/common/basemodel/attention/triton/fp.py +++ b/lightllm/common/basemodel/attention/triton/fp.py @@ -114,7 +114,6 @@ def decode_att( v: torch.Tensor, att_control: AttControl = AttControl(), alloc_func=torch.empty, - req_to_token_indexs=None, ): if att_control.use_alibi: assert att_control.use_sliding_window is False, "alibi + sliding_window not supported" @@ -134,12 +133,7 @@ def decode_att( return self._normal_decode_flash_decoding_att(q=q, k=k, v=v, alloc_func=alloc_func) elif q_head_num > k_head_num: return self._normal_decode_gqa_flash_decoding_att( - q=q, - k=k, - v=v, - att_control=att_control, - alloc_func=alloc_func, - req_to_token_indexs=req_to_token_indexs, + q=q, k=k, v=v, att_control=att_control, alloc_func=alloc_func ) else: raise NotImplementedError("error") @@ -201,7 +195,6 @@ def _normal_decode_gqa_flash_decoding_att( v: torch.Tensor, att_control: AttControl = AttControl(), alloc_func=torch.empty, - req_to_token_indexs=None, ): from ...triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import ( gqa_token_decode_attention_flash_decoding, @@ -222,7 +215,6 @@ def _normal_decode_gqa_flash_decoding_att( out=out, alloc_tensor_func=alloc_func, sliding_window=sliding_window, - req_to_token_indexs=req_to_token_indexs, ) return out diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py index d4a806d6f4..c52da54553 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py @@ -9,7 +9,6 @@ def gqa_token_decode_attention_flash_decoding( out=None, alloc_tensor_func=torch.empty, sliding_window=(-1, -1), - req_to_token_indexs=None, ): batch_size = infer_state.batch_size q_head_num, head_dim = q.shape[1], q.shape[2] @@ -35,14 +34,11 @@ def gqa_token_decode_attention_flash_decoding( mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device="cuda") mid_o_logexpsum = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device="cuda") - if req_to_token_indexs is None: - req_to_token_indexs = infer_state.req_manager.req_to_token_indexs - flash_decode_stage1( q=q.view(calcu_shape1), k=cache_k, v=cache_v, - Req_to_tokens=req_to_token_indexs, + Req_to_tokens=infer_state.req_manager.get_decode_kv_indexs(use_sliding_window=sliding_window != (-1, -1)), B_req_idx=infer_state.b_req_idx, B_Seqlen=infer_state.b_seq_len, max_len_in_batch=infer_state.max_kv_seq_len, diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index 372ca9deb3..6d523cc101 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -70,6 +70,10 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana def alloc(self): return self.req_list.alloc() + def get_decode_kv_indexs(self, use_sliding_window: bool = False): + """Return the physical KV index table used by decode attention.""" + return self.req_to_token_indexs + def free(self, free_req_indexes: List[int], free_token_index): for req_index in free_req_indexes: self.req_list.free(req_index) diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index 1910017650..af29f1a065 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -158,6 +158,11 @@ def get_layer_kv(self, layer_index: int): head_num = self.sliding_config.sliding_head_num return layer_buffer[:, :head_num], layer_buffer[:, head_num:] + def get_decode_kv_indexs(self, use_sliding_window: bool = False): + if use_sliding_window: + return self.req_to_sliding_window_indexs + return super().get_decode_kv_indexs(use_sliding_window=use_sliding_window) + def commit_layer_state(self, layer_index: int, infer_state): local_layer = self.sliding_config.get_sliding_layer_index(layer_index) commit_sliding_window_state( diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 9ebf82230a..6375d8f872 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -235,14 +235,7 @@ def _token_attention_kernel( _k, _v = self._get_layer_kv(infer_state) _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, - req_to_token_indexs=(infer_state.req_manager.req_to_sliding_window_indexs if self.is_sliding else None), - ) + 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 and not self.is_kv_shared_: infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) return o_tensor.view(q.shape) From f80ec447febe252352384cc23721a18a6fb801b9 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Sun, 6 Sep 2026 16:41:04 +0800 Subject: [PATCH 03/14] refactor: keep sliding runtime state model-local --- .../gqa/flash_decoding/gqa_flash_decoding.py | 2 +- .../hybrid_sliding_mem_manager.py | 4 ++-- lightllm/common/req_manager/base.py | 4 ---- lightllm/common/req_manager/sliding_window.py | 10 ++-------- lightllm/models/gemma4/infer_struct.py | 16 ++++++++++++++++ 5 files changed, 21 insertions(+), 15 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py index c52da54553..59a7d4f742 100644 --- a/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py +++ b/lightllm/common/basemodel/triton_kernel/att/decode_att/gqa/flash_decoding/gqa_flash_decoding.py @@ -38,7 +38,7 @@ def gqa_token_decode_attention_flash_decoding( q=q.view(calcu_shape1), k=cache_k, v=cache_v, - Req_to_tokens=infer_state.req_manager.get_decode_kv_indexs(use_sliding_window=sliding_window != (-1, -1)), + Req_to_tokens=infer_state.req_manager.req_to_token_indexs, B_req_idx=infer_state.b_req_idx, B_Seqlen=infer_state.b_seq_len, max_len_in_batch=infer_state.max_kv_seq_len, 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 index 3caa7dcaab..f1fee86baf 100644 --- a/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -48,8 +48,8 @@ def _free_buffers(self): self.linear_att_big_page_buffers = None def write_to_shm(self, req_manager): - # Host-side checkpoints are local to the inference process. Excluding - # them also preserves their pinned allocation during serialization. + # Page checkpoints are process-local GPU runtime state and must not be + # serialized into the shared-memory view of the memory manager. big_page_buffers = self.hybrid_att_big_page_buffers legacy_big_page_buffers = self.linear_att_big_page_buffers self.hybrid_att_big_page_buffers = None diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index 6d523cc101..372ca9deb3 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -70,10 +70,6 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana def alloc(self): return self.req_list.alloc() - def get_decode_kv_indexs(self, use_sliding_window: bool = False): - """Return the physical KV index table used by decode attention.""" - return self.req_to_token_indexs - def free(self, free_req_indexes: List[int], free_token_index): for req_index in free_req_indexes: self.req_list.free(req_index) diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index af29f1a065..c5410f29b7 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -18,7 +18,7 @@ class SlidingWindowStateCacheManager: - """Pinned host storage for request-level sliding-window checkpoints.""" + """GPU storage for request-level sliding-window checkpoints.""" def __init__(self, size: int, sliding_config: "SlidingWindowCacheConfig", keep_num: int = 0): self.size = size @@ -28,8 +28,7 @@ def __init__(self, size: int, sliding_config: "SlidingWindowCacheConfig", keep_n self.state_cache = torch.empty( (size, *sliding_config.get_state_shape()), dtype=sliding_config.dtype, - device="cpu", - pin_memory=True, + device="cuda", ) self.clear_to_init_state() @@ -158,11 +157,6 @@ def get_layer_kv(self, layer_index: int): head_num = self.sliding_config.sliding_head_num return layer_buffer[:, :head_num], layer_buffer[:, head_num:] - def get_decode_kv_indexs(self, use_sliding_window: bool = False): - if use_sliding_window: - return self.req_to_sliding_window_indexs - return super().get_decode_kv_indexs(use_sliding_window=use_sliding_window) - def commit_layer_state(self, layer_index: int, infer_state): local_layer = self.sliding_config.get_sliding_layer_index(layer_index) commit_sliding_window_state( diff --git a/lightllm/models/gemma4/infer_struct.py b/lightllm/models/gemma4/infer_struct.py index d18a470dc5..cfbe745b52 100644 --- a/lightllm/models/gemma4/infer_struct.py +++ b/lightllm/models/gemma4/infer_struct.py @@ -1,3 +1,5 @@ +import copy + import torch from lightllm.common.basemodel import InferStateInfo from lightllm.models.gemma4.triton_kernel.build_b_image_token_end import build_b_image_token_end @@ -46,6 +48,20 @@ def init_some_extra_state(self, model): self.req_manager.prepare_sliding_window(self) return + def init_att_state(self): + if not self.is_prefill: + # Keep the common attention path unchanged: its decode kernels read + # req_to_token_indexs from infer_state.req_manager. The sliding + # state receives a shallow model-side view whose table addresses + # the request-window KV buffer; the full-attention state continues + # to use this infer state and the virtual token table. + sliding_infer_state = copy.copy(self) + sliding_req_manager = copy.copy(self.req_manager) + sliding_req_manager.req_to_token_indexs = self.req_manager.req_to_sliding_window_indexs + sliding_infer_state.req_manager = sliding_req_manager + self.decode_att_state.infer_state = sliding_infer_state + return super().init_att_state() + 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) From 87ee9b413974d4448daac67fede041ca10620193 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Sun, 6 Sep 2026 10:26:56 +0000 Subject: [PATCH 04/14] fix: harden hybrid sliding-window state lifecycle --- .../hybrid_sliding_mem_manager.py | 89 +++++++--- .../kv_cache_mem_manager/operator/__init__.py | 1 - lightllm/common/req_manager/__init__.py | 3 +- lightllm/common/req_manager/hybrid_att.py | 19 +-- lightllm/common/req_manager/linear_att.py | 27 +-- lightllm/common/req_manager/sliding_window.py | 74 ++------ .../sliding_window_cache_manager/__init__.py | 3 +- .../sliding_window_cache_manager/config.py | 44 +---- .../state_cache.py | 44 +++++ lightllm/models/gemma4/kv_layout.py | 20 +++ .../layer_infer/transformer_layer_infer.py | 26 ++- lightllm/models/gemma4/model.py | 8 +- .../server/router/model_infer/infer_batch.py | 24 ++- .../model_infer/mode_backend/base_backend.py | 7 +- .../mode_backend/chunked_prefill/impl.py | 4 +- .../mode_backend/dp_backend/impl.py | 8 +- test/kernel/test_sliding_window_state.py | 161 ++++++++++++++++++ test/utils/test_sliding_window_cache.py | 130 ++++++++++++++ .../mode_backend/test_multi_level_kv_cache.py | 18 +- 19 files changed, 523 insertions(+), 187 deletions(-) create mode 100644 lightllm/common/sliding_window_cache_manager/state_cache.py create mode 100644 lightllm/models/gemma4/kv_layout.py create mode 100644 test/kernel/test_sliding_window_state.py create mode 100644 test/utils/test_sliding_window_cache.py 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 index f1fee86baf..58a376fa84 100644 --- a/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -1,11 +1,17 @@ +import torch +import torch.distributed as dist import triton -from lightllm.common.req_manager.sliding_window import SlidingWindowStateCacheManager +from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger +from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory from .mem_manager import MemoryManager from .operator.hybrid_sliding import HybridSlidingMemOperator +logger = init_logger(__name__) + class HybridSlidingMemoryManager(MemoryManager): """Token-granular full KV plus request-granular sliding-window KV.""" @@ -14,6 +20,10 @@ class HybridSlidingMemoryManager(MemoryManager): def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): self.sliding_config = sliding_config + args = get_env_start_args() + self.enable_prompt_cache = args.use_dynamic_prompt_cache + self.small_page_num = args.linear_att_cache_size if self.enable_prompt_cache else 0 + 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, @@ -24,17 +34,67 @@ def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): mem_fraction=mem_fraction, ) + def _big_page_num(self, token_num): + return max(1, triton.cdiv(token_num, self.big_page_token_num)) if self.enable_prompt_cache else 0 + + def _cache_nbytes(self, token_num): + # Runtime windows already exist when profiling. Reserve BOTH GPU page + # pools here, plus the full-KV hold token and the final partial big page. + return (token_num + 1) * self.get_cell_size() + ( + self.small_page_num + self._big_page_num(token_num) + ) * self.sliding_config.get_state_nbytes() + + def _profile_token_num(self, available_bytes): + if self._cache_nbytes(1) > available_bytes: + raise ValueError( + "Insufficient GPU memory for sliding-window checkpoints and full KV: " + f"{available_bytes / 1024 ** 3:.2f} GiB available, " + f"{self.small_page_num} small pages at " + f"{self.sliding_config.get_state_nbytes() / 1024 ** 2:.2f} MiB/page. " + "Reduce --linear_att_cache_size or --running_max_req_size." + ) + low, high = 1, available_bytes // self.get_cell_size() + while low < high: + mid = (low + high + 1) // 2 + if self._cache_nbytes(mid) <= available_bytes: + low = mid + else: + high = mid - 1 + return low + + def profile_size(self, mem_fraction): + torch.cuda.empty_cache() + world_size = dist.get_world_size() + available_memory = get_available_gpu_memory(world_size) + if self.size is None: + available_memory -= get_total_gpu_memory() * (1 - mem_fraction) + self.size = self._profile_token_num(int(available_memory * 1024 ** 3)) + if world_size > 1: + size_tensor = torch.tensor(self.size, dtype=torch.int64, device="cuda") + dist.all_reduce(size_tensor, op=dist.ReduceOp.MIN) + self.size = size_tensor.item() + elif self._cache_nbytes(self.size) > int(available_memory * 1024 ** 3): + raise ValueError( + "Requested full KV and sliding-window checkpoints exceed available GPU memory; " + "reduce --max_total_token_num, --linear_att_cache_size or --running_max_req_size." + ) + logger.info( + f"Sliding-window cache budget: {self.size} full-KV tokens, " + f"{self._big_page_num(self.size)} big pages, {self.small_page_num} small pages, " + f"{self._cache_nbytes(self.size) / 1024 ** 3:.2f} GiB (runtime windows already allocated)" + ) + def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): super()._init_buffers(size, dtype, head_num, head_dim, layer_num) - big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size + # Keep the existing radix-cache contract; no second alias is needed. + self.linear_att_big_page_buffers = SlidingWindowStateCacheManager( + size=self._big_page_num(size), + sliding_config=self.sliding_config, ) - self.hybrid_att_big_page_buffers = SlidingWindowStateCacheManager( - size=max(1, triton.cdiv(self.size, big_page_token_num)), + self.sliding_small_page_buffers = SlidingWindowStateCacheManager( + size=self.small_page_num, sliding_config=self.sliding_config, ) - # Compatibility with the unchanged big/small-page radix implementation. - self.linear_att_big_page_buffers = self.hybrid_att_big_page_buffers def get_att_input_params(self, layer_index: int): return super().get_att_input_params(self.sliding_config.get_full_layer_index(layer_index)) @@ -44,18 +104,5 @@ def get_full_cache_layer_index(self, layer_index: int): def _free_buffers(self): super()._free_buffers() - self.hybrid_att_big_page_buffers = None - self.linear_att_big_page_buffers = None - - def write_to_shm(self, req_manager): - # Page checkpoints are process-local GPU runtime state and must not be - # serialized into the shared-memory view of the memory manager. - big_page_buffers = self.hybrid_att_big_page_buffers - legacy_big_page_buffers = self.linear_att_big_page_buffers - self.hybrid_att_big_page_buffers = None self.linear_att_big_page_buffers = None - try: - return super().write_to_shm(req_manager) - finally: - self.hybrid_att_big_page_buffers = big_page_buffers - self.linear_att_big_page_buffers = legacy_big_page_buffers + self.sliding_small_page_buffers = None diff --git a/lightllm/common/kv_cache_mem_manager/operator/__init__.py b/lightllm/common/kv_cache_mem_manager/operator/__init__.py index 26740d567d..85c37ad39b 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/__init__.py +++ b/lightllm/common/kv_cache_mem_manager/operator/__init__.py @@ -2,7 +2,6 @@ from .normal import NormalMemOperator from .quant import QuantScaleMemOperator, PPLInt4KVMemOperator, PPLInt8KVMemOperator from .linear_att import LinearAttMemOperator -from .hybrid_sliding import HybridSlidingMemOperator from .deepseek import ( Deepseek2MemOperator, Deepseek3_2MemOperator, diff --git a/lightllm/common/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index 11b92fb15a..2d8f8113b4 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -2,7 +2,7 @@ from .hybrid_att import HybridAttentionReqManager from .linear_att import ReqManagerForMamba from .req_sampling_params import ReqSamplingParamsManager -from .sliding_window import ReqManagerForSlidingWindow, SlidingWindowStateCacheManager +from .sliding_window import ReqManagerForSlidingWindow __all__ = [ "ReqManager", @@ -10,5 +10,4 @@ "ReqManagerForMamba", "ReqManagerForSlidingWindow", "ReqSamplingParamsManager", - "SlidingWindowStateCacheManager", ] diff --git a/lightllm/common/req_manager/hybrid_att.py b/lightllm/common/req_manager/hybrid_att.py index d24b427ea6..49cc5b6b8c 100644 --- a/lightllm/common/req_manager/hybrid_att.py +++ b/lightllm/common/req_manager/hybrid_att.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List, Union +from typing import TYPE_CHECKING, List import torch @@ -18,11 +18,9 @@ class HybridAttentionReqManager(ReqManager, ABC): this interface and may have a different physical granularity. """ - is_linear_attention = False - @abstractmethod def create_state_cache_manager(self, size: int): - """Create checkpoint storage used by request-state page boundaries.""" + """Return checkpoint storage used by request-state page boundaries.""" @abstractmethod def init_hybrid_attention_state(self, req: "InferReq"): @@ -37,10 +35,9 @@ def restore_small_page_state(self, req: "InferReq", small_page_buffers): """Restore runtime state from a small-page checkpoint.""" @abstractmethod - def copy_runtime_state_to_cache( - self, - req_indexes: Union[List[int], torch.Tensor], - buffer_indexes: List[int], - state_cache_manager, - ): - """Copy selected request runtime states into host-side page buffers.""" + def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): + """Save selected checkpoints; CPU request IDs avoid device-to-host synchronization.""" + + @abstractmethod + def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffers): + """Save a request's final small-page checkpoint in the layout's storage.""" diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index e33a9d7d65..f21a2ec1bd 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, List, Union +from typing import TYPE_CHECKING, List import torch @@ -15,8 +15,6 @@ class ReqManagerForMamba(HybridAttentionReqManager): - is_linear_attention = True - def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig): super().__init__(max_request_num, max_sequence_length, mem_manager) self.mtp_step = get_env_start_args().mtp_step @@ -67,23 +65,17 @@ def restore_small_page_state(self, req: "InferReq", small_page_buffers): linear_att_small_page_buffers=small_page_buffers, ) - def copy_runtime_state_to_cache( - self, - req_indexes: Union[List[int], torch.Tensor], - buffer_indexes: list[int], - state_cache_manager: LinearAttCacheManager, - ): - assert len(req_indexes) == len(buffer_indexes) + def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): + assert len(b_req_idx) == len(buffer_indexes) if not any(buffer_idx != -1 for buffer_idx in buffer_indexes): return from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer - if not isinstance(req_indexes, torch.Tensor): - req_indexes = torch.tensor(req_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) + state_cache_manager = self.mem_manager.linear_att_big_page_buffers copy_linear_att_state_to_kv_buffer( - b_req_idx=req_indexes, + b_req_idx=b_req_idx, big_page_buffer_ids=buffer_indexes, gpu_conv_state=self.req_to_conv_state.buffer, gpu_ssm_state=self.req_to_ssm_state.buffer, @@ -93,6 +85,15 @@ def copy_runtime_state_to_cache( ) return + def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffers: LinearAttCacheManager): + # Preserve main's small-page copies, including the MTP conv-state crop. + conv_cache_width = self.linear_config.get_conv_state_shape()[-1] + gpu_conv_state = self.req_to_conv_state.buffer[:, req_idx, ..., :conv_cache_width] + gpu_ssm_state = self.req_to_ssm_state.buffer[:, req_idx * (self.mtp_step + 1), ...] + dst_conv_state, dst_ssm_state = small_page_buffers.get_state_cache(buffer_idx=buffer_idx) + dst_conv_state.copy_(gpu_conv_state, non_blocking=True) + dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) + def init_linear_att_state(self, req: "InferReq"): conv_index = req.req_idx ssm_start = req.req_idx * (self.mtp_step + 1) diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index c5410f29b7..346c83ed07 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -1,5 +1,4 @@ -import collections -from typing import TYPE_CHECKING, List, Optional, Union +from typing import TYPE_CHECKING, List, Optional import torch @@ -7,6 +6,7 @@ commit_sliding_window_state, prepare_sliding_window_indexes, ) +from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager from .hybrid_att import HybridAttentionReqManager @@ -17,49 +17,6 @@ from lightllm.server.router.model_infer.infer_batch import InferReq -class SlidingWindowStateCacheManager: - """GPU storage for request-level sliding-window checkpoints.""" - - def __init__(self, size: int, sliding_config: "SlidingWindowCacheConfig", keep_num: int = 0): - self.size = size - self.keep_num = keep_num - self.sliding_config = sliding_config - assert 0 <= keep_num <= size - self.state_cache = torch.empty( - (size, *sliding_config.get_state_shape()), - dtype=sliding_config.dtype, - device="cuda", - ) - self.clear_to_init_state() - - def get_state_cache(self, buffer_idx: int): - return self.state_cache[buffer_idx] - - def alloc_one_state_cache(self) -> Optional[int]: - return None if not self.free_list else self.free_list.popleft() - - def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: - if need_size > len(self.free_list): - return None - return [self.free_list.popleft() for _ in range(need_size)] - - def free_state_cache(self, free_indexes: List[int]): - alloc_upper_bound = self.size - self.keep_num - assert all(0 <= idx < alloc_upper_bound for idx in free_indexes) - self.free_list.extend(free_indexes) - assert len(self.free_list) <= alloc_upper_bound - - def get_free_cache_num(self): - return len(self.free_list) - - def get_used_cache_num(self): - return self.size - len(self.free_list) - - def clear_to_init_state(self): - self.state_cache.zero_() - self.free_list = collections.deque(range(self.size - self.keep_num)) - - class ReqManagerForSlidingWindow(HybridAttentionReqManager): """Token-granular virtual addresses plus request-granular sliding KV.""" @@ -94,14 +51,15 @@ def __init__( ) def create_state_cache_manager(self, size: int): - return SlidingWindowStateCacheManager(size=size, sliding_config=self.sliding_config) + # Allocated with full KV and big pages, within the same GPU budget. + return self.mem_manager.sliding_small_page_buffers def init_hybrid_attention_state(self, req: "InferReq"): start = req.req_idx * self.sliding_window self.req_to_sliding_window[:, start : start + self.sliding_window].zero_() def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - self._restore_state(req.req_idx, self.mem_manager.hybrid_att_big_page_buffers, big_page_buffer_idx) + self._restore_state(req.req_idx, self.mem_manager.linear_att_big_page_buffers, big_page_buffer_idx) def restore_small_page_state(self, req: "InferReq", small_page_buffers): self._restore_state(req.req_idx, small_page_buffers, req.shared_kv_node.small_page_buffer_idx) @@ -113,23 +71,19 @@ def _restore_state(self, req_idx: int, state_cache_manager, buffer_idx: int): non_blocking=True, ) - def copy_runtime_state_to_cache( - self, - req_indexes: Union[List[int], torch.Tensor], - buffer_indexes: List[int], - state_cache_manager: SlidingWindowStateCacheManager, - ): + def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): assert len(req_indexes) == len(buffer_indexes) - if isinstance(req_indexes, torch.Tensor): - req_indexes = req_indexes.tolist() for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): if buffer_idx == -1: continue - start = req_idx * self.sliding_window - state_cache_manager.get_state_cache(buffer_idx).copy_( - self.req_to_sliding_window[:, start : start + self.sliding_window], - non_blocking=True, - ) + self.save_small_page_state(req_idx, buffer_idx, self.mem_manager.linear_att_big_page_buffers) + + def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffers: SlidingWindowStateCacheManager): + start = req_idx * self.sliding_window + small_page_buffers.get_state_cache(buffer_idx).copy_( + self.req_to_sliding_window[:, start : start + self.sliding_window], + non_blocking=True, + ) def prepare_sliding_window(self, infer_state): q_token_num = infer_state.input_ids.shape[0] diff --git a/lightllm/common/sliding_window_cache_manager/__init__.py b/lightllm/common/sliding_window_cache_manager/__init__.py index 540971d4d8..b05865d3bc 100644 --- a/lightllm/common/sliding_window_cache_manager/__init__.py +++ b/lightllm/common/sliding_window_cache_manager/__init__.py @@ -1,4 +1,5 @@ from .config import SlidingWindowCacheConfig +from .state_cache import SlidingWindowStateCacheManager -__all__ = ["SlidingWindowCacheConfig"] +__all__ = ["SlidingWindowCacheConfig", "SlidingWindowStateCacheManager"] diff --git a/lightllm/common/sliding_window_cache_manager/config.py b/lightllm/common/sliding_window_cache_manager/config.py index 5f06c33711..57005170b2 100644 --- a/lightllm/common/sliding_window_cache_manager/config.py +++ b/lightllm/common/sliding_window_cache_manager/config.py @@ -1,5 +1,5 @@ import dataclasses -from typing import Dict, List +from typing import Dict import torch @@ -8,8 +8,8 @@ class SlidingWindowCacheConfig: """Physical cache layout for a full + sliding-window transformer.""" - layer_types: List[str] - num_kv_shared_layers: int + 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 @@ -19,38 +19,12 @@ class SlidingWindowCacheConfig: def __post_init__(self): assert self.sliding_window > 0 - assert "sliding_attention" in self.layer_types - assert "full_attention" in self.layer_types - cutoff = len(self.layer_types) - self.num_kv_shared_layers - assert 0 < cutoff <= len(self.layer_types) - - self.sliding_layer_to_cache_index: Dict[int, int] = {} - self.full_layer_to_cache_index: Dict[int, int] = {} - owner_to_index = {"sliding_attention": {}, "full_attention": {}} - next_index = {"sliding_attention": 0, "full_attention": 0} - - for layer_idx, layer_type in enumerate(self.layer_types[:cutoff]): - assert layer_type in owner_to_index, f"unsupported attention layer type: {layer_type}" - owner_to_index[layer_type][layer_idx] = next_index[layer_type] - next_index[layer_type] += 1 - - for layer_idx, layer_type in enumerate(self.layer_types): - if layer_idx < cutoff: - owner = layer_idx - else: - owner = next(idx for idx in range(cutoff - 1, -1, -1) if self.layer_types[idx] == layer_type) - cache_index = owner_to_index[layer_type][owner] - if layer_type == "sliding_attention": - self.sliding_layer_to_cache_index[layer_idx] = cache_index - else: - self.full_layer_to_cache_index[layer_idx] = cache_index - - self.sliding_layer_num = next_index["sliding_attention"] - self.full_layer_num = next_index["full_attention"] - - @property - def all_layer_num(self): - return len(self.layer_types) + 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_sliding_layer_index(self, layer_index: int) -> int: return self.sliding_layer_to_cache_index[layer_index] diff --git a/lightllm/common/sliding_window_cache_manager/state_cache.py b/lightllm/common/sliding_window_cache_manager/state_cache.py new file mode 100644 index 0000000000..c60748bf97 --- /dev/null +++ b/lightllm/common/sliding_window_cache_manager/state_cache.py @@ -0,0 +1,44 @@ +import collections +from typing import List, Optional + +import torch + +from .config import SlidingWindowCacheConfig + + +class SlidingWindowStateCacheManager: + """GPU storage for immutable request-level sliding-window checkpoints.""" + + def __init__(self, size: int, sliding_config: SlidingWindowCacheConfig): + self.size = size + assert size >= 0 + self.state_cache = torch.empty( + (size, *sliding_config.get_state_shape()), dtype=sliding_config.dtype, device="cuda" + ) + self.clear_to_init_state() + + def get_state_cache(self, buffer_idx: int): + return self.state_cache[buffer_idx] + + def alloc_one_state_cache(self) -> Optional[int]: + return None if not self.free_list else self.free_list.popleft() + + def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: + if need_size > len(self.free_list): + return None + return [self.free_list.popleft() for _ in range(need_size)] + + def free_state_cache(self, free_indexes: List[int]): + assert all(0 <= idx < self.size for idx in free_indexes) + self.free_list.extend(free_indexes) + assert len(self.free_list) <= self.size + + def get_free_cache_num(self): + return len(self.free_list) + + def get_used_cache_num(self): + return self.size - len(self.free_list) + + def clear_to_init_state(self): + self.state_cache.zero_() + self.free_list = collections.deque(range(self.size)) diff --git a/lightllm/models/gemma4/kv_layout.py b/lightllm/models/gemma4/kv_layout.py new file mode 100644 index 0000000000..7f02998064 --- /dev/null +++ b/lightllm/models/gemma4/kv_layout.py @@ -0,0 +1,20 @@ +def get_kv_cache_layout(config): + """Map Gemma's shared tail layers to physical owners and their last readers.""" + 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 = {} + last_reader = {} + 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(set(cache_map.values())) + else: + cache_map[layer_index] = cache_map[last_owner[layer_type]] + owner = last_owner[layer_type] + owners.append(owner) + last_reader[owner] = layer_index + return layer_maps, owners, last_reader diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 6375d8f872..c8a02b1a42 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -6,6 +6,7 @@ from lightllm.common.basemodel.infer_struct import InferStateInfo 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, ) @@ -69,20 +70,13 @@ 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}" - ) + _, kv_owners, last_reader = get_kv_cache_layout(network_config) + kv_owner = kv_owners[layer_num] + self.is_kv_shared_ = kv_owner != layer_num + self.kv_share_target_layer_ = kv_owner if self.is_kv_shared_ else None + # A chunk must not overwrite the history ring until every shared-KV + # consumer has read it. This also keeps graph capture/replay layer-local. + self.commit_sliding_state_ = self.is_sliding and last_reader[kv_owner] == 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 @@ -214,7 +208,7 @@ def _context_attention_kernel( infer_state.b_image_token_end, sliding_window=sw, ) - if not self.is_kv_shared_: + if self.commit_sliding_state_: infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) return o_tensor.view(q.shape) @@ -236,7 +230,7 @@ def _token_attention_kernel( _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 and not self.is_kv_shared_: + if self.commit_sliding_state_: infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) return o_tensor.view(q.shape) diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index 907ad10ed6..cfbc40f110 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -9,6 +9,7 @@ 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 get_kv_cache_layout 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 @@ -88,15 +89,18 @@ def _verify_params(self): 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 _get_sliding_cache_config(self): if hasattr(self, "sliding_cache_config"): return self.sliding_cache_config num_global_kv = self.config.get("num_global_key_value_heads") or self.config["num_key_value_heads"] + layer_maps, _, _ = get_kv_cache_layout(self.config) self.sliding_cache_config = SlidingWindowCacheConfig( - layer_types=self.config["layer_types"], - num_kv_shared_layers=self.config.get("num_kv_shared_layers") or 0, + sliding_layer_to_cache_index=layer_maps["sliding_attention"], + full_layer_to_cache_index=layer_maps["full_attention"], sliding_window=self.config["sliding_window"], sliding_head_num=self.config["num_key_value_heads"] // self.tp_world_size_, sliding_head_dim=self.config["head_dim"], diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 2da7577064..e126ead141 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -398,7 +398,7 @@ def get_can_alloc_token_num(self): ) return self.req_manager.mem_manager.allocator.can_use_mem_size + radix_cache_unref_token_num - def copy_hybrid_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): + def copy_linear_att_state_to_cache_buffer(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_mixed_model or self.radix_cache is None: return @@ -418,11 +418,12 @@ def copy_hybrid_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L big_page_buffer_ids.append(-1) assert len(b_req_idx) == len(big_page_buffer_ids) - self.req_manager.copy_runtime_state_to_cache( - req_indexes=b_req_idx, - buffer_indexes=big_page_buffer_ids, - state_cache_manager=self.radix_cache.linear_att_big_page_buffers, - ) + if any(buffer_id != -1 for buffer_id in big_page_buffer_ids): + self.req_manager.save_big_page_states( + b_req_idx=b_req_idx, + req_indexes=[req.req_idx for req in reqs], + buffer_indexes=big_page_buffer_ids, + ) assert not self.args.disable_chunked_prefill, "chunked prefill must be enabled for hybrid attention models" @@ -438,16 +439,13 @@ def copy_hybrid_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L ) if req.tail_linear_att_small_page_buffer_id is not None: dst_buffer_idx = req.tail_linear_att_small_page_buffer_id - self.req_manager.copy_runtime_state_to_cache( - req_indexes=[req.req_idx], - buffer_indexes=[dst_buffer_idx], - state_cache_manager=self.radix_cache.linear_att_small_page_buffers, + self.req_manager.save_small_page_state( + req_idx=req.req_idx, + buffer_idx=dst_buffer_idx, + small_page_buffers=self.radix_cache.linear_att_small_page_buffers, ) return - # Compatibility for out-of-tree backends while the hybrid name becomes canonical. - copy_linear_att_state_to_cache_buffer = copy_hybrid_att_state_to_cache_buffer - g_infer_context = InferenceContext() diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 69f117a8d0..d1f97c7b22 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -154,12 +154,11 @@ def init_model(self, kvargs): self.is_hybrid_att_mixed_model = isinstance(self.model.req_manager, HybridAttentionReqManager) if self.is_hybrid_att_mixed_model: - self.hybrid_att_cache_manager = self.model.req_manager.create_state_cache_manager( + self.linear_att_cache_manager = self.model.req_manager.create_state_cache_manager( size=self.args.linear_att_cache_size ) else: - self.hybrid_att_cache_manager = None - self.linear_att_cache_manager = self.hybrid_att_cache_manager if self.is_linear_att_mixed_model else None + self.linear_att_cache_manager = None if not self.use_dynamic_prompt_cache: self.radix_cache = None @@ -172,7 +171,7 @@ def init_model(self, kvargs): hash_page_size=self.args.linear_att_hash_page_size, big_page_num=self.args.linear_att_page_block_num, kv_cache_mem_manager=self.model.mem_manager, - linear_att_small_page_buffers=self.hybrid_att_cache_manager, + linear_att_small_page_buffers=self.linear_att_cache_manager, ) else: self.radix_cache = RadixCache( diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 50a0e66e8d..4d09476849 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -119,7 +119,7 @@ def prefill_normal( b_prefill_has_output_cpu=model_input.b_prefill_has_output_cpu, mask_func=self.prefill_mask_func, ) - g_infer_context.copy_hybrid_att_state_to_cache_buffer( + g_infer_context.copy_linear_att_state_to_cache_buffer( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) @@ -215,7 +215,7 @@ def prefill_mtp( target_model_output=model_output, target_next_token_ids=next_token_ids, ) - g_infer_context.copy_hybrid_att_state_to_cache_buffer( + g_infer_context.copy_linear_att_state_to_cache_buffer( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 63cbfe74e8..a7a00e5e83 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -195,7 +195,7 @@ def prefill_normal( b_prefill_has_output_cpu=model_input.b_prefill_has_output_cpu, mask_func=None, ) - g_infer_context.copy_hybrid_att_state_to_cache_buffer( + g_infer_context.copy_linear_att_state_to_cache_buffer( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) @@ -317,7 +317,7 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer ) if g_infer_context.is_hybrid_att_mixed_model: - g_infer_context.copy_hybrid_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() @@ -440,7 +440,7 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] target_next_token_ids=next_token_ids, ) if req_num > 0: - g_infer_context.copy_hybrid_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() @@ -696,7 +696,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I ) if req_num > 0 and g_infer_context.is_hybrid_att_mixed_model: - g_infer_context.copy_hybrid_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) + g_infer_context.copy_linear_att_state_to_cache_buffer(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() diff --git a/test/kernel/test_sliding_window_state.py b/test/kernel/test_sliding_window_state.py new file mode 100644 index 0000000000..a7baa46bec --- /dev/null +++ b/test/kernel/test_sliding_window_state.py @@ -0,0 +1,161 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( + commit_sliding_window_state, + prepare_sliding_window_indexes, +) +from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow +from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig, SlidingWindowStateCacheManager +from lightllm.models.gemma4.kv_layout import get_kv_cache_layout +from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer +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") + + +@pytest.mark.parametrize("history_len", [0, 511, 512, 513, 1024]) +@pytest.mark.parametrize("q_len", [1, 31, 256, 768]) +def test_ring_attention_and_commit_match_token_cache(history_len, q_len): + torch.manual_seed(42) + window, head_dim, req_idx = 512, 64, 1 + seq_len = history_len + q_len + scratch_start = 3 * window + reference = torch.randn((seq_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) + runtime = torch.zeros((scratch_start + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) + old_positions = torch.arange(max(0, history_len - window), history_len, device="cuda") + runtime[req_idx * window + old_positions % window] = reference[old_positions] + runtime[scratch_start:] = reference[history_len:] + mapping = torch.full((3, seq_len), -1, device="cuda", dtype=torch.int32) + b_req = torch.tensor([req_idx], device="cuda", dtype=torch.int32) + b_seq = torch.tensor([seq_len], device="cuda", dtype=torch.int32) + b_q = torch.tensor([q_len], device="cuda", dtype=torch.int32) + b_start = torch.tensor([0], device="cuda", dtype=torch.int32) + b_history = torch.tensor([history_len], device="cuda", dtype=torch.int32) + image_end = torch.zeros(q_len, device="cuda", dtype=torch.int32) + prepare_sliding_window_indexes(mapping, b_req, b_seq, b_q, b_start, window, scratch_start, q_len) + q = torch.randn((q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) + actual, expected = torch.empty_like(q), torch.empty_like(q) + reference_mapping = torch.arange(seq_len, device="cuda", dtype=torch.int32).expand(3, -1) + for buffer, indexes, output in [(runtime, mapping, actual), (reference, reference_mapping, expected)]: + context_attention_fwd_gemma4_mm( + q, + buffer[:, :1], + buffer[:, 1:], + output, + b_req, + b_start, + b_seq, + b_history, + q_len, + indexes, + image_end, + sliding_window=(window - 1, 0), + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + commit_sliding_window_state(runtime, b_req, b_seq, b_q, b_start, window, scratch_start, q_len) + positions = torch.arange(max(0, seq_len - window), seq_len, device="cuda") + torch.testing.assert_close(runtime[req_idx * window + positions % window], reference[positions], atol=0, rtol=0) + assert torch.count_nonzero(runtime[:window]).item() == 0 + + +def test_shared_kv_is_committed_only_after_last_reader_and_snapshot_is_independent(): + window, history_len, q_len, head_dim = 512, 512, 256, 64 + layout, owners, last_reader = get_kv_cache_layout( + {"layer_types": ["sliding_attention", "full_attention", "sliding_attention"], "num_kv_shared_layers": 1} + ) + config = SlidingWindowCacheConfig( + layout["sliding_attention"], layout["full_attention"], window, 1, head_dim, 1, 64, torch.bfloat16 + ) + manager = object.__new__(ReqManagerForSlidingWindow) + manager.sliding_config = config + manager.sliding_window = window + manager.scratch_token_num = q_len + manager.scratch_start = 2 * window + manager.req_to_sliding_window = torch.zeros( + (1, 2 * window + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16 + ) + manager.req_to_sliding_window[0, manager.scratch_start :, 1] = 1 + manager.req_to_sliding_window_indexs = torch.zeros((2, history_len + q_len), device="cuda", dtype=torch.int32) + int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + state = SimpleNamespace( + input_ids=int_tensor([0] * q_len), + b_req_idx=int_tensor([0]), + b_seq_len=int_tensor([history_len + q_len]), + b_q_seq_len=int_tensor([q_len]), + b_q_start_loc=int_tensor([0]), + b_ready_cache_len=int_tensor([history_len]), + max_q_seq_len=q_len, + b_image_token_end=int_tensor([0] * q_len), + req_manager=manager, + ) + manager.prepare_sliding_window(state) + q = torch.zeros((q_len, 1, head_dim), device="cuda", dtype=torch.bfloat16) + outputs = [] + for index in [0, 2]: + layer = object.__new__(Gemma4TransformerLayerInfer) + layer.layer_num_, layer.is_sliding, layer.sliding_window_ = index, True, window + layer.is_kv_shared_, layer.kv_share_target_layer_ = index != owners[index], owners[index] + layer.commit_sliding_state_ = last_reader[owners[index]] == index + layer.tp_q_head_num_, layer.head_dim_ = 1, head_dim + layer.alloc_tensor = lambda shape, dtype: torch.empty(shape, dtype=dtype, device="cuda") + outputs.append(layer._context_attention_kernel(q, None, state, None)) + if index == 0: + assert torch.count_nonzero(manager.req_to_sliding_window[:, :window]).item() == 0 + torch.testing.assert_close(outputs[0], outputs[1], atol=0, rtol=0) + assert outputs[1][0, 0, 0].item() == 1 / window + assert manager.req_to_sliding_window[0, 0, 1, 0].item() == 1 + + pages = SlidingWindowStateCacheManager(2, config) + page = pages.alloc_one_state_cache() + manager.save_small_page_state(0, page, pages) + saved = pages.get_state_cache(page).clone() + manager.req_to_sliding_window[:, :window].fill_(7) + torch.testing.assert_close(pages.get_state_cache(page), saved, atol=0, rtol=0) + manager._restore_state(1, pages, page) + torch.testing.assert_close(manager.req_to_sliding_window[:, window : 2 * window], saved, atol=0, rtol=0) + pages.free_state_cache([page]) + assert pages.get_free_cache_num() == 2 + + +def test_empty_snapshot_does_not_read_gpu_request_ids(): + manager = object.__new__(ReqManagerForSlidingWindow) + # No runtime or page pool is needed for a no-op. In particular, no .tolist() + # or other GPU operation should be performed on b_req_idx. + manager.save_big_page_states(object(), [0, 1], [-1, -1]) + + +def test_batched_window_commit_with_hold_request_and_cuda_graph_replay(): + window, scratch_start, head_dim = 32, 4 * 32, 64 + req_ids, lengths, q_lengths = [2, 0, 3], [86, 5, 100], [6, 5, 32] + starts = [0, 6, 11] + int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + b_req, b_seq, b_q, b_start = map(int_tensor, [req_ids, lengths, q_lengths, starts]) + mapping = torch.full((4, 128), -1, device="cuda", dtype=torch.int32) + runtime = torch.zeros((scratch_start + sum(q_lengths), 2, head_dim), device="cuda", dtype=torch.bfloat16) + runtime[scratch_start:] = torch.randn_like(runtime[scratch_start:]) + + def forward(): + prepare_sliding_window_indexes(mapping, b_req, b_seq, b_q, b_start, window, scratch_start, max(q_lengths)) + commit_sliding_window_state(runtime, b_req, b_seq, b_q, b_start, window, scratch_start, max(q_lengths)) + + forward() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + forward() + # Replay with changed GPU state, including a padded/hold request ID. + runtime[:scratch_start].zero_() + runtime[scratch_start:].mul_(2) + graph.replay() + for req, seq, q_len, start in zip(req_ids, lengths, q_lengths, starts): + positions = torch.arange(seq - q_len, seq, device="cuda") + expected = runtime[scratch_start + start : scratch_start + start + q_len] + torch.testing.assert_close(runtime[req * window + positions % window], expected, atol=0, rtol=0) + torch.testing.assert_close( + mapping[req, positions], + torch.arange(scratch_start + start, scratch_start + start + q_len, device="cuda", dtype=torch.int32), + ) + assert torch.count_nonzero(runtime[window : 2 * window]).item() == 0 diff --git a/test/utils/test_sliding_window_cache.py b/test/utils/test_sliding_window_cache.py new file mode 100644 index 0000000000..b7dc1472db --- /dev/null +++ b/test/utils/test_sliding_window_cache.py @@ -0,0 +1,130 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager +from lightllm.common.req_manager.linear_att import ReqManagerForMamba +from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig +from lightllm.models.gemma4.kv_layout import get_kv_cache_layout + + +@pytest.mark.parametrize("layer_num,shared,sliding_num,full_num", [(42, 18, 20, 4), (60, 0, 50, 10)]) +def test_gemma_physical_owners_and_last_readers(layer_num, shared, sliding_num, full_num): + layer_types = ["sliding_attention"] * 5 + ["full_attention"] + maps, owners, last_readers = get_kv_cache_layout( + {"layer_types": layer_types * (layer_num // 6), "num_kv_shared_layers": shared} + ) + assert len(set(maps["sliding_attention"].values())) == sliding_num + assert len(set(maps["full_attention"].values())) == full_num + for index, owner in enumerate(owners): + assert owner <= index <= last_readers[owner] + if shared: + assert owners[40] == 22 and last_readers[22] == 40 + assert owners[41] == 23 and last_readers[23] == 41 + + +def _memory_manager(big_page_tokens=2048, small_pages=8, enabled=True): + manager = object.__new__(HybridSlidingMemoryManager) + manager.head_num, manager.head_dim, manager.layer_num, manager.dtype = 1, 512, 10, torch.bfloat16 + manager.sliding_config = SlidingWindowCacheConfig( + {i: i for i in range(50)}, {50 + i: i for i in range(10)}, 1024, 4, 256, 1, 512, torch.bfloat16 + ) + manager.big_page_token_num, manager.small_page_num, manager.enable_prompt_cache = ( + big_page_tokens, + small_pages, + enabled, + ) + return manager + + +@pytest.mark.parametrize("token_num", [1, 2047, 2048, 2049, 8192]) +def test_profile_accounts_for_small_big_partial_page_and_hold_token(token_num): + manager = _memory_manager() + assert manager.sliding_config.get_state_nbytes() == 200 * 1024 ** 2 + expected = (token_num + 1) * 20480 + (8 + (token_num + 2047) // 2048) * 200 * 1024 ** 2 + assert manager._cache_nbytes(token_num) == expected + assert manager._profile_token_num(expected) == token_num + if token_num > 1: + assert manager._profile_token_num(expected - 1) < token_num + + +def test_profile_reports_impossible_checkpoint_budget(): + manager = _memory_manager(small_pages=512) + with pytest.raises(ValueError, match="linear_att_cache_size"): + manager._profile_token_num(80 * 1024 ** 3) + + +def test_disabled_prompt_cache_does_not_reserve_pages(): + manager = _memory_manager(small_pages=0, enabled=False) + assert manager._cache_nbytes(4096) == 4097 * manager.get_cell_size() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_profiled_gpu_pools_match_reserved_bytes_and_are_reused(monkeypatch): + import lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager as memory_module + from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow + + manager = _memory_manager(big_page_tokens=32, small_pages=2) + manager.head_num, manager.head_dim, manager.layer_num = 1, 64, 1 + manager.sliding_config = SlidingWindowCacheConfig({0: 0}, {1: 0}, 32, 1, 64, 1, 64, torch.bfloat16) + manager.size = None + budget = manager._cache_nbytes(65) + monkeypatch.setattr(memory_module.dist, "get_world_size", lambda: 1) + monkeypatch.setattr(memory_module, "get_available_gpu_memory", lambda world_size: budget / 1024 ** 3) + monkeypatch.setattr(memory_module, "get_total_gpu_memory", lambda: 1) + manager.profile_size(1.0) + assert manager.size == 65 + manager._init_buffers(manager.size, manager.dtype, manager.head_num, manager.head_dim, manager.layer_num) + allocated = sum( + t.numel() * t.element_size() + for t in [ + manager.kv_buffer, + manager.linear_att_big_page_buffers.state_cache, + manager.sliding_small_page_buffers.state_cache, + ] + ) + assert allocated == budget + req_manager = object.__new__(ReqManagerForSlidingWindow) + req_manager.mem_manager = manager + assert req_manager.create_state_cache_manager(2) is manager.sliding_small_page_buffers + manager.size = 1000000 + with pytest.raises(ValueError, match="exceed available GPU memory"): + manager.profile_size(1.0) + + +@pytest.mark.parametrize("mtp_step", [0, 2]) +def test_linear_small_page_preserves_main_copy_and_mtp_crop(mtp_step): + manager = object.__new__(ReqManagerForMamba) + manager.mtp_step = mtp_step + manager.linear_config = SimpleNamespace(get_conv_state_shape=lambda: (3, 4)) + conv = torch.arange(2 * 3 * 3 * (4 + mtp_step)).reshape(2, 3, 3, 4 + mtp_step) + ssm = torch.arange(2 * 3 * (mtp_step + 1) * 5).reshape(2, 3 * (mtp_step + 1), 5) + manager.req_to_conv_state, manager.req_to_ssm_state = SimpleNamespace(buffer=conv), SimpleNamespace(buffer=ssm) + dst_conv, dst_ssm = torch.empty((2, 3, 4), dtype=conv.dtype), torch.empty((2, 5), dtype=ssm.dtype) + pages = SimpleNamespace(get_state_cache=lambda buffer_idx: (dst_conv, dst_ssm)) + manager.save_small_page_state(1, 0, pages) + torch.testing.assert_close(dst_conv, conv[:, 1, :, :4]) + torch.testing.assert_close(dst_ssm, ssm[:, mtp_step + 1]) + + +@pytest.mark.parametrize("unsupported_mode", ["enable_dp_prompt_cache_fetch", "diverse_mode"]) +def test_unsupported_sliding_state_transfer_modes_fail_before_loading_weights(monkeypatch, unsupported_mode): + import lightllm.models.gemma4.model as gemma_model + + model = object.__new__(gemma_model.Gemma4TpPartModel) + model.load_way, model.tp_world_size_ = "HF", 2 + model.config = {"num_attention_heads": 8, "num_key_value_heads": 2, "num_hidden_layers": 42} + args = SimpleNamespace( + mtp_step=0, + enable_cpu_cache=False, + disable_chunked_prefill=False, + run_mode="normal", + llm_kv_type="None", + enable_dp_prompt_cache_fetch=False, + diverse_mode=False, + ) + setattr(args, unsupported_mode, True) + monkeypatch.setattr(gemma_model, "get_env_start_args", lambda: args) + with pytest.raises(AssertionError, match="does not support"): + model._verify_params() diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index 96f295efd2..e5ccaa3a35 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -131,11 +131,13 @@ def start_offload(req, cpu_kv_cache_stream): assert len(module.cpu_cache_handle_queue) == 2 -def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): +@pytest.mark.parametrize("is_linear", [True, False]) +def test_non_gpu_hybrid_cache_tiers_release_pending_state_pages(is_linear): freed_small_pages = [] freed_big_pages = [] context = InferenceContext() - context.is_linear_att_mixed_model = True + context.is_linear_att_mixed_model = is_linear + context.is_hybrid_att_mixed_model = True context.req_manager = SimpleNamespace(req_to_token_indexs=torch.tensor([[10, 11, 12]])) context.radix_cache = SimpleNamespace( linear_att_small_page_buffers=SimpleNamespace(free_state_cache=freed_small_pages.extend), @@ -157,3 +159,15 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): assert freed_big_pages == [8, 9] assert req.tail_linear_att_small_page_buffer_id is None assert req.linear_att_len_to_big_page_id == {} + + +def test_hybrid_snapshot_outside_page_boundaries_does_not_touch_runtime(): + context = InferenceContext() + context.is_hybrid_att_mixed_model = True + context.args = SimpleNamespace( + linear_att_hash_page_size=32, linear_att_page_block_num=8, disable_chunked_prefill=False + ) + context.radix_cache = SimpleNamespace() + context.req_manager = None # Any attempted snapshot would fail. + reqs = [SimpleNamespace(req_idx=0, get_chuncked_input_token_len=lambda: 17, linear_att_cache_len=32)] + context.copy_linear_att_state_to_cache_buffer(b_req_idx=[0], reqs=reqs) From 8c25c4de0ce025e785b90ad1ae3df8a897a0056b Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Mon, 7 Sep 2026 06:21:00 +0000 Subject: [PATCH 05/14] fix: bound sliding window commits to retained tokens --- .../source/cookbook/gemma4_hybrid_cache.rst | 62 ++++++++++++++ docs/CN/source/index.rst | 1 + .../triton_kernel/sliding_window_state.py | 8 +- lightllm/models/gemma4/model.py | 6 ++ test/kernel/test_sliding_window_state.py | 44 ++++++++-- test/utils/test_sliding_window_cache.py | 36 +++++++++ .../model_infer/test_hybrid_state_cache.py | 80 +++++++++++++++++++ 7 files changed, 226 insertions(+), 11 deletions(-) create mode 100644 docs/CN/source/cookbook/gemma4_hybrid_cache.rst create mode 100644 unit_tests/server/router/model_infer/test_hybrid_state_cache.py diff --git a/docs/CN/source/cookbook/gemma4_hybrid_cache.rst b/docs/CN/source/cookbook/gemma4_hybrid_cache.rst new file mode 100644 index 0000000000..66cc8e3929 --- /dev/null +++ b/docs/CN/source/cookbook/gemma4_hybrid_cache.rst @@ -0,0 +1,62 @@ +Gemma4 hybrid 缓存参数 +===================== + +Gemma4 的 full-attention 层保留 token 粒度的 KV，sliding-attention 层使用 +GPU 请求窗口及临时 chunk 区域。大小页继续沿用现有 linear hybrid 的虚拟 +token 边界和匹配规则；窗口运行态与缓存快照分别存储。 + +显存容量 +-------- + +当前 sliding 的小页快照在 GPU 上，不同于 linear 小页的 CPU 状态池。 +一个快照包含所有物理 sliding KV 层的窗口，其每卡大小为： + +.. code-block:: text + + 物理 sliding 层数 × window × 2 × 每卡 KV heads × head_dim × dtype 字节数 + +共享 KV 的逻辑层不重复计费。例如 Gemma4-31B、TP4、BF16 下，一个窗口快照为： + +.. code-block:: text + + 50 × 1024 × 2 × 4 × 256 × 2 = 200 MiB / GPU + +需要分别考虑以下容量： + +- ``--running_max_req_size`` 决定请求窗口数量，另有一个 hold request 窗口。 +- ``--batch_max_tokens`` 等参数决定临时 chunk 区域大小，不能忽略长 chunk 的显存。 +- ``--linear_att_cache_size`` 是小页状态槽位数，不是 token 数。 +- 大页状态池、full KV、模型权重及计算工作区也占用显存。 + +在 DP1 下，未指定 ``linear_att_cache_size`` 时，它默认是 +``running_max_req_size`` 的两倍。上述 31B 配置若使用 256 个请求槽位， +仅请求窗口和 512 个小页就约需 150 GiB，尚未包含权重和临时 chunk 区域。 +因此应显式按显存预算设置这两个参数；不要直接套用 full-attention 模型的并发槽位配置。 +内存管理器会计入大小页池并检查容量，不会静默缩减用户指定的小页数量。 + +切分与性能比较 +-------------- + +- ``--linear_att_hash_page_size`` 是当前 hybrid 树的 hash 分块粒度，默认 512。 +- 大页覆盖的 token 数是 + ``linear_att_hash_page_size * linear_att_page_block_num``。 + 例如 hash 为 512、block num 为 32 时，大页为 16384 token。 +- ``--chunked_prefill_size`` 是单轮上限，不是每轮固定长度。 + 当前流程还会在大页边界和请求尾部 checkpoint 处截断。 +- 请求尾部 checkpoint 为 ``floor((prompt_len - 1) / hash_page_size) * hash_page_size``。 + 不会在每一个 hash 分块末尾都保存请求快照；尾部 checkpoint 若落在 chunk 内， + 可能多产生一轮 prefill。 + +性能比较应报告 chunk、大小页参数、缓存实际命中长度和 CUDA Graph 配置。 +256-token chunk、32-token hash 等边界压力配置不能作为常规部署速度的唯一基线。 +大页 16384 与 chunk 4096/8192 在从零开始时对齐，但尾部 checkpoint 仍可能额外截断。 + +当前能力边界 +------------ + +本实现尚不支持 CPU cache、PD、MTP、量化 KV、DP prompt-cache fetch 和 diverse mode， +并要求启用 chunked prefill。带跨层 KV sharing 的配置不支持 microbatch overlap： +临时 KV 区域尚未按微批隔离。这里的 microbatch overlap 不包括普通的 CPU/GPU 调度重叠。 + +任意长度的 request-level 小页、输入与输出双 checkpoint，以及纯 sliding 模型的 +提前淘汰策略不属于本次 hybrid 接入的范围。 diff --git a/docs/CN/source/index.rst b/docs/CN/source/index.rst index 8f79e5126f..c35e15217e 100755 --- a/docs/CN/source/index.rst +++ b/docs/CN/source/index.rst @@ -65,6 +65,7 @@ Lightllm 整合了众多的开源方案的优点，包括但不限于 FasterTran GLM-4.7-Flash 部署 Qwen3.5 部署 + Gemma4 hybrid 缓存参数 .. toctree:: :maxdepth: 1 diff --git a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py index 37a7693c96..bd940950d7 100644 --- a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py @@ -102,14 +102,15 @@ def _commit_sliding_window_state( BLOCK_D: tl.constexpr, ): batch_idx = tl.program_id(0) - q_offset = tl.program_id(1) + tail_offset = tl.program_id(1) head_idx = tl.program_id(2) req_idx = tl.load(BReqIdx + batch_idx) seq_len = tl.load(BSeqLen + batch_idx) q_len = tl.load(BQSeqLen + batch_idx) q_start = tl.load(BQStartLoc + batch_idx) + q_offset = tl.maximum(q_len - WINDOW, 0) + tail_offset pos = seq_len - q_len + q_offset - mask_token = (q_offset < q_len) & (pos >= seq_len - WINDOW) + mask_token = q_offset < q_len src_token = SCRATCH_START + q_start + q_offset dst_token = req_idx * WINDOW + pos % WINDOW dims = tl.arange(0, BLOCK_D) @@ -132,7 +133,8 @@ def commit_sliding_window_state( max_q_seq_len, ): block_d = triton.next_power_of_2(layer_buffer.shape[-1]) - grid = (b_req_idx.shape[0], max_q_seq_len, layer_buffer.shape[1]) + # Earlier chunk tokens are only used by attention, never by the next step. + grid = (b_req_idx.shape[0], min(max_q_seq_len, sliding_window), layer_buffer.shape[1]) _commit_sliding_window_state[grid]( layer_buffer, b_req_idx, diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index cfbc40f110..be3885cc13 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -84,6 +84,12 @@ def _verify_params(self): f"num_kv_shared_layers={kv_shared} out of range for " f"num_hidden_layers={self.config['num_hidden_layers']}" ) + if kv_shared: + # Shared layers retain the owner's scratch KV across layers. Two + # interleaved microbatches would overwrite the same scratch slots. + assert not ( + args.enable_prefill_microbatch_overlap or args.enable_decode_microbatch_overlap + ), "Gemma-4 shared sliding-window KV does not support microbatch overlap yet" assert args.mtp_step == 0, "Gemma-4 hybrid sliding-window cache does not support MTP yet" assert not args.enable_cpu_cache, "Gemma-4 hybrid sliding-window cache does not support CPU cache" assert not args.disable_chunked_prefill, "Gemma-4 hybrid sliding-window cache requires chunked prefill" diff --git a/test/kernel/test_sliding_window_state.py b/test/kernel/test_sliding_window_state.py index a7baa46bec..1f8fdc70b6 100644 --- a/test/kernel/test_sliding_window_state.py +++ b/test/kernel/test_sliding_window_state.py @@ -127,13 +127,19 @@ def test_empty_snapshot_does_not_read_gpu_request_ids(): manager.save_big_page_states(object(), [0, 1], [-1, -1]) -def test_batched_window_commit_with_hold_request_and_cuda_graph_replay(): - window, scratch_start, head_dim = 32, 4 * 32, 64 - req_ids, lengths, q_lengths = [2, 0, 3], [86, 5, 100], [6, 5, 32] - starts = [0, 6, 11] +@pytest.mark.parametrize( + "window,q_lengths", + [(32, [6, 5, 32]), (512, [4096, 1, 513]), (512, [1, 8192, 511]), (1024, [8192, 4096, 1])], +) +def test_batched_window_commit_with_hold_request_and_cuda_graph_replay(window, q_lengths): + scratch_start, head_dim = 4 * window, 64 + req_ids = [2, 0, 3] + lengths = [q_lengths[0] + 2 * window + 3, q_lengths[1], q_lengths[2] + window - 1] + replay_lengths = [length + window + 7 for length in lengths] + starts = [0, q_lengths[0], q_lengths[0] + q_lengths[1]] int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) b_req, b_seq, b_q, b_start = map(int_tensor, [req_ids, lengths, q_lengths, starts]) - mapping = torch.full((4, 128), -1, device="cuda", dtype=torch.int32) + mapping = torch.full((4, max(replay_lengths)), -1, device="cuda", dtype=torch.int32) runtime = torch.zeros((scratch_start + sum(q_lengths), 2, head_dim), device="cuda", dtype=torch.bfloat16) runtime[scratch_start:] = torch.randn_like(runtime[scratch_start:]) @@ -149,13 +155,35 @@ def forward(): # Replay with changed GPU state, including a padded/hold request ID. runtime[:scratch_start].zero_() runtime[scratch_start:].mul_(2) + b_seq.copy_(int_tensor(replay_lengths)) graph.replay() - for req, seq, q_len, start in zip(req_ids, lengths, q_lengths, starts): + for req, seq, q_len, start in zip(req_ids, replay_lengths, q_lengths, starts): + tail_start = max(q_len - window, 0) + positions = torch.arange(seq - q_len + tail_start, seq, device="cuda") + expected_ring = torch.zeros_like(runtime[req * window : (req + 1) * window]) + expected_ring[positions % window] = runtime[scratch_start + start + tail_start : scratch_start + start + q_len] + torch.testing.assert_close(runtime[req * window : (req + 1) * window], expected_ring, atol=0, rtol=0) positions = torch.arange(seq - q_len, seq, device="cuda") - expected = runtime[scratch_start + start : scratch_start + start + q_len] - torch.testing.assert_close(runtime[req * window + positions % window], expected, atol=0, rtol=0) torch.testing.assert_close( mapping[req, positions], torch.arange(scratch_start + start, scratch_start + start + q_len, device="cuda", dtype=torch.int32), ) assert torch.count_nonzero(runtime[window : 2 * window]).item() == 0 + + +@pytest.mark.parametrize("max_q_seq_len", [1, 511, 512, 513, 4096, 8192]) +def test_commit_grid_is_bounded_by_window(monkeypatch, max_q_seq_len): + import lightllm.common.basemodel.triton_kernel.sliding_window_state as state_kernel + + grids = [] + + class RecordingKernel: + def __getitem__(self, grid): + grids.append(grid) + return lambda *args, **kwargs: None + + monkeypatch.setattr(state_kernel, "_commit_sliding_window_state", RecordingKernel()) + layer_buffer = SimpleNamespace(shape=(8192, 2, 64), stride=lambda: (128, 64, 1)) + req_ids = SimpleNamespace(shape=(3,)) + state_kernel.commit_sliding_window_state(layer_buffer, req_ids, None, None, None, 512, 2048, max_q_seq_len) + assert grids == [(3, min(max_q_seq_len, 512), 2)] diff --git a/test/utils/test_sliding_window_cache.py b/test/utils/test_sliding_window_cache.py index b7dc1472db..3cab1fd56f 100644 --- a/test/utils/test_sliding_window_cache.py +++ b/test/utils/test_sliding_window_cache.py @@ -128,3 +128,39 @@ def test_unsupported_sliding_state_transfer_modes_fail_before_loading_weights(mo monkeypatch.setattr(gemma_model, "get_env_start_args", lambda: args) with pytest.raises(AssertionError, match="does not support"): model._verify_params() + + +@pytest.mark.parametrize("shared_layers", [0, 18]) +@pytest.mark.parametrize( + "overlap_mode", [None, "enable_prefill_microbatch_overlap", "enable_decode_microbatch_overlap"] +) +def test_shared_kv_rejects_interleaved_microbatches(monkeypatch, shared_layers, overlap_mode): + import lightllm.models.gemma4.model as gemma_model + + model = object.__new__(gemma_model.Gemma4TpPartModel) + model.load_way, model.tp_world_size_ = "HF", 2 + model.config = { + "num_attention_heads": 8, + "num_key_value_heads": 2, + "num_hidden_layers": 42, + "num_kv_shared_layers": shared_layers, + } + args = SimpleNamespace( + mtp_step=0, + enable_cpu_cache=False, + disable_chunked_prefill=False, + run_mode="normal", + llm_kv_type="None", + enable_dp_prompt_cache_fetch=False, + diverse_mode=False, + enable_prefill_microbatch_overlap=False, + enable_decode_microbatch_overlap=False, + ) + if overlap_mode is not None: + setattr(args, overlap_mode, True) + monkeypatch.setattr(gemma_model, "get_env_start_args", lambda: args) + if shared_layers and overlap_mode is not None: + with pytest.raises(AssertionError, match="shared sliding-window KV does not support microbatch overlap"): + model._verify_params() + else: + model._verify_params() diff --git a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py new file mode 100644 index 0000000000..03fc6b3e25 --- /dev/null +++ b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py @@ -0,0 +1,80 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow +from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager +from lightllm.server.router.model_infer.infer_batch import InferenceContext + + +@pytest.mark.parametrize("is_hybrid,radix_cache", [(False, object()), (True, None)]) +def test_snapshot_without_hybrid_cache_returns_before_reading_requests(is_hybrid, radix_cache): + context = InferenceContext(is_hybrid_att_mixed_model=is_hybrid, radix_cache=radix_cache) + + # Neither argument supports iteration or len: an early return must not inspect them. + context.copy_linear_att_state_to_cache_buffer(b_req_idx=object(), reqs=object()) + + +@pytest.mark.parametrize("chunk_end,cache_len", [(17, 32), (768, 544)]) +def test_snapshot_outside_cacheable_boundaries_does_not_allocate_or_copy(chunk_end, cache_len): + context = InferenceContext(is_hybrid_att_mixed_model=True, radix_cache=object()) + context.args = SimpleNamespace( + linear_att_hash_page_size=32, linear_att_page_block_num=8, disable_chunked_prefill=False + ) + req = SimpleNamespace( + req_idx=0, + get_chuncked_input_token_len=lambda: chunk_end, + linear_att_cache_len=cache_len, + linear_att_len_to_big_page_id={}, + tail_linear_att_small_page_buffer_id=None, + ) + + # The radix object has no allocator and req_manager is None, so either access fails. + context.copy_linear_att_state_to_cache_buffer(b_req_idx=[0], reqs=[req]) + + assert req.linear_att_len_to_big_page_id == {} + assert req.tail_linear_att_small_page_buffer_id is None + + +def _cpu_state_cache(size): + pages = object.__new__(SlidingWindowStateCacheManager) + pages.size = size + pages.state_cache = torch.empty((size, 2, 4, 2, 4), dtype=torch.float32) + pages.clear_to_init_state() + return pages + + +def test_sliding_big_snapshot_skips_invalid_requests_and_copies_only_selected_page(): + pages = _cpu_state_cache(3) + manager = object.__new__(ReqManagerForSlidingWindow) + manager.sliding_window = 4 + manager.mem_manager = SimpleNamespace(linear_att_big_page_buffers=pages) + manager.req_to_sliding_window = torch.arange(2 * 12 * 2 * 4, dtype=torch.float32).reshape(2, 12, 2, 4) + expected = manager.req_to_sliding_window[:, 4:8].clone() + + # Skipped request IDs are deliberately out of range; GPU request IDs must not be read. + manager.save_big_page_states(b_req_idx=object(), req_indexes=[999, 1, 888], buffer_indexes=[-1, 2, -1]) + + torch.testing.assert_close(pages.get_state_cache(2), expected, atol=0, rtol=0) + assert torch.count_nonzero(pages.state_cache[:2]).item() == 0 + manager.req_to_sliding_window.fill_(-1) + torch.testing.assert_close(pages.get_state_cache(2), expected, atol=0, rtol=0) + + +def test_sliding_state_pool_exhaustion_and_released_slot_reuse(): + pages = _cpu_state_cache(2) + assert pages.alloc_state_cache(3) is None + assert pages.get_free_cache_num() == 2 + assert pages.alloc_state_cache(2) == [0, 1] + assert pages.get_used_cache_num() == 2 + assert pages.alloc_one_state_cache() is None + + pages.free_state_cache([1]) + assert pages.get_free_cache_num() == 1 + assert pages.alloc_one_state_cache() == 1 + assert pages.alloc_one_state_cache() is None + + pages.free_state_cache([0, 1]) + assert pages.get_free_cache_num() == 2 + assert pages.get_used_cache_num() == 0 From fd7c6535cd04b58ddbdbbab24c5319f815893003 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Mon, 7 Sep 2026 07:16:29 +0000 Subject: [PATCH 06/14] docs: defer Gemma hybrid cache guide until page redesign --- .../source/cookbook/gemma4_hybrid_cache.rst | 62 ------------------- docs/CN/source/index.rst | 1 - 2 files changed, 63 deletions(-) delete mode 100644 docs/CN/source/cookbook/gemma4_hybrid_cache.rst diff --git a/docs/CN/source/cookbook/gemma4_hybrid_cache.rst b/docs/CN/source/cookbook/gemma4_hybrid_cache.rst deleted file mode 100644 index 66cc8e3929..0000000000 --- a/docs/CN/source/cookbook/gemma4_hybrid_cache.rst +++ /dev/null @@ -1,62 +0,0 @@ -Gemma4 hybrid 缓存参数 -===================== - -Gemma4 的 full-attention 层保留 token 粒度的 KV，sliding-attention 层使用 -GPU 请求窗口及临时 chunk 区域。大小页继续沿用现有 linear hybrid 的虚拟 -token 边界和匹配规则；窗口运行态与缓存快照分别存储。 - -显存容量 --------- - -当前 sliding 的小页快照在 GPU 上，不同于 linear 小页的 CPU 状态池。 -一个快照包含所有物理 sliding KV 层的窗口，其每卡大小为： - -.. code-block:: text - - 物理 sliding 层数 × window × 2 × 每卡 KV heads × head_dim × dtype 字节数 - -共享 KV 的逻辑层不重复计费。例如 Gemma4-31B、TP4、BF16 下，一个窗口快照为： - -.. code-block:: text - - 50 × 1024 × 2 × 4 × 256 × 2 = 200 MiB / GPU - -需要分别考虑以下容量： - -- ``--running_max_req_size`` 决定请求窗口数量，另有一个 hold request 窗口。 -- ``--batch_max_tokens`` 等参数决定临时 chunk 区域大小，不能忽略长 chunk 的显存。 -- ``--linear_att_cache_size`` 是小页状态槽位数，不是 token 数。 -- 大页状态池、full KV、模型权重及计算工作区也占用显存。 - -在 DP1 下，未指定 ``linear_att_cache_size`` 时，它默认是 -``running_max_req_size`` 的两倍。上述 31B 配置若使用 256 个请求槽位， -仅请求窗口和 512 个小页就约需 150 GiB，尚未包含权重和临时 chunk 区域。 -因此应显式按显存预算设置这两个参数；不要直接套用 full-attention 模型的并发槽位配置。 -内存管理器会计入大小页池并检查容量，不会静默缩减用户指定的小页数量。 - -切分与性能比较 --------------- - -- ``--linear_att_hash_page_size`` 是当前 hybrid 树的 hash 分块粒度，默认 512。 -- 大页覆盖的 token 数是 - ``linear_att_hash_page_size * linear_att_page_block_num``。 - 例如 hash 为 512、block num 为 32 时，大页为 16384 token。 -- ``--chunked_prefill_size`` 是单轮上限，不是每轮固定长度。 - 当前流程还会在大页边界和请求尾部 checkpoint 处截断。 -- 请求尾部 checkpoint 为 ``floor((prompt_len - 1) / hash_page_size) * hash_page_size``。 - 不会在每一个 hash 分块末尾都保存请求快照；尾部 checkpoint 若落在 chunk 内， - 可能多产生一轮 prefill。 - -性能比较应报告 chunk、大小页参数、缓存实际命中长度和 CUDA Graph 配置。 -256-token chunk、32-token hash 等边界压力配置不能作为常规部署速度的唯一基线。 -大页 16384 与 chunk 4096/8192 在从零开始时对齐，但尾部 checkpoint 仍可能额外截断。 - -当前能力边界 ------------- - -本实现尚不支持 CPU cache、PD、MTP、量化 KV、DP prompt-cache fetch 和 diverse mode， -并要求启用 chunked prefill。带跨层 KV sharing 的配置不支持 microbatch overlap： -临时 KV 区域尚未按微批隔离。这里的 microbatch overlap 不包括普通的 CPU/GPU 调度重叠。 - -任意长度的 request-level 小页、输入与输出双 checkpoint，以及纯 sliding 模型的 -提前淘汰策略不属于本次 hybrid 接入的范围。 diff --git a/docs/CN/source/index.rst b/docs/CN/source/index.rst index c35e15217e..8f79e5126f 100755 --- a/docs/CN/source/index.rst +++ b/docs/CN/source/index.rst @@ -65,7 +65,6 @@ Lightllm 整合了众多的开源方案的优点，包括但不限于 FasterTran GLM-4.7-Flash 部署 Qwen3.5 部署 - Gemma4 hybrid 缓存参数 .. toctree:: :maxdepth: 1 From f32a59728ac01686232b316ce9f348c966e8f39c Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Mon, 7 Sep 2026 07:47:34 +0000 Subject: [PATCH 07/14] feat: support CPU cache for hybrid sliding windows --- .../sliding_window_cpu_cache_copy.py | 209 ++++++++++++++ .../hybrid_sliding_mem_manager.py | 15 +- .../operator/hybrid_sliding.py | 110 +++++++- .../sliding_window_cache_manager/config.py | 22 ++ .../state_cache.py | 12 +- lightllm/models/gemma4/kv_layout.py | 23 ++ lightllm/models/gemma4/model.py | 19 +- lightllm/server/api_start.py | 8 +- lightllm/server/core/objs/req.py | 4 +- .../mode_backend/multi_level_kv_cache.py | 5 +- lightllm/utils/kv_cache_utils.py | 41 ++- ...test_sliding_window_cpu_cache_attention.py | 106 +++++++ .../test_sliding_window_cpu_cache_copy.py | 170 ++++++++++++ test/utils/test_sliding_cpu_cache_meta.py | 155 +++++++++++ test/utils/test_sliding_window_cache.py | 37 ++- .../mode_backend/test_multi_level_kv_cache.py | 2 +- .../model_infer/test_hybrid_state_cache.py | 17 +- .../test_sliding_cpu_cache_loading.py | 125 +++++++++ .../test_sliding_cpu_cache_operator.py | 259 ++++++++++++++++++ 19 files changed, 1294 insertions(+), 45 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/sliding_window_cpu_cache_copy.py create mode 100644 test/kernel/test_sliding_window_cpu_cache_attention.py create mode 100644 test/kernel/test_sliding_window_cpu_cache_copy.py create mode 100644 test/utils/test_sliding_cpu_cache_meta.py create mode 100644 unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py create mode 100644 unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py 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..0769efdcd8 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_cpu_cache_copy.py @@ -0,0 +1,209 @@ +import torch +import triton +import triton.language as tl + +from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig + + +@triton.jit +def _copy_sliding_window_cpu_cache( + mem_indexes, + page_indexes, + page_readies, + big_page_buffer_ids, + gpu_full_att_kv_state, + gpu_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 + gpu_ptr = gpu_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(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) + + +def _copy_state_cache( + mem_indexes, + page_indexes, + page_readies, + big_page_buffer_ids, + gpu_full_att_kv_state, + gpu_sliding_state, + cpu_cache_tensor, + tp_rank, + tp_world_size, + big_page_token_num, + sliding_config, + 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 + assert 0 <= tp_rank < tp_world_size + assert big_page_token_num > 0 and grid_num > 0 + if page_num == 0: + return + + assert gpu_full_att_kv_state.shape[0] == sliding_config.full_layer_num + assert gpu_full_att_kv_state.shape[2:] == ( + 2 * sliding_config.full_head_num, + sliding_config.full_head_dim, + ) + assert gpu_sliding_state.shape[1:] == sliding_config.get_state_shape() + assert gpu_full_att_kv_state.dtype == gpu_sliding_state.dtype == sliding_config.dtype + assert gpu_full_att_kv_state.is_contiguous() and gpu_sliding_state.is_contiguous() + assert cpu_cache_tensor.is_contiguous() + + cpu_cache = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint8) + full_bytes = sliding_config.get_cpu_cache_full_att_bytes(big_page_token_num, tp_world_size) + window_bytes = sliding_config.get_cpu_cache_state_bytes(tp_world_size) + assert cpu_cache.shape[1] == sliding_config.get_cpu_cache_big_page_bytes(big_page_token_num, tp_world_size) + # Packing preserves the original bit patterns, including BF16/FP16 NaNs. + # Gemma's K+V head rows and checkpoint tensors are all uint64-aligned. + full_state = gpu_full_att_kv_state.flatten(2).view(torch.uint64) + window_state = gpu_sliding_state.flatten(1).view(torch.uint64) + cpu_cache = cpu_cache.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] + assert full_rank_size * tp_world_size * 8 == full_bytes + assert window_rank_size * tp_world_size * 8 == window_bytes + + _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=sliding_config.full_layer_num, + FULL_TOKEN_LAYER_SIZE=full_state.shape[2], + FULL_RANK_SIZE=full_rank_size, + FULL_TOTAL_SIZE=full_bytes // 8, + 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, + gpu_sliding_state: torch.Tensor, + cpu_cache_tensor: torch.Tensor, + tp_rank: int, + tp_world_size: int, + big_page_token_num: int, + sliding_config: SlidingWindowCacheConfig, + grid_num: int = 12, +): + """Pack full KV and a raw ring checkpoint 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, + gpu_sliding_state, + cpu_cache_tensor, + tp_rank, + tp_world_size, + big_page_token_num, + sliding_config, + 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, + gpu_sliding_state: torch.Tensor, + cpu_cache_tensor: torch.Tensor, + tp_rank: int, + tp_world_size: int, + big_page_token_num: int, + sliding_config: SlidingWindowCacheConfig, + grid_num: int = 12, +): + """Restore full KV and the unchanged ring checkpoint from a packed CPU page.""" + _copy_state_cache( + mem_indexes, + page_indexes, + None, + big_page_buffer_ids, + gpu_full_att_kv_state, + gpu_sliding_state, + cpu_cache_tensor, + tp_rank, + tp_world_size, + big_page_token_num, + sliding_config, + offload=False, + grid_num=grid_num, + ) 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 index 58a376fa84..fbaa51db2f 100644 --- a/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -21,8 +21,9 @@ class HybridSlidingMemoryManager(MemoryManager): def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): self.sliding_config = sliding_config args = get_env_start_args() - self.enable_prompt_cache = args.use_dynamic_prompt_cache + self.enable_prompt_cache = not args.disable_dynamic_prompt_cache self.small_page_num = args.linear_att_cache_size if self.enable_prompt_cache else 0 + self.cpu_cache_temp_page_num = 2 if args.enable_cpu_cache else 0 self.big_page_token_num = args.linear_att_page_block_num * args.linear_att_hash_page_size super().__init__( size=size, @@ -39,9 +40,10 @@ def _big_page_num(self, token_num): def _cache_nbytes(self, token_num): # Runtime windows already exist when profiling. Reserve BOTH GPU page - # pools here, plus the full-KV hold token and the final partial big page. + # pools here, plus the full-KV hold token, final partial big page, and + # separate CPU-cache load/offload staging states when enabled. return (token_num + 1) * self.get_cell_size() + ( - self.small_page_num + self._big_page_num(token_num) + self.small_page_num + self._big_page_num(token_num) + self.cpu_cache_temp_page_num ) * self.sliding_config.get_state_nbytes() def _profile_token_num(self, available_bytes): @@ -81,6 +83,7 @@ def profile_size(self, mem_fraction): logger.info( f"Sliding-window cache budget: {self.size} full-KV tokens, " f"{self._big_page_num(self.size)} big pages, {self.small_page_num} small pages, " + f"{self.cpu_cache_temp_page_num} CPU-cache staging states, " f"{self._cache_nbytes(self.size) / 1024 ** 3:.2f} GiB (runtime windows already allocated)" ) @@ -88,9 +91,13 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): super()._init_buffers(size, dtype, head_num, head_dim, layer_num) # Keep the existing radix-cache contract; no second alias is needed. self.linear_att_big_page_buffers = SlidingWindowStateCacheManager( - size=self._big_page_num(size), + size=self._big_page_num(size) + self.cpu_cache_temp_page_num, sliding_config=self.sliding_config, + keep_num=self.cpu_cache_temp_page_num, ) + if self.cpu_cache_temp_page_num: + self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.linear_att_big_page_buffers.size - 2 + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.linear_att_big_page_buffers.size - 1 self.sliding_small_page_buffers = SlidingWindowStateCacheManager( size=self.small_page_num, sliding_config=self.sliding_config, diff --git a/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py index 4dafe300b4..07f35cd1c8 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py @@ -1,10 +1,118 @@ import torch +import triton + +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size +from lightllm.utils.envs_utils import get_env_start_args from .normal import NormalMemOperator class HybridSlidingMemOperator(NormalMemOperator): - """GPU KV operations for the token-granular full-attention cache.""" + """Full-KV operations and CPU transfers of hybrid sliding checkpoints.""" + + 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 + + args = get_env_start_args() + page_size = args.cpu_cache_token_page_size + assert mem_indexes.is_cuda and page_indexes.is_cuda + assert page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num + assert len(mem_indexes) % args.linear_att_hash_page_size == 0 + assert triton.cdiv(len(mem_indexes), page_size) == len(page_indexes) + if not len(page_indexes): + return + + mem_manager = self.mem_manager + 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.linear_att_big_page_buffers.alloc_one_state_cache() + assert page_id is not None + req.linear_att_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, + gpu_sliding_state=mem_manager.linear_att_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, + sliding_config=mem_manager.sliding_config, + ) + # 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, + ) + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + args = get_env_start_args() + page_size = args.cpu_cache_token_page_size + assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda + assert page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num + assert len(mem_indexes) % args.linear_att_hash_page_size == 0 + assert triton.cdiv(len(mem_indexes), page_size) == len(page_indexes) == len(page_readies) + if not len(page_indexes): + return + + mem_manager = self.mem_manager + 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.linear_att_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_linear_att_small_page_buffer_id is not None + temp_id = mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID + src_state = radix_cache.linear_att_small_page_buffers.get_state_cache( + req.tail_linear_att_small_page_buffer_id + ) + mem_manager.linear_att_big_page_buffers.get_state_cache(temp_id).copy_(src_state, non_blocking=True) + big_page_ids.append(temp_id) + + assert len(big_page_ids) == len(page_indexes) + 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, + gpu_sliding_state=mem_manager.linear_att_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, + sliding_config=mem_manager.sliding_config, + ) 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 diff --git a/lightllm/common/sliding_window_cache_manager/config.py b/lightllm/common/sliding_window_cache_manager/config.py index 57005170b2..e9599efa7b 100644 --- a/lightllm/common/sliding_window_cache_manager/config.py +++ b/lightllm/common/sliding_window_cache_manager/config.py @@ -45,3 +45,25 @@ def get_state_nbytes(self): for dim in self.get_state_shape(): elements *= dim return elements * self.dtype.itemsize + + def get_cpu_cache_full_att_bytes(self, big_page_token_num: int, tp_world_size: int): + assert big_page_token_num > 0 and tp_world_size > 0 + 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): + assert tp_world_size > 0 + return self.get_state_nbytes() * tp_world_size + + def get_cpu_cache_big_page_bytes(self, big_page_token_num: int, tp_world_size: int): + # 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 diff --git a/lightllm/common/sliding_window_cache_manager/state_cache.py b/lightllm/common/sliding_window_cache_manager/state_cache.py index c60748bf97..1ad8c76cc2 100644 --- a/lightllm/common/sliding_window_cache_manager/state_cache.py +++ b/lightllm/common/sliding_window_cache_manager/state_cache.py @@ -9,9 +9,10 @@ class SlidingWindowStateCacheManager: """GPU storage for immutable request-level sliding-window checkpoints.""" - def __init__(self, size: int, sliding_config: SlidingWindowCacheConfig): + def __init__(self, size: int, sliding_config: SlidingWindowCacheConfig, keep_num: int = 0): self.size = size - assert size >= 0 + self.keep_num = keep_num + assert 0 <= keep_num <= size self.state_cache = torch.empty( (size, *sliding_config.get_state_shape()), dtype=sliding_config.dtype, device="cuda" ) @@ -29,9 +30,10 @@ def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: return [self.free_list.popleft() for _ in range(need_size)] def free_state_cache(self, free_indexes: List[int]): - assert all(0 <= idx < self.size for idx in free_indexes) + alloc_size = self.size - self.keep_num + assert all(0 <= idx < alloc_size for idx in free_indexes) self.free_list.extend(free_indexes) - assert len(self.free_list) <= self.size + assert len(self.free_list) <= alloc_size def get_free_cache_num(self): return len(self.free_list) @@ -41,4 +43,4 @@ def get_used_cache_num(self): def clear_to_init_state(self): self.state_cache.zero_() - self.free_list = collections.deque(range(self.size)) + self.free_list = collections.deque(range(self.size - self.keep_num)) diff --git a/lightllm/models/gemma4/kv_layout.py b/lightllm/models/gemma4/kv_layout.py index 7f02998064..d363cb8bd0 100644 --- a/lightllm/models/gemma4/kv_layout.py +++ b/lightllm/models/gemma4/kv_layout.py @@ -1,3 +1,6 @@ +from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig + + def get_kv_cache_layout(config): """Map Gemma's shared tail layers to physical owners and their last readers.""" layer_types = config["layer_types"] @@ -18,3 +21,23 @@ def get_kv_cache_layout(config): owners.append(owner) last_reader[owner] = layer_index return layer_maps, owners, last_reader + + +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 tp_world_size > 0 + 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/model.py b/lightllm/models/gemma4/model.py index be3885cc13..603bec119b 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -5,11 +5,10 @@ from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager from lightllm.common.req_manager import ReqManagerForSlidingWindow -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig 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 get_kv_cache_layout +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 @@ -91,7 +90,8 @@ def _verify_params(self): args.enable_prefill_microbatch_overlap or args.enable_decode_microbatch_overlap ), "Gemma-4 shared sliding-window KV does not support microbatch overlap yet" assert args.mtp_step == 0, "Gemma-4 hybrid sliding-window cache does not support MTP yet" - assert not args.enable_cpu_cache, "Gemma-4 hybrid sliding-window cache does not support CPU cache" + 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" @@ -102,18 +102,7 @@ def _verify_params(self): def _get_sliding_cache_config(self): if hasattr(self, "sliding_cache_config"): return self.sliding_cache_config - num_global_kv = self.config.get("num_global_key_value_heads") or self.config["num_key_value_heads"] - layer_maps, _, _ = get_kv_cache_layout(self.config) - self.sliding_cache_config = SlidingWindowCacheConfig( - sliding_layer_to_cache_index=layer_maps["sliding_attention"], - full_layer_to_cache_index=layer_maps["full_attention"], - sliding_window=self.config["sliding_window"], - sliding_head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - sliding_head_dim=self.config["head_dim"], - full_head_num=num_global_kv // self.tp_world_size_, - full_head_dim=self.config["global_head_dim"], - dtype=self.data_type, - ) + self.sliding_cache_config = build_sliding_cache_config(self.config, self.tp_world_size_, self.data_type) return self.sliding_cache_config def _init_req_manager(self): diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 6f8973425f..75fa2a035b 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -19,7 +19,7 @@ from lightllm.utils.config_utils import ( has_audio_module, has_vision_module, - is_linear_att_mixed_model, + is_hybrid_att_mixed_model, auto_set_max_req_total_len, auto_set_fused_shared_experts, auto_set_response_parsers, @@ -258,7 +258,7 @@ def _launch_subprocesses(args: StartArgs): # linear att cache 参数自动设置 if args.linear_att_cache_size is None: - # linear_att_cache_size 只会在 qwen3.5 等混合线性层模型中生效。 + # 混合 attention 模型使用请求级状态缓存，保留原有 linear_att 参数名。 default_cache_size = args.running_max_req_size * 2 dp_size_in_node = max(1, args.dp // args.nnodes) per_dp_cache_size = max(1, math.ceil(args.running_max_req_size / dp_size_in_node) * 2) @@ -270,9 +270,9 @@ def _launch_subprocesses(args: StartArgs): # 避免请求释放时将不完整的大页 state 写入 radix cache 并触发断言。 args.linear_att_page_block_num = 10000000 - if args.enable_cpu_cache and is_linear_att_mixed_model(args.model_dir): + if args.enable_cpu_cache and is_hybrid_att_mixed_model(args.model_dir): args.cpu_cache_token_page_size = args.linear_att_hash_page_size * args.linear_att_page_block_num - logger.info(f"set cpu_cache_token_page_size to {args.cpu_cache_token_page_size} for linear hybrid att model") + logger.info(f"set cpu_cache_token_page_size to {args.cpu_cache_token_page_size} for hybrid attention model") # help to manage data stored on Ceph if "s3://" in args.model_dir: diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index a0993099eb..5805d291df 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -11,7 +11,7 @@ from lightllm.server.req_id_generator import convert_sub_id_to_group_id from lightllm.utils.envs_utils import get_unique_server_name from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.config_utils import is_hybrid_att_mixed_model, is_linear_att_mixed_model +from lightllm.utils.config_utils import is_hybrid_att_mixed_model from lightllm.utils.kv_cache_utils import compute_token_list_hash from typing import Any, Dict, List, Union from lightllm.utils.log_utils import init_logger @@ -218,7 +218,7 @@ def init( args = get_env_start_args() if is_hybrid_att_mixed_model(args.model_dir): self._fill_linear_att_token_hash() - if args.enable_cpu_cache and is_linear_att_mixed_model(args.model_dir): + if args.enable_cpu_cache: cpu_cache_hash_list, cpu_cache_page_len_list = self._calcu_linear_att_cpu_cache_page_len_list() self.token_hash_list = TokenHashList() self.token_hash_list.clear() diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 00489b9c27..cc7ec435d0 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -8,7 +8,6 @@ from collections import deque from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient -from lightllm.utils.config_utils import is_linear_att_mixed_model from lightllm.utils.envs_utils import get_env_start_args from ..infer_batch import InferReq from lightllm.utils.dist_utils import create_new_group_for_current_dp @@ -176,7 +175,7 @@ def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> L continue # 过滤不适合进行 kv 卸载到 cpu cache 的请求。 - if g_infer_context.is_linear_att_mixed_model: + if g_infer_context.is_hybrid_att_mixed_model: offload_limit_size = self.args.linear_att_hash_page_size else: offload_limit_size = self.args.cpu_cache_token_page_size @@ -309,7 +308,7 @@ def _start_kv_cache_offload_task( return trans_task def _handle_linear_att_last_page(self, req: InferReq, move_block_size: int, page_len_list: List[int]) -> int: - if not g_infer_context.is_linear_att_mixed_model: + if not g_infer_context.is_hybrid_att_mixed_model: return move_block_size if move_block_size == 0: diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index e81caafe7a..8c5e887b9d 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -16,7 +16,15 @@ get_added_mtp_kv_layer_num, ) from lightllm.utils.log_utils import init_logger -from lightllm.utils.config_utils import get_num_key_value_heads, get_head_dim, get_layer_num, is_linear_att_mixed_model +from lightllm.utils.config_utils import ( + get_config_json, + get_num_key_value_heads, + get_head_dim, + get_layer_num, + is_linear_att_mixed_model, + is_sliding_att_mixed_model, + is_hybrid_att_mixed_model, +) from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class from lightllm.common.kv_cache_mem_manager import ( MemoryManager, @@ -31,6 +39,7 @@ from lightllm.utils.auto_shm_cleanup import register_sysv_shm_for_cleanup from lightllm.utils.dist_utils import get_current_device_id from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager logger = init_logger(__name__) @@ -67,10 +76,31 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": if is_linear_att_mixed_model(args.model_dir): # 对于 qwen3.5 等 linear att 混合模型的特殊处理。 mem_manager_class = Qwen3NextMemManager + elif is_sliding_att_mixed_model(args.model_dir): + mem_manager_class = HybridSlidingMemoryManager else: mem_manager_class = select_mem_manager_class() - if mem_manager_class is Qwen3NextMemManager: + if mem_manager_class is HybridSlidingMemoryManager: + from lightllm.models.gemma4.kv_layout import build_sliding_cache_config + + model_config = get_config_json(args.model_dir) + text_config = model_config.get("text_config", model_config) + tp_world_size = args.tp // args.dp + sliding_config = build_sliding_cache_config(text_config, tp_world_size, get_llm_data_type()) + 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 + cpu_cache_meta = CpuKVCacheMeta( + page_num=0, + token_page_size=1, + layer_num=1, + num_heads=1, + head_dim=sliding_config.get_cpu_cache_big_page_bytes(big_page_token_num, tp_world_size), + data_type=torch.uint8, + scale_head_dim=0, + scale_data_type=get_llm_data_type(), + ) + elif mem_manager_class is Qwen3NextMemManager: linear_config = LinearAttCacheConfig.load_from_args() cpu_cache_meta = CpuKVCacheMeta( page_num=0, @@ -121,10 +151,9 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": if args.mtp_mode is not None: # TODO 可能会存在不同mtp模式的精度问题 - if not is_linear_att_mixed_model(args.model_dir): - # 对于非 linear att 混合模型，需要额外增加 mtp 的 kv 层数， - # 对于 linear att 混合模型，如qwen 3.5 mtp，已经将 kv 数据 - # 打包成一个块了，所以不需要额外增加，其 layer_num 一直都保持为 1 + if not is_hybrid_att_mixed_model(args.model_dir): + # 普通 token KV 需要额外增加 MTP 的 KV 层数；hybrid cache + # 已将所有 payload 打包成字节页，其 layer_num 始终为 1。 cpu_cache_meta.layer_num += get_added_mtp_kv_layer_num() cpu_cache_page_num = int( diff --git a/test/kernel/test_sliding_window_cpu_cache_attention.py b/test/kernel/test_sliding_window_cpu_cache_attention.py new file mode 100644 index 0000000000..f52084a648 --- /dev/null +++ b/test/kernel/test_sliding_window_cpu_cache_attention.py @@ -0,0 +1,106 @@ +from types import SimpleNamespace + +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, +) +from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow +from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig, SlidingWindowStateCacheManager +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") + + +@pytest.mark.parametrize("window,history_len", [(512, 256), (512, 512), (512, 544), (1024, 1056)]) +@pytest.mark.parametrize("q_len", [1, 31]) +def test_cpu_window_checkpoint_resumes_attention_exactly(window, history_len, q_len): + torch.manual_seed(42) + page_size, head_dim, req_idx = 512, 64, 1 + seq_len = history_len + q_len + # Logical layer 3 shares layer 2's KV: only physical owners are stored. + config = SlidingWindowCacheConfig({0: 0, 2: 1, 3: 1}, {1: 0}, window, 1, head_dim, 1, head_dim, torch.bfloat16) + reference = torch.randn((2, seq_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) + full_kv = torch.randn((1, history_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) + expected_full_kv = full_kv.clone() + endpoints = list(range(page_size, history_len + 1, page_size)) + if history_len % page_size: + endpoints.append(history_len) + pages = SlidingWindowStateCacheManager(len(endpoints), config) + for page_id, endpoint in enumerate(endpoints): + positions = torch.arange(max(0, endpoint - window), endpoint, device="cuda") + pages.state_cache[page_id, :, positions % window] = reference[:, positions] + + int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + mem_indexes = torch.full((len(endpoints) * page_size,), -1, device="cuda", dtype=torch.int32) + mem_indexes[:history_len] = torch.arange(history_len, device="cuda", dtype=torch.int32) + page_ids = int_tensor(list(range(len(endpoints)))) + cpu_cache = torch.zeros( + (len(endpoints), 1, 1, 1, config.get_cpu_cache_big_page_bytes(page_size, 1)), + dtype=torch.uint8, + pin_memory=True, + ) + copy_args = dict( + mem_indexes=mem_indexes, + page_indexes=page_ids, + big_page_buffer_ids=page_ids, + gpu_full_att_kv_state=full_kv, + gpu_sliding_state=pages.state_cache, + cpu_cache_tensor=cpu_cache, + tp_rank=0, + tp_world_size=1, + big_page_token_num=page_size, + sliding_config=config, + ) + copy_kv_buffer_to_cpu_cache(page_readies=torch.zeros_like(page_ids, dtype=torch.bool), **copy_args) + full_kv.fill_(-7) + pages.state_cache.fill_(-9) + copy_cpu_cache_to_kv_buffer(**copy_args) + torch.testing.assert_close(full_kv, expected_full_kv, atol=0, rtol=0) + + manager = object.__new__(ReqManagerForSlidingWindow) + manager.sliding_config, manager.sliding_window = config, window + manager.scratch_token_num, manager.scratch_start = q_len, 3 * window + manager.mem_manager = SimpleNamespace(linear_att_big_page_buffers=pages) + manager.req_to_sliding_window = torch.zeros( + (2, manager.scratch_start + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16 + ) + manager.req_to_sliding_window_indexs = torch.full((3, seq_len), -1, device="cuda", dtype=torch.int32) + manager.restore_big_page_state(len(endpoints) - 1, SimpleNamespace(req_idx=req_idx)) + manager.req_to_sliding_window[:, manager.scratch_start :] = reference[:, history_len:] + state = SimpleNamespace( + input_ids=int_tensor([0] * q_len), + b_req_idx=int_tensor([req_idx]), + b_seq_len=int_tensor([seq_len]), + b_q_seq_len=int_tensor([q_len]), + b_q_start_loc=int_tensor([0]), + max_q_seq_len=q_len, + ) + manager.prepare_sliding_window(state) + q = torch.randn((q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) + reference_indexes = torch.arange(seq_len, device="cuda", dtype=torch.int32).expand(3, -1) + image_end = int_tensor([0] * q_len) + for layer_index in [0, 2, 3]: + physical_layer = config.get_sliding_layer_index(layer_index) + actual, expected = torch.empty_like(q), torch.empty_like(q) + for kv, indexes, output in [ + (manager.req_to_sliding_window[physical_layer], manager.req_to_sliding_window_indexs, actual), + (reference[physical_layer], reference_indexes, expected), + ]: + context_attention_fwd_gemma4_mm( + q, + kv[:, :1], + kv[:, 1:], + output, + state.b_req_idx, + state.b_q_start_loc, + state.b_seq_len, + int_tensor([history_len]), + q_len, + indexes, + image_end, + sliding_window=(window - 1, 0), + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) 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..7f16fc64ca --- /dev/null +++ b/test/kernel/test_sliding_window_cpu_cache_copy.py @@ -0,0 +1,170 @@ +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, +) +from lightllm.common.sliding_window_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) + 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(), + gpu_sliding_state=window_cpu.cuda(), + cpu_cache_tensor=cpu_cache, + tp_rank=rank, + tp_world_size=tp_world_size, + big_page_token_num=big_page_tokens, + sliding_config=config, + 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_gpu = expected_full.cuda(), expected_window.cuda() + 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, + gpu_sliding_state=window_gpu, + cpu_cache_tensor=cpu_cache, + tp_rank=rank, + tp_world_size=tp_world_size, + big_page_token_num=big_page_tokens, + sliding_config=config, + grid_num=3, + ) + torch.cuda.synchronize() + _assert_same_bits(full_gpu, expected_full) + _assert_same_bits(window_gpu, expected_window) + _assert_same_bits(cpu_cache.view(cpu_page_num, page_bytes), expected_cache) + + +def test_empty_copy_is_a_noop(): + config = SlidingWindowCacheConfig({0: 0}, {1: 0}, 8, 1, 64, 1, 64, torch.bfloat16) + 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, + gpu_sliding_state=None, + cpu_cache_tensor=None, + tp_rank=0, + tp_world_size=1, + big_page_token_num=16, + sliding_config=config, + ) + copy_kv_buffer_to_cpu_cache(page_readies=indexes, **kwargs) + copy_cpu_cache_to_kv_buffer(**kwargs) diff --git a/test/utils/test_sliding_cpu_cache_meta.py b/test/utils/test_sliding_cpu_cache_meta.py new file mode 100644 index 0000000000..5ec84b9166 --- /dev/null +++ b/test/utils/test_sliding_cpu_cache_meta.py @@ -0,0 +1,155 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig +from lightllm.models.gemma4.kv_layout import build_sliding_cache_config + + +def _gemma_config(shared): + layer_num = 42 if shared else 60 + return { + "model_type": "gemma4_text", + "num_hidden_layers": layer_num, + "num_attention_heads": 8, + "num_key_value_heads": 2 if shared else 16, + "num_global_key_value_heads": None if shared else 4, + "head_dim": 256, + "global_head_dim": 512, + "sliding_window": 512 if shared else 1024, + "layer_types": (["sliding_attention"] * 5 + ["full_attention"]) * (layer_num // 6), + "num_kv_shared_layers": 18 if shared else 0, + } + + +@pytest.mark.parametrize("shared,tp_world_size", [(True, 1), (True, 2), (False, 1), (False, 2), (False, 4)]) +def test_cpu_page_layout_uses_physical_owners_and_all_tp_shards(shared, tp_world_size): + config = _gemma_config(shared) + layout = build_sliding_cache_config(config, tp_world_size, torch.bfloat16) + full_layers, sliding_layers = (4, 20) if shared else (10, 50) + full_heads = config["num_global_key_value_heads"] or config["num_key_value_heads"] + big_page_tokens = 2048 + full_bytes = full_layers * big_page_tokens * 2 * full_heads * config["global_head_dim"] * 2 + state_bytes = sliding_layers * config["sliding_window"] * 2 * config["num_key_value_heads"] * config["head_dim"] * 2 + assert layout.get_cpu_cache_full_att_bytes(big_page_tokens, tp_world_size) == full_bytes + assert layout.get_cpu_cache_state_bytes(tp_world_size) == state_bytes + assert layout.get_cpu_cache_big_page_bytes(big_page_tokens, tp_world_size) == full_bytes + state_bytes + + +def test_cpu_page_payload_is_aligned_without_changing_section_sizes(): + layout = SlidingWindowCacheConfig({0: 0}, {1: 0}, 1, 1, 3, 1, 5, torch.bfloat16) + assert layout.get_cpu_cache_full_att_bytes(3, 1) == 60 + assert layout.get_cpu_cache_state_bytes(1) == 12 + assert layout.get_cpu_cache_big_page_bytes(3, 1) == 80 + + +@pytest.mark.parametrize("shared,tp_world_size", [(True, 4), (False, 8)]) +def test_cpu_layout_rejects_replicated_kv_heads(shared, tp_world_size): + with pytest.raises(AssertionError, match="KV heads must be divisible"): + build_sliding_cache_config(_gemma_config(shared), tp_world_size, torch.bfloat16) + + +@pytest.mark.parametrize("wrapped", [False, True]) +def test_sliding_cpu_meta_is_a_flat_global_payload(monkeypatch, wrapped): + import lightllm.utils.kv_cache_utils as cache_utils + + config = _gemma_config(True) + layout = build_sliding_cache_config(config, 2, torch.bfloat16) + page_bytes = layout.get_cpu_cache_big_page_bytes(2048, 2) + args = SimpleNamespace( + model_dir="gemma-test", + enable_cpu_cache=True, + tp=4, + dp=2, + linear_att_hash_page_size=128, + linear_att_page_block_num=16, + cpu_cache_token_page_size=2048, + cpu_cache_storage_size=3 * page_bytes / 1024 ** 3, + mtp_mode=None, + ) + monkeypatch.setattr(cache_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr(cache_utils, "is_linear_att_mixed_model", lambda _: False) + monkeypatch.setattr(cache_utils, "is_sliding_att_mixed_model", lambda _: True) + monkeypatch.setattr(cache_utils, "get_llm_data_type", lambda: torch.bfloat16) + monkeypatch.setattr(cache_utils, "get_config_json", lambda _: {"text_config": config} if wrapped else config) + meta = cache_utils.calcu_cpu_cache_meta.__wrapped__() + assert meta.data_type == torch.uint8 + assert (meta.layer_num, meta.token_page_size, meta.num_heads) == (1, 1, 1) + assert meta.head_dim == meta.calcu_one_page_size() == page_bytes + assert meta.page_num == 3 + assert args.cpu_cache_token_page_size == 2048 + + +def test_hybrid_request_initializes_cpu_hashes_without_linear_state(monkeypatch): + import lightllm.server.core.objs.req as req_module + + prompt = list(range(18)) + args = SimpleNamespace( + model_dir="gemma-test", + mtp_step=0, + enable_cpu_cache=True, + linear_att_hash_page_size=4, + linear_att_page_block_num=3, + cpu_cache_token_page_size=12, + ) + monkeypatch.setattr(req_module, "get_env_start_args", lambda: args) + monkeypatch.setattr(req_module, "is_hybrid_att_mixed_model", lambda _: True) + req = SimpleNamespace(index_in_shm_mem=0, ref_count=0) + req.create_logprobs_shm_array = lambda: None + req.create_prompt_ids_shm_array = lambda: setattr( + req, "shm_prompt_ids", SimpleNamespace(arr=np.empty(2048, dtype=np.int64)) + ) + req.post_init = lambda: None + req.get_prompt_ids = lambda: prompt + req._fill_linear_att_token_hash = lambda: req_module.Req._fill_linear_att_token_hash(req) + req._calcu_linear_att_cpu_cache_page_len_list = lambda: req_module.Req._calcu_linear_att_cpu_cache_page_len_list( + req + ) + req_module.Req.init(req, 0, prompt, req_module.SamplingParams(), tokenizer=None, chunked_prefill_size=16) + hashes = req.linear_att_token_hash_list.get_all() + assert len(hashes) == 4 + assert req.token_hash_list.get_all() == [hashes[2], hashes[3]] + assert req.token_hash_page_len_list.get_all() == [12, 16] + assert req.cpu_cache_match_page_indexes.get_all() == [] + + +@pytest.mark.parametrize("disable_tail,tail_buffer", [(False, None), (False, 3), (True, 3)]) +def test_sliding_offload_uses_existing_hybrid_tail_policy(monkeypatch, disable_tail, tail_buffer): + import lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache as cache_module + + module = object.__new__(cache_module.MultiLevelKvCacheModule) + module.args = SimpleNamespace(cpu_cache_token_page_size=16, disable_linear_att_small_page_cpu_cache=disable_tail) + monkeypatch.setattr(cache_module.g_infer_context, "is_linear_att_mixed_model", False) + monkeypatch.setattr(cache_module.g_infer_context, "is_hybrid_att_mixed_model", True) + req = SimpleNamespace(tail_linear_att_small_page_buffer_id=tail_buffer) + expected_pages = 2 if not disable_tail and tail_buffer is not None else 1 + assert module._handle_linear_att_last_page(req, 2, [16, 20]) == expected_pages + + +@pytest.mark.parametrize("enable_cpu_cache,disable_gpu_cache", [(True, False), (True, True), (False, True)]) +def test_cpu_state_transfer_requires_gpu_prefix_cache(monkeypatch, enable_cpu_cache, disable_gpu_cache): + import lightllm.models.gemma4.model as gemma_model + + model = object.__new__(gemma_model.Gemma4TpPartModel) + model.load_way, model.tp_world_size_ = "HF", 2 + model.config = _gemma_config(True) + args = SimpleNamespace( + mtp_step=0, + enable_cpu_cache=enable_cpu_cache, + disable_dynamic_prompt_cache=disable_gpu_cache, + disable_chunked_prefill=False, + run_mode="normal", + llm_kv_type="None", + enable_dp_prompt_cache_fetch=False, + diverse_mode=False, + enable_prefill_microbatch_overlap=False, + enable_decode_microbatch_overlap=False, + ) + monkeypatch.setattr(gemma_model, "get_env_start_args", lambda: args) + if enable_cpu_cache and disable_gpu_cache: + with pytest.raises(AssertionError, match="CPU cache requires GPU prefix cache"): + model._verify_params() + else: + model._verify_params() diff --git a/test/utils/test_sliding_window_cache.py b/test/utils/test_sliding_window_cache.py index 3cab1fd56f..84cb494f18 100644 --- a/test/utils/test_sliding_window_cache.py +++ b/test/utils/test_sliding_window_cache.py @@ -24,7 +24,7 @@ def test_gemma_physical_owners_and_last_readers(layer_num, shared, sliding_num, assert owners[41] == 23 and last_readers[23] == 41 -def _memory_manager(big_page_tokens=2048, small_pages=8, enabled=True): +def _memory_manager(big_page_tokens=2048, small_pages=8, enabled=True, cpu_cache=False): manager = object.__new__(HybridSlidingMemoryManager) manager.head_num, manager.head_dim, manager.layer_num, manager.dtype = 1, 512, 10, torch.bfloat16 manager.sliding_config = SlidingWindowCacheConfig( @@ -35,6 +35,7 @@ def _memory_manager(big_page_tokens=2048, small_pages=8, enabled=True): small_pages, enabled, ) + manager.cpu_cache_temp_page_num = 2 if cpu_cache else 0 return manager @@ -60,12 +61,42 @@ def test_disabled_prompt_cache_does_not_reserve_pages(): assert manager._cache_nbytes(4096) == 4097 * manager.get_cell_size() +def test_cpu_cache_reserves_two_additional_window_checkpoints(): + gpu_only = _memory_manager() + cpu_cache = _memory_manager(cpu_cache=True) + assert cpu_cache._cache_nbytes(4096) == ( + gpu_only._cache_nbytes(4096) + 2 * cpu_cache.sliding_config.get_state_nbytes() + ) + + +@pytest.mark.parametrize("disabled", [False, True]) +def test_page_pools_follow_active_prompt_cache_flag(monkeypatch, disabled): + import lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager as memory_module + + args = SimpleNamespace( + use_dynamic_prompt_cache=False, + disable_dynamic_prompt_cache=disabled, + enable_cpu_cache=False, + linear_att_cache_size=3, + linear_att_hash_page_size=32, + linear_att_page_block_num=8, + ) + monkeypatch.setattr(memory_module, "get_env_start_args", lambda: args) + monkeypatch.setattr(memory_module.MemoryManager, "__init__", lambda self, **kwargs: None) + config = SlidingWindowCacheConfig({0: 0}, {1: 0}, 32, 1, 64, 1, 64, torch.bfloat16) + manager = HybridSlidingMemoryManager(size=256, sliding_config=config) + assert manager.enable_prompt_cache is not disabled + assert manager.small_page_num == (0 if disabled else 3) + assert manager._big_page_num(256) == (0 if disabled else 1) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_profiled_gpu_pools_match_reserved_bytes_and_are_reused(monkeypatch): +@pytest.mark.parametrize("cpu_cache", [False, True]) +def test_profiled_gpu_pools_match_reserved_bytes_and_are_reused(monkeypatch, cpu_cache): import lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager as memory_module from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow - manager = _memory_manager(big_page_tokens=32, small_pages=2) + manager = _memory_manager(big_page_tokens=32, small_pages=2, cpu_cache=cpu_cache) manager.head_num, manager.head_dim, manager.layer_num = 1, 64, 1 manager.sliding_config = SlidingWindowCacheConfig({0: 0}, {1: 0}, 32, 1, 64, 1, 64, torch.bfloat16) manager.size = None diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index e5ccaa3a35..f9aebf83bc 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -98,7 +98,7 @@ def start_offload(req, cpu_kv_cache_stream): return SimpleNamespace(req=req) module._start_kv_cache_offload_task = start_offload - monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "is_linear_att_mixed_model", False) + monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "is_hybrid_att_mixed_model", False) monkeypatch.setattr( multi_level_kv_cache_impl.g_infer_context, "get_cpu_kv_cache_stream", diff --git a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py index 03fc6b3e25..471df7ae09 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py @@ -37,9 +37,10 @@ def test_snapshot_outside_cacheable_boundaries_does_not_allocate_or_copy(chunk_e assert req.tail_linear_att_small_page_buffer_id is None -def _cpu_state_cache(size): +def _cpu_state_cache(size, keep_num=0): pages = object.__new__(SlidingWindowStateCacheManager) pages.size = size + pages.keep_num = keep_num pages.state_cache = torch.empty((size, 2, 4, 2, 4), dtype=torch.float32) pages.clear_to_init_state() return pages @@ -78,3 +79,17 @@ def test_sliding_state_pool_exhaustion_and_released_slot_reuse(): pages.free_state_cache([0, 1]) assert pages.get_free_cache_num() == 2 assert pages.get_used_cache_num() == 0 + + +def test_sliding_state_pool_preserves_cpu_transfer_slots(): + pages = _cpu_state_cache(5, keep_num=2) + assert pages.alloc_state_cache(3) == [0, 1, 2] + assert pages.alloc_one_state_cache() is None + for reserved_id in [3, 4]: + with pytest.raises(AssertionError): + pages.free_state_cache([reserved_id]) + pages.free_state_cache([0, 1, 2]) + assert pages.get_free_cache_num() == 3 + pages.clear_to_init_state() + assert pages.alloc_state_cache(3) == [0, 1, 2] + assert pages.alloc_one_state_cache() is None diff --git a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py new file mode 100644 index 0000000000..8a4247dac8 --- /dev/null +++ b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py @@ -0,0 +1,125 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel import sliding_window_cpu_cache_copy as copy_kernels +from lightllm.common.kv_cache_mem_manager.operator import hybrid_sliding as operator_module +from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow +from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as cache_module + + +@pytest.mark.parametrize( + "gpu_prefix,cpu_prefix,expected_endpoints", + [(288, 736, {512: 0}), (288, 768, {512: 1, 768: 0}), (544, 736, {})], +) +def test_cpu_load_prepends_partial_gpu_page_and_restores_absolute_checkpoint( + monkeypatch, gpu_prefix, cpu_prefix, expected_endpoints +): + # Run the real public loader, sliding operator and runtime restore with + # CPU tensors. Only the transfer kernel and CUDA/distributed calls are + # replaced; GPU kernel and stream behavior have separate coverage. + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, non_blocking=False: self) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None)) + monkeypatch.setattr(cache_module.dist, "barrier", lambda group: None) + args = SimpleNamespace(cpu_cache_token_page_size=256, linear_att_hash_page_size=32, linear_att_page_block_num=8) + monkeypatch.setattr(operator_module, "get_env_start_args", lambda: args) + monkeypatch.setattr(operator_module, "get_current_rank_in_dp", lambda: 0) + monkeypatch.setattr(operator_module, "get_dp_world_size", lambda: 1) + + allocated_tokens, evicted_tokens, dereferenced_pages, transfers = [], [], [], [] + states = torch.zeros((6, 1, 4, 2, 4)) + free_ids = iter(range(4)) + state_pool = SimpleNamespace( + state_cache=states, + alloc_one_state_cache=lambda: next(free_ids), + get_state_cache=lambda index: states[index], + ) + + def alloc(need_size): + allocated_tokens.append(need_size) + return torch.arange(1000, 1000 + need_size, dtype=torch.int32) + + mem_manager = SimpleNamespace( + alloc=alloc, + sliding_config=object(), + kv_buffer=object(), + linear_att_big_page_buffers=state_pool, + CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID=4, + ) + mem_manager.operator = operator_module.HybridSlidingMemOperator(mem_manager) + req_manager = object.__new__(ReqManagerForSlidingWindow) + req_manager.mem_manager, req_manager.sliding_window = mem_manager, 4 + req_manager.req_to_sliding_window = torch.full((1, 8, 2, 4), -1.0) + req_manager.req_to_token_indexs = torch.full((2, 1024), -1, dtype=torch.int32) + req_manager.req_to_token_indexs[1, :gpu_prefix] = torch.arange(10000, 10000 + gpu_prefix, dtype=torch.int32) + original_mapping = req_manager.req_to_token_indexs.clone() + radix_cache = SimpleNamespace( + free_radix_cache_to_get_enough_token=lambda need_token_num: evicted_tokens.append(need_token_num) + ) + monkeypatch.setattr(cache_module.g_infer_context, "req_manager", req_manager) + monkeypatch.setattr(cache_module.g_infer_context, "radix_cache", radix_cache) + monkeypatch.setattr(cache_module.g_infer_context, "get_can_alloc_token_num", lambda: 2048) + + req = SimpleNamespace( + req_idx=1, + cur_kv_len=gpu_prefix, + linear_att_len_to_big_page_id={}, + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + shm_req=SimpleNamespace( + input_len=cpu_prefix + 1, + disk_prompt_cache_len=0, + cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: [4, 8, 12]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [256, 512, cpu_prefix]), + ), + ) + + def load(**kwargs): + # cur_kv_len must already be the absolute CPU endpoint when the + # operator assigns full-page checkpoints, not the old GPU hit length. + assert req.cur_kv_len == cpu_prefix + transfers.append(kwargs) + for state_id, cpu_page in zip(kwargs["big_page_buffer_ids"], kwargs["page_indexes"]): + kwargs["gpu_sliding_state"][state_id].fill_(cpu_page.item()) + + monkeypatch.setattr(copy_kernels, "copy_cpu_cache_to_kv_buffer", load) + module = object.__new__(cache_module.MultiLevelKvCacheModule) + module.backend = SimpleNamespace( + is_master_in_dp=True, + radix_cache=radix_cache, + model=SimpleNamespace(mem_manager=mem_manager, req_manager=req_manager), + ) + module.need_sync_compute_stream = lambda: False + module.init_sync_group = object() + module.cpu_cache_client = SimpleNamespace( + cpu_kv_cache_tensor=object(), + lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), + deref_pages=lambda page_list: dereferenced_pages.extend(page_list), + ) + + module.load_cpu_cache_to_reqs([req]) + + need_tokens = cpu_prefix - gpu_prefix + page_start = gpu_prefix // 256 * 256 + new_indexes = torch.arange(1000, 1000 + need_tokens, dtype=torch.int32) + expected_transfer = torch.cat([original_mapping[1, page_start:gpu_prefix], new_indexes]) + padding = (-len(expected_transfer)) % 256 + assert allocated_tokens == evicted_tokens == [need_tokens] + assert len(transfers) == 1 + assert transfers[0]["mem_indexes"].tolist() == expected_transfer.tolist() + [-1] * padding + assert transfers[0]["page_indexes"].tolist() == [4, 8, 12][gpu_prefix // 256 :] + assert req.linear_att_len_to_big_page_id == expected_endpoints + assert 4 not in req.linear_att_len_to_big_page_id.values() + if cpu_prefix % 256: + assert transfers[0]["big_page_buffer_ids"][-1].item() == 4 + torch.testing.assert_close(req_manager.req_to_token_indexs[1, :gpu_prefix], original_mapping[1, :gpu_prefix]) + torch.testing.assert_close(req_manager.req_to_token_indexs[1, gpu_prefix:cpu_prefix], new_indexes) + assert torch.all(req_manager.req_to_token_indexs[1, cpu_prefix:] == -1) + assert torch.all(req_manager.req_to_sliding_window[:, :4] == -1) + assert torch.all(req_manager.req_to_sliding_window[:, 4:8] == 12) + assert req.shm_req.cpu_prompt_cache_len == need_tokens + assert req.shm_req.shm_cur_kv_len == cpu_prefix + # Dereference all matched pages, including the page already covered by + # the GPU prefix and omitted from the actual transfer. + assert dereferenced_pages == [4, 8, 12] diff --git a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py new file mode 100644 index 0000000000..fedcccabca --- /dev/null +++ b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py @@ -0,0 +1,259 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel import sliding_window_cpu_cache_copy as copy_kernels +from lightllm.common.kv_cache_mem_manager.operator import hybrid_sliding as operator_module +from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow +from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig, SlidingWindowStateCacheManager +from lightllm.server.router.model_infer.infer_batch import g_infer_context + + +def _cpu_pages(size, keep_num=0): + pages = object.__new__(SlidingWindowStateCacheManager) + pages.size, pages.keep_num = size, keep_num + pages.state_cache = torch.empty((size, 1, 4, 2, 4), dtype=torch.float32) + pages.clear_to_init_state() + return pages + + +@pytest.fixture +def sliding_operator(monkeypatch): + # Exercise page ownership and transfer orchestration on CPU; kernel tests + # separately cover real CUDA pointers, byte layout, and stream ordering. + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, non_blocking=False: self) + monkeypatch.setattr( + operator_module, + "get_env_start_args", + lambda: SimpleNamespace(cpu_cache_token_page_size=8, linear_att_hash_page_size=2, linear_att_page_block_num=4), + ) + monkeypatch.setattr(operator_module, "get_current_rank_in_dp", lambda: 0) + monkeypatch.setattr(operator_module, "get_dp_world_size", lambda: 1) + manager = SimpleNamespace( + sliding_config=object(), + kv_buffer=torch.zeros((1, 33, 2, 4)), + linear_att_big_page_buffers=_cpu_pages(6, keep_num=2), + CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID=4, + CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID=5, + ) + req_manager = object.__new__(ReqManagerForSlidingWindow) + req_manager.mem_manager = manager + req_manager.sliding_window = 4 + req_manager.req_to_sliding_window = torch.full((1, 8, 2, 4), -1.0) + small_pages = _cpu_pages(2) + small_pages.get_state_cache(0).fill_(17) + radix = SimpleNamespace( + linear_att_small_page_buffers=small_pages, + get_big_page_ids_by_node=lambda node: [] if node is None else node.big_page_ids.copy(), + ) + monkeypatch.setattr(g_infer_context, "req_manager", req_manager) + monkeypatch.setattr(g_infer_context, "radix_cache", radix) + return operator_module.HybridSlidingMemOperator(manager), req_manager, small_pages + + +@pytest.mark.parametrize( + "cached_tokens,token_num,expected_endpoints,has_tail", + [(0, 6, {}, True), (0, 16, {8: 1, 16: 0}, False), (0, 22, {8: 1, 16: 0}, True), (8, 22, {16: 1, 24: 0}, True)], +) +def test_load_restores_last_checkpoint_without_owning_tail_staging_slot( + monkeypatch, sliding_operator, cached_tokens, token_num, expected_endpoints, has_tail +): + operator, req_manager, _ = sliding_operator + page_num = (token_num + 7) // 8 + captures = [] + + def load(**kwargs): + captures.append(kwargs) + for buffer_id, page_id in zip(kwargs["big_page_buffer_ids"].tolist(), kwargs["page_indexes"].tolist()): + kwargs["gpu_sliding_state"][buffer_id].fill_(10 + page_id) + + monkeypatch.setattr(copy_kernels, "copy_cpu_cache_to_kv_buffer", load) + req = SimpleNamespace(req_idx=1, cur_kv_len=cached_tokens + token_num, linear_att_len_to_big_page_id={}) + operator.load_cpu_cache_to_gpu( + torch.arange(cached_tokens, cached_tokens + token_num, dtype=torch.int32), + torch.arange(cached_tokens // 8, cached_tokens // 8 + page_num, dtype=torch.int32), + SimpleNamespace(cpu_kv_cache_tensor=object()), + req, + ) + + assert req.linear_att_len_to_big_page_id == expected_endpoints + assert len(captures) == 1 + padded_indexes = captures[0]["mem_indexes"] + assert padded_indexes.tolist() == list(range(cached_tokens, cached_tokens + token_num)) + [-1] * ( + page_num * 8 - token_num + ) + assert 4 not in req.linear_att_len_to_big_page_id.values() + assert operator.mem_manager.linear_att_big_page_buffers.get_free_cache_num() == 4 - token_num // 8 + if has_tail: + assert captures[0]["big_page_buffer_ids"][-1].item() == 4 + torch.testing.assert_close( + req_manager.req_to_sliding_window[:, 4:8], + torch.full((1, 4, 2, 4), 10.0 + cached_tokens // 8 + page_num - 1), + atol=0, + rtol=0, + ) + assert torch.all(req_manager.req_to_sliding_window[:, :4] == -1) + + +@pytest.mark.parametrize("token_num", [16, 22]) +def test_offload_combines_shared_owned_and_tail_checkpoints(monkeypatch, sliding_operator, token_num): + operator, _, small_pages = sliding_operator + captures = [] + monkeypatch.setattr(copy_kernels, "copy_kv_buffer_to_cpu_cache", lambda **kwargs: captures.append(kwargs)) + big_pages = operator.mem_manager.linear_att_big_page_buffers + big_pages.get_state_cache(1).fill_(11) + big_pages.get_state_cache(3).fill_(13) + req = SimpleNamespace( + shared_kv_node=SimpleNamespace(big_page_ids=[1]), + linear_att_len_to_big_page_id={16: 3}, + tail_linear_att_small_page_buffer_id=0 if token_num % 8 else None, + ) + page_num = (token_num + 7) // 8 + ready = torch.tensor([True] + [False] * (page_num - 1)) + operator.offload_gpu_kv_to_cpu_cache( + torch.arange(token_num, dtype=torch.int32), + torch.arange(page_num, dtype=torch.int32), + ready, + SimpleNamespace(cpu_kv_cache_tensor=object()), + req, + ) + + assert len(captures) == 1 + assert captures[0]["big_page_buffer_ids"].tolist() == ([1, 3, 5] if token_num % 8 else [1, 3]) + assert captures[0]["mem_indexes"].tolist() == list(range(token_num)) + [-1] * (page_num * 8 - token_num) + assert captures[0]["page_readies"] is ready + assert req.linear_att_len_to_big_page_id == {16: 3} + assert torch.count_nonzero(big_pages.get_state_cache(4)) == 0 + if token_num % 8: + torch.testing.assert_close(big_pages.get_state_cache(5), small_pages.get_state_cache(0), atol=0, rtol=0) + small_pages.get_state_cache(0).zero_() + assert torch.all(big_pages.get_state_cache(5) == 17) + + +def test_reserved_state_slots_are_never_allocated_or_freed(): + pages = _cpu_pages(4, keep_num=2) + assert pages.alloc_state_cache(2) == [0, 1] + assert pages.alloc_one_state_cache() is None + for reserved_id in [2, 3]: + with pytest.raises(AssertionError): + pages.free_state_cache([reserved_id]) + pages.free_state_cache([0, 1]) + assert pages.get_free_cache_num() == 2 + assert pages.get_used_cache_num() == 2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_real_cpu_transfers_reuse_separate_load_and_offload_slots_across_streams(monkeypatch): + request_num, page_size, tail_len, window = 8, 8, 6, 4 + config = SlidingWindowCacheConfig({0: 0}, {1: 0}, window, 1, 8, 1, 8, torch.bfloat16) + full_bytes = config.get_cpu_cache_full_att_bytes(page_size, 1) + state_bytes = config.get_cpu_cache_state_bytes(1) + cpu_cache = torch.zeros( + (request_num * 2, config.get_cpu_cache_big_page_bytes(page_size, 1)), + dtype=torch.uint8, + device="cpu", + pin_memory=True, + ) + + def full_page(page_id): + return cpu_cache[page_id, :full_bytes].view(config.dtype).view(page_size, 1, 2, 8) + + def window_page(page_id): + return ( + cpu_cache[page_id, full_bytes : full_bytes + state_bytes].view(config.dtype).view(config.get_state_shape()) + ) + + # The load stream reads already-ready pages, while offload writes disjoint + # CPU pages. Both directions share the same big-state pool, as in serving. + for req_idx in range(request_num): + full_page(req_idx).fill_(10 + req_idx) + window_page(req_idx).fill_(200 + req_idx) + + big_pages = SlidingWindowStateCacheManager(2, config, keep_num=2) + small_pages = SlidingWindowStateCacheManager(request_num, config) + manager = SimpleNamespace( + sliding_config=config, + kv_buffer=torch.zeros((1, request_num * tail_len * 2, 2, 8), dtype=config.dtype, device="cuda"), + linear_att_big_page_buffers=big_pages, + CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID=0, + CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID=1, + ) + req_manager = object.__new__(ReqManagerForSlidingWindow) + req_manager.mem_manager, req_manager.sliding_window = manager, window + req_manager.req_to_sliding_window = torch.zeros((1, request_num * window, 2, 8), dtype=config.dtype, device="cuda") + monkeypatch.setattr(g_infer_context, "req_manager", req_manager) + monkeypatch.setattr( + g_infer_context, + "radix_cache", + SimpleNamespace(linear_att_small_page_buffers=small_pages, get_big_page_ids_by_node=lambda node: []), + ) + monkeypatch.setattr( + operator_module, + "get_env_start_args", + lambda: SimpleNamespace( + cpu_cache_token_page_size=page_size, linear_att_hash_page_size=2, linear_att_page_block_num=4 + ), + ) + monkeypatch.setattr(operator_module, "get_current_rank_in_dp", lambda: 0) + monkeypatch.setattr(operator_module, "get_dp_world_size", lambda: 1) + operator = operator_module.HybridSlidingMemOperator(manager) + client = SimpleNamespace(cpu_kv_cache_tensor=cpu_cache) + transfers = [] + for req_idx in range(request_num): + start = req_idx * tail_len + manager.kv_buffer[:, start : start + tail_len].fill_(30 + req_idx) + small_page_id = small_pages.alloc_one_state_cache() + small_pages.get_state_cache(small_page_id).fill_(100 + req_idx) + req = SimpleNamespace( + req_idx=req_idx, + cur_kv_len=tail_len, + shared_kv_node=None, + linear_att_len_to_big_page_id={}, + tail_linear_att_small_page_buffer_id=small_page_id, + ) + source_indexes = torch.arange(start, start + tail_len, dtype=torch.int32, device="cuda") + load_indexes = source_indexes + request_num * tail_len + load_page = torch.tensor([req_idx], dtype=torch.int32, device="cuda") + offload_page = torch.tensor([request_num + req_idx], dtype=torch.int32, device="cuda") + ready = torch.tensor([False], dtype=torch.bool, device="cuda") + transfers.append((req, source_indexes, load_indexes, load_page, offload_page, ready)) + + offload_stream, load_stream = torch.cuda.Stream(), torch.cuda.Stream() + offload_stream.wait_stream(torch.cuda.current_stream()) + load_stream.wait_stream(torch.cuda.current_stream()) + for req, source_indexes, load_indexes, load_page, offload_page, ready in transfers: + with torch.cuda.stream(offload_stream): + operator.offload_gpu_kv_to_cpu_cache(source_indexes, offload_page, ready, client, req) + with torch.cuda.stream(load_stream): + operator.load_cpu_cache_to_gpu(load_indexes, load_page, client, req) + # No per-request wait: each reserved slot has been reused eight times. + offload_stream.synchronize() + load_stream.synchronize() + + for req_idx, (req, source_indexes, load_indexes, _, _, _) in enumerate(transfers): + assert req.linear_att_len_to_big_page_id == {} + torch.testing.assert_close( + manager.kv_buffer[:, load_indexes], + torch.full((1, tail_len, 2, 8), 10 + req_idx, dtype=config.dtype, device="cuda"), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + req_manager.req_to_sliding_window[:, req_idx * window : (req_idx + 1) * window], + torch.full(config.get_state_shape(), 200 + req_idx, dtype=config.dtype, device="cuda"), + atol=0, + rtol=0, + ) + assert torch.all(full_page(request_num + req_idx)[:tail_len] == 30 + req_idx) + assert torch.count_nonzero(full_page(request_num + req_idx)[tail_len:]) == 0 + assert torch.all(window_page(request_num + req_idx) == 100 + req_idx) + assert torch.all(full_page(req_idx) == 10 + req_idx) + assert torch.all(window_page(req_idx) == 200 + req_idx) + assert torch.all(manager.kv_buffer[:, source_indexes] == 30 + req_idx) + assert big_pages.get_free_cache_num() == 0 + assert big_pages.alloc_one_state_cache() is None + for slot in [manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID, manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID]: + with pytest.raises(AssertionError): + big_pages.free_state_cache([slot]) From adb4155d31b1c2170bc897b8f90f197de7d5f026 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Mon, 7 Sep 2026 08:41:36 +0000 Subject: [PATCH 08/14] refactor: simplify hybrid cache state interfaces --- .../hybrid_sliding_mem_manager.py | 63 +++++------ .../operator/hybrid_sliding.py | 4 + .../operator/linear_att.py | 2 +- lightllm/common/req_manager/hybrid_att.py | 10 +- lightllm/common/req_manager/linear_att.py | 24 +--- .../layer_infer/transformer_layer_infer.py | 6 +- .../server/router/model_infer/infer_batch.py | 2 +- lightllm/utils/config_utils.py | 19 +--- test/kernel/test_sliding_window_state.py | 19 ++++ test/utils/test_sliding_cpu_cache_meta.py | 18 +++ test/utils/test_sliding_window_cache.py | 104 +++++++++++++----- .../model_infer/test_hybrid_state_cache.py | 68 ++++++++++++ 12 files changed, 227 insertions(+), 112 deletions(-) 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 index fbaa51db2f..0e15b5d329 100644 --- a/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -35,63 +35,55 @@ def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): mem_fraction=mem_fraction, ) - def _big_page_num(self, token_num): - return max(1, triton.cdiv(token_num, self.big_page_token_num)) if self.enable_prompt_cache else 0 - - def _cache_nbytes(self, token_num): - # Runtime windows already exist when profiling. Reserve BOTH GPU page - # pools here, plus the full-KV hold token, final partial big page, and - # separate CPU-cache load/offload staging states when enabled. - return (token_num + 1) * self.get_cell_size() + ( - self.small_page_num + self._big_page_num(token_num) + self.cpu_cache_temp_page_num - ) * self.sliding_config.get_state_nbytes() - - def _profile_token_num(self, available_bytes): - if self._cache_nbytes(1) > available_bytes: - raise ValueError( - "Insufficient GPU memory for sliding-window checkpoints and full KV: " - f"{available_bytes / 1024 ** 3:.2f} GiB available, " - f"{self.small_page_num} small pages at " - f"{self.sliding_config.get_state_nbytes() / 1024 ** 2:.2f} MiB/page. " - "Reduce --linear_att_cache_size or --running_max_req_size." - ) - low, high = 1, available_bytes // self.get_cell_size() - while low < high: - mid = (low + high + 1) // 2 - if self._cache_nbytes(mid) <= available_bytes: - low = mid - else: - high = mid - 1 - return low - def profile_size(self, mem_fraction): torch.cuda.empty_cache() world_size = dist.get_world_size() available_memory = get_available_gpu_memory(world_size) if self.size is None: available_memory -= get_total_gpu_memory() * (1 - mem_fraction) - self.size = self._profile_token_num(int(available_memory * 1024 ** 3)) + available_bytes = int(available_memory * 1024 ** 3) + cell_size = self.get_cell_size() + state_bytes = self.sliding_config.get_state_nbytes() + # Runtime windows already exist. Reserve the hold token, small pages + # and CPU-transfer slots before sizing full KV and big checkpoints. + fixed_bytes = cell_size + (self.small_page_num + self.cpu_cache_temp_page_num) * state_bytes + big_page_state_bytes = state_bytes if self.enable_prompt_cache else 0 + if self.size is None: + if available_bytes < fixed_bytes + cell_size + big_page_state_bytes: + raise ValueError( + "Insufficient GPU memory for sliding-window checkpoints and full KV; " + "reduce --linear_att_cache_size or --running_max_req_size." + ) + # Each complete page costs B full-KV tokens plus one checkpoint. + # A partial page also needs one checkpoint before it can hold tokens. + page_bytes = self.big_page_token_num * cell_size + big_page_state_bytes + page_num, tail_bytes = divmod(available_bytes - fixed_bytes, page_bytes) + self.size = page_num * self.big_page_token_num + max(0, (tail_bytes - big_page_state_bytes) // cell_size) if world_size > 1: size_tensor = torch.tensor(self.size, dtype=torch.int64, device="cuda") dist.all_reduce(size_tensor, op=dist.ReduceOp.MIN) self.size = size_tensor.item() - elif self._cache_nbytes(self.size) > int(available_memory * 1024 ** 3): + + big_page_num = triton.cdiv(self.size, self.big_page_token_num) if self.enable_prompt_cache else 0 + cache_bytes = fixed_bytes + self.size * cell_size + big_page_num * state_bytes + if cache_bytes > available_bytes: raise ValueError( "Requested full KV and sliding-window checkpoints exceed available GPU memory; " "reduce --max_total_token_num, --linear_att_cache_size or --running_max_req_size." ) logger.info( f"Sliding-window cache budget: {self.size} full-KV tokens, " - f"{self._big_page_num(self.size)} big pages, {self.small_page_num} small pages, " + f"{big_page_num} big pages, {self.small_page_num} small pages, " f"{self.cpu_cache_temp_page_num} CPU-cache staging states, " - f"{self._cache_nbytes(self.size) / 1024 ** 3:.2f} GiB (runtime windows already allocated)" + f"{cache_bytes / 1024 ** 3:.2f} GiB (runtime windows already allocated)" ) def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): super()._init_buffers(size, dtype, head_num, head_dim, layer_num) + big_page_num = triton.cdiv(size, self.big_page_token_num) if self.enable_prompt_cache else 0 # Keep the existing radix-cache contract; no second alias is needed. self.linear_att_big_page_buffers = SlidingWindowStateCacheManager( - size=self._big_page_num(size) + self.cpu_cache_temp_page_num, + size=big_page_num + self.cpu_cache_temp_page_num, sliding_config=self.sliding_config, keep_num=self.cpu_cache_temp_page_num, ) @@ -106,9 +98,6 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): def get_att_input_params(self, layer_index: int): return super().get_att_input_params(self.sliding_config.get_full_layer_index(layer_index)) - def get_full_cache_layer_index(self, layer_index: int): - return self.sliding_config.get_full_layer_index(layer_index) - def _free_buffers(self): super()._free_buffers() self.linear_att_big_page_buffers = None diff --git a/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py index 07f35cd1c8..da27b7efec 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py @@ -10,6 +10,10 @@ 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.get_full_layer_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, diff --git a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py index 49e2549265..497c042314 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py +++ b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py @@ -92,7 +92,7 @@ def load_cpu_cache_to_gpu( from lightllm.server.router.model_infer.infer_batch import g_infer_context - g_infer_context.req_manager.copy_big_page_buffer_to_linear_att_state( + g_infer_context.req_manager.restore_big_page_state( big_page_buffer_idx=big_page_buffer_ids_cpu[-1], req=req, ) diff --git a/lightllm/common/req_manager/hybrid_att.py b/lightllm/common/req_manager/hybrid_att.py index 49cc5b6b8c..8f37ee07df 100644 --- a/lightllm/common/req_manager/hybrid_att.py +++ b/lightllm/common/req_manager/hybrid_att.py @@ -11,11 +11,13 @@ class HybridAttentionReqManager(ReqManager, ABC): - """Request manager contract for token/full + request-state attention models. + """混合 attention 的请求运行态与大小页 checkpoint 管理接口。 - The token index table remains the virtual, token-granular address space used - by prefix-cache matching. The non-full attention state is managed through - this interface and may have a different physical granularity. + 大小页沿同一虚拟 token 索引空间匹配前缀，full attention KV 保持 token 粒度存储。 + linear/sliding-window 状态在大页边界及请求可缓存尾部的小页边界保存 checkpoint， + 缓存命中后，再将相应 checkpoint 恢复到请求运行态。 + + 公共缓存流程负责大小页分配、边界、匹配与淘汰；各实现负责状态存储和保存/恢复。 """ @abstractmethod diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index f21a2ec1bd..e6beeea461 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -53,18 +53,6 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_con def create_state_cache_manager(self, size: int): return LinearAttCacheManager(size=size, linear_config=self.linear_config) - def init_hybrid_attention_state(self, req: "InferReq"): - return self.init_linear_att_state(req) - - def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - return self.copy_big_page_buffer_to_linear_att_state(big_page_buffer_idx=big_page_buffer_idx, req=req) - - def restore_small_page_state(self, req: "InferReq", small_page_buffers): - return self.copy_small_page_buffer_to_linear_att_state( - req=req, - linear_att_small_page_buffers=small_page_buffers, - ) - def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): assert len(b_req_idx) == len(buffer_indexes) if not any(buffer_idx != -1 for buffer_idx in buffer_indexes): @@ -94,7 +82,7 @@ def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffer dst_conv_state.copy_(gpu_conv_state, non_blocking=True) dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) - def init_linear_att_state(self, req: "InferReq"): + def init_hybrid_attention_state(self, req: "InferReq"): conv_index = req.req_idx ssm_start = req.req_idx * (self.mtp_step + 1) self.req_to_conv_state.buffer[:, conv_index, ...].fill_(0) @@ -114,7 +102,7 @@ def get_mamba_cache(self, layer_idx_in_all: int): ssm_states = self.req_to_ssm_state.buffer[layer_idx_in_linear] return conv_states, ssm_states - def copy_big_page_buffer_to_linear_att_state(self, big_page_buffer_idx: int, req: "InferReq"): + def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): big_page_buffers: LinearAttCacheManager = self.mem_manager.linear_att_big_page_buffers conv_state, ssm_state = big_page_buffers.get_state_cache(buffer_idx=big_page_buffer_idx) @@ -127,12 +115,8 @@ def copy_big_page_buffer_to_linear_att_state(self, big_page_buffer_idx: int, req self.req_to_mtp_state_index[req.req_idx] = 0 return - def copy_small_page_buffer_to_linear_att_state( - self, req: "InferReq", linear_att_small_page_buffers: LinearAttCacheManager - ): - conv_state, ssm_state = linear_att_small_page_buffers.get_state_cache( - buffer_idx=req.shared_kv_node.small_page_buffer_idx - ) + def restore_small_page_state(self, req: "InferReq", small_page_buffers: LinearAttCacheManager): + conv_state, ssm_state = small_page_buffers.get_state_cache(buffer_idx=req.shared_kv_node.small_page_buffer_idx) conv_dest = req.req_idx ssm_dest = req.req_idx * (self.mtp_step + 1) conv_cache_width = conv_state.shape[-1] diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index c8a02b1a42..3206be89a8 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -158,11 +158,7 @@ def _post_cache_kv(self, cache_kv, infer_state, layer_weight): infer_state.req_manager.req_to_sliding_window[layer_idx], ) return - infer_state.mem_manager.operator.copy_kv_to_mem_manager( - layer_index=infer_state.mem_manager.get_full_cache_layer_index(self.layer_num_), - mem_index=infer_state.mem_index, - kv=cache_kv, - ) + super()._post_cache_kv(cache_kv, infer_state, layer_weight) # ----- Attention kernels (sliding window + per-layer KV reshape) --- diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index e126ead141..c7f535c2e9 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -742,7 +742,7 @@ def _linear_match_radix_cache(self): destination_indexes=tail_mems, ) - self.shared_kv_node = share_node # 只是为了保证 copy_small_page_buffer_to_linear_att_state 正确调用 + self.shared_kv_node = share_node # 只是为了保证 restore_small_page_state 正确调用 g_infer_context.req_manager.restore_small_page_state( req=self, small_page_buffers=g_infer_context.radix_cache.linear_att_small_page_buffers, diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 4dfabe66e9..961cead912 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -467,24 +467,7 @@ def is_linear_att_mixed_model(model_path: str) -> bool: @lru_cache(maxsize=None) def is_sliding_att_mixed_model(model_path: str) -> bool: - try: - config_json = get_config_json(model_path) - llm_config = config_json.get("text_config", config_json) - model_type = config_json.get("model_type") or llm_config.get("model_type") - layer_types = set(llm_config.get("layer_types", [])) - # Keep the shared-request ABI opt-in aligned with models that actually - # instantiate ReqManagerForSlidingWindow. Other architectures may use - # the same layer-type strings while retaining token-granular KV. - return ( - model_type in {"gemma4", "gemma4_text"} - and { - "full_attention", - "sliding_attention", - }.issubset(layer_types) - ) - except Exception: - logger.info(f"model path: {model_path} does not have hybrid sliding-window attention") - return False + return get_model_type(model_path) in {"gemma4", "gemma4_text"} def is_hybrid_att_mixed_model(model_path: str) -> bool: diff --git a/test/kernel/test_sliding_window_state.py b/test/kernel/test_sliding_window_state.py index 1f8fdc70b6..4a8ab7599a 100644 --- a/test/kernel/test_sliding_window_state.py +++ b/test/kernel/test_sliding_window_state.py @@ -8,6 +8,7 @@ prepare_sliding_window_indexes, ) from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow +from lightllm.common.kv_cache_mem_manager.operator.hybrid_sliding import HybridSlidingMemOperator from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig, SlidingWindowStateCacheManager from lightllm.models.gemma4.kv_layout import get_kv_cache_layout from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer @@ -16,6 +17,24 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("layer_index,is_shared", [(5, False), (11, False), (17, True)]) +def test_full_kv_write_maps_logical_layer_once_and_skips_shared_readers(layer_index, is_shared): + config = SlidingWindowCacheConfig({0: 0}, {5: 0, 11: 1, 17: 1}, 32, 1, 64, 1, 64, torch.bfloat16) + mem_manager = SimpleNamespace( + sliding_config=config, kv_buffer=torch.zeros((2, 8, 2, 64), dtype=torch.bfloat16, device="cuda") + ) + mem_manager.operator = HybridSlidingMemOperator(mem_manager) + layer = object.__new__(Gemma4TransformerLayerInfer) + layer.layer_num_, layer.is_sliding, layer.is_kv_shared_ = layer_index, False, is_shared + indexes = torch.tensor([1, 3], dtype=torch.int32, device="cuda") + kv = torch.randn((2, 2, 64), dtype=torch.bfloat16, device="cuda") + layer._post_cache_kv(kv, SimpleNamespace(mem_manager=mem_manager, mem_index=indexes), None) + expected = torch.zeros_like(mem_manager.kv_buffer) + if not is_shared: + expected[config.get_full_layer_index(layer_index), indexes] = kv + torch.testing.assert_close(mem_manager.kv_buffer, expected, atol=0, rtol=0) + + @pytest.mark.parametrize("history_len", [0, 511, 512, 513, 1024]) @pytest.mark.parametrize("q_len", [1, 31, 256, 768]) def test_ring_attention_and_commit_match_token_cache(history_len, q_len): diff --git a/test/utils/test_sliding_cpu_cache_meta.py b/test/utils/test_sliding_cpu_cache_meta.py index 5ec84b9166..8f98fa5e3f 100644 --- a/test/utils/test_sliding_cpu_cache_meta.py +++ b/test/utils/test_sliding_cpu_cache_meta.py @@ -8,6 +8,24 @@ from lightllm.models.gemma4.kv_layout import build_sliding_cache_config +@pytest.mark.parametrize( + "model_config,expected", + [ + ({"model_type": "gemma4"}, True), + ({"model_type": "gemma4_text"}, True), + ({"text_config": {"model_type": "gemma4_text"}}, True), + ({"model_type": "gemma4", "layer_types": ["full_attention"]}, True), + ({"model_type": "gemma3", "layer_types": ["sliding_attention", "full_attention"]}, False), + ({"model_type": "qwen3_5"}, False), + ], +) +def test_sliding_cache_architecture_is_selected_by_model_type(monkeypatch, model_config, expected): + import lightllm.utils.config_utils as config_utils + + monkeypatch.setattr(config_utils, "get_config_json", lambda _: model_config) + assert config_utils.is_sliding_att_mixed_model.__wrapped__("test-model") is expected + + def _gemma_config(shared): layer_num = 42 if shared else 60 return { diff --git a/test/utils/test_sliding_window_cache.py b/test/utils/test_sliding_window_cache.py index 84cb494f18..fc066bb3a4 100644 --- a/test/utils/test_sliding_window_cache.py +++ b/test/utils/test_sliding_window_cache.py @@ -26,6 +26,7 @@ def test_gemma_physical_owners_and_last_readers(layer_num, shared, sliding_num, def _memory_manager(big_page_tokens=2048, small_pages=8, enabled=True, cpu_cache=False): manager = object.__new__(HybridSlidingMemoryManager) + manager.size = None manager.head_num, manager.head_dim, manager.layer_num, manager.dtype = 1, 512, 10, torch.bfloat16 manager.sliding_config = SlidingWindowCacheConfig( {i: i for i in range(50)}, {50 + i: i for i in range(10)}, 1024, 4, 256, 1, 512, torch.bfloat16 @@ -39,34 +40,89 @@ def _memory_manager(big_page_tokens=2048, small_pages=8, enabled=True, cpu_cache return manager +def _required_bytes(manager, token_num): + big_pages = (token_num + manager.big_page_token_num - 1) // manager.big_page_token_num + state_pages = manager.small_page_num + manager.cpu_cache_temp_page_num + if manager.enable_prompt_cache: + state_pages += big_pages + return (token_num + 1) * manager.get_cell_size() + state_pages * manager.sliding_config.get_state_nbytes() + + +def _profile_with_budget(monkeypatch, manager, available_bytes, mem_fraction=1.0, total_bytes=16 * 1024 ** 3): + import lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager as memory_module + + monkeypatch.setattr(memory_module.dist, "get_world_size", lambda: 1) + monkeypatch.setattr(memory_module, "get_available_gpu_memory", lambda world_size: available_bytes / 1024 ** 3) + monkeypatch.setattr(memory_module, "get_total_gpu_memory", lambda: total_bytes / 1024 ** 3) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + manager.profile_size(mem_fraction) + + @pytest.mark.parametrize("token_num", [1, 2047, 2048, 2049, 8192]) -def test_profile_accounts_for_small_big_partial_page_and_hold_token(token_num): +def test_profile_accounts_for_small_big_partial_page_and_hold_token(monkeypatch, token_num): manager = _memory_manager() assert manager.sliding_config.get_state_nbytes() == 200 * 1024 ** 2 expected = (token_num + 1) * 20480 + (8 + (token_num + 2047) // 2048) * 200 * 1024 ** 2 - assert manager._cache_nbytes(token_num) == expected - assert manager._profile_token_num(expected) == token_num + assert _required_bytes(manager, token_num) == expected + _profile_with_budget(monkeypatch, manager, expected) + assert manager.size == token_num + manager.size = None if token_num > 1: - assert manager._profile_token_num(expected - 1) < token_num + _profile_with_budget(monkeypatch, manager, expected - 1) + assert manager.size == token_num - 1 + else: + with pytest.raises(ValueError, match="Insufficient GPU memory"): + _profile_with_budget(monkeypatch, manager, expected - 1) + + +@pytest.mark.parametrize("leftover", [1, 200 * 1024 ** 2 - 1, 200 * 1024 ** 2]) +def test_profile_cannot_start_next_page_without_checkpoint_and_token_budget(monkeypatch, leftover): + manager = _memory_manager() + _profile_with_budget(monkeypatch, manager, _required_bytes(manager, 2048) + leftover) + assert manager.size == 2048 -def test_profile_reports_impossible_checkpoint_budget(): +@pytest.mark.parametrize("available_bytes", [-1, 0, 80 * 1024 ** 3]) +def test_profile_reports_impossible_checkpoint_budget(monkeypatch, available_bytes): manager = _memory_manager(small_pages=512) with pytest.raises(ValueError, match="linear_att_cache_size"): - manager._profile_token_num(80 * 1024 ** 3) + _profile_with_budget(monkeypatch, manager, available_bytes) -def test_disabled_prompt_cache_does_not_reserve_pages(): +def test_disabled_prompt_cache_does_not_reserve_pages(monkeypatch): manager = _memory_manager(small_pages=0, enabled=False) - assert manager._cache_nbytes(4096) == 4097 * manager.get_cell_size() + _profile_with_budget(monkeypatch, manager, 4097 * manager.get_cell_size()) + assert manager.size == 4096 -def test_cpu_cache_reserves_two_additional_window_checkpoints(): +def test_cpu_cache_reserves_two_additional_window_checkpoints(monkeypatch): gpu_only = _memory_manager() cpu_cache = _memory_manager(cpu_cache=True) - assert cpu_cache._cache_nbytes(4096) == ( - gpu_only._cache_nbytes(4096) + 2 * cpu_cache.sliding_config.get_state_nbytes() - ) + gpu_budget = _required_bytes(gpu_only, 4096) + cpu_budget = gpu_budget + 2 * cpu_cache.sliding_config.get_state_nbytes() + _profile_with_budget(monkeypatch, gpu_only, gpu_budget) + _profile_with_budget(monkeypatch, cpu_cache, cpu_budget) + assert gpu_only.size == cpu_cache.size == 4096 + + +def test_explicit_size_is_checked_against_complete_cache_budget(monkeypatch): + manager = _memory_manager(cpu_cache=True) + manager.size = 2049 + budget = _required_bytes(manager, manager.size) + _profile_with_budget(monkeypatch, manager, budget) + assert manager.size == 2049 + with pytest.raises(ValueError, match="exceed available GPU memory"): + _profile_with_budget(monkeypatch, manager, budget - 1) + + +@pytest.mark.parametrize("explicit_size", [None, 2049]) +def test_mem_fraction_reserves_headroom_only_for_automatic_size(monkeypatch, explicit_size): + manager = _memory_manager() + manager.size = explicit_size + total_bytes, mem_fraction = 8 * 1024 ** 3, 0.5 + budget = _required_bytes(manager, 2048) + total_bytes // 2 + _profile_with_budget(monkeypatch, manager, budget, mem_fraction=mem_fraction, total_bytes=total_bytes) + assert manager.size == (2048 if explicit_size is None else explicit_size) @pytest.mark.parametrize("disabled", [False, True]) @@ -87,26 +143,25 @@ def test_page_pools_follow_active_prompt_cache_flag(monkeypatch, disabled): manager = HybridSlidingMemoryManager(size=256, sliding_config=config) assert manager.enable_prompt_cache is not disabled assert manager.small_page_num == (0 if disabled else 3) - assert manager._big_page_num(256) == (0 if disabled else 1) + assert manager.big_page_token_num == 256 + assert manager.cpu_cache_temp_page_num == 0 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("cpu_cache", [False, True]) -def test_profiled_gpu_pools_match_reserved_bytes_and_are_reused(monkeypatch, cpu_cache): - import lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager as memory_module +@pytest.mark.parametrize("enabled,cpu_cache", [(False, False), (True, False), (True, True)]) +def test_profiled_gpu_pools_match_reserved_bytes_and_are_reused(monkeypatch, enabled, cpu_cache): from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow - manager = _memory_manager(big_page_tokens=32, small_pages=2, cpu_cache=cpu_cache) + manager = _memory_manager(big_page_tokens=32, small_pages=2 if enabled else 0, enabled=enabled, cpu_cache=cpu_cache) manager.head_num, manager.head_dim, manager.layer_num = 1, 64, 1 manager.sliding_config = SlidingWindowCacheConfig({0: 0}, {1: 0}, 32, 1, 64, 1, 64, torch.bfloat16) - manager.size = None - budget = manager._cache_nbytes(65) - monkeypatch.setattr(memory_module.dist, "get_world_size", lambda: 1) - monkeypatch.setattr(memory_module, "get_available_gpu_memory", lambda world_size: budget / 1024 ** 3) - monkeypatch.setattr(memory_module, "get_total_gpu_memory", lambda: 1) - manager.profile_size(1.0) + budget = _required_bytes(manager, 65) + _profile_with_budget(monkeypatch, manager, budget) assert manager.size == 65 manager._init_buffers(manager.size, manager.dtype, manager.head_num, manager.head_dim, manager.layer_num) + assert manager.linear_att_big_page_buffers.size == (3 if enabled else 0) + (2 if cpu_cache else 0) + assert manager.linear_att_big_page_buffers.get_free_cache_num() == (3 if enabled else 0) + assert manager.sliding_small_page_buffers.size == (2 if enabled else 0) allocated = sum( t.numel() * t.element_size() for t in [ @@ -119,9 +174,6 @@ def test_profiled_gpu_pools_match_reserved_bytes_and_are_reused(monkeypatch, cpu req_manager = object.__new__(ReqManagerForSlidingWindow) req_manager.mem_manager = manager assert req_manager.create_state_cache_manager(2) is manager.sliding_small_page_buffers - manager.size = 1000000 - with pytest.raises(ValueError, match="exceed available GPU memory"): - manager.profile_size(1.0) @pytest.mark.parametrize("mtp_step", [0, 2]) diff --git a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py index 471df7ae09..0e4267ef5f 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py @@ -3,11 +3,79 @@ import pytest import torch +from lightllm.common.req_manager.linear_att import ReqManagerForMamba from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager from lightllm.server.router.model_infer.infer_batch import InferenceContext +def _cpu_linear_req_manager(mtp_step): + manager = object.__new__(ReqManagerForMamba) + manager.mtp_step = mtp_step + manager.req_to_conv_state = SimpleNamespace( + buffer=torch.arange(2 * 3 * 3 * (4 + mtp_step), dtype=torch.float32).reshape(2, 3, 3, 4 + mtp_step) + ) + manager.req_to_ssm_state = SimpleNamespace( + buffer=torch.arange(2 * 3 * (mtp_step + 1) * 5, dtype=torch.float32).reshape(2, 3 * (mtp_step + 1), 5) + ) + manager.req_to_mtp_state_index = torch.full((3,), mtp_step, dtype=torch.int32) if mtp_step else None + return manager + + +@pytest.mark.parametrize("mtp_step", [0, 2]) +def test_linear_init_clears_entire_request_state_and_resets_mtp_index(mtp_step): + manager = _cpu_linear_req_manager(mtp_step) + req = SimpleNamespace(req_idx=1) + conv = manager.req_to_conv_state.buffer + ssm = manager.req_to_ssm_state.buffer + ssm_start = req.req_idx * (mtp_step + 1) + conv[:, req.req_idx] = float("nan") + ssm[:, ssm_start : ssm_start + mtp_step + 1] = float("nan") + expected_conv, expected_ssm = conv.clone(), ssm.clone() + expected_conv[:, req.req_idx].zero_() + expected_ssm[:, ssm_start : ssm_start + mtp_step + 1].zero_() + + manager.init_hybrid_attention_state(req=req) + + torch.testing.assert_close(conv, expected_conv, atol=0, rtol=0) + torch.testing.assert_close(ssm, expected_ssm, atol=0, rtol=0) + if mtp_step: + torch.testing.assert_close( + manager.req_to_mtp_state_index, torch.tensor([mtp_step, 0, mtp_step], dtype=torch.int32) + ) + else: + assert manager.req_to_mtp_state_index is None + + +@pytest.mark.parametrize("mtp_step", [0, 2]) +@pytest.mark.parametrize("page_kind", ["big", "small"]) +def test_linear_restore_preserves_mtp_conv_tail_and_noncanonical_ssm_rows(mtp_step, page_kind): + manager = _cpu_linear_req_manager(mtp_step) + req = SimpleNamespace(req_idx=1, shared_kv_node=SimpleNamespace(small_page_buffer_idx=2)) + conv_pages = torch.arange(3 * 2 * 3 * 4, dtype=torch.float32).reshape(3, 2, 3, 4) + 1000 + ssm_pages = torch.arange(3 * 2 * 5, dtype=torch.float32).reshape(3, 2, 5) + 2000 + pages = SimpleNamespace(get_state_cache=lambda buffer_idx: (conv_pages[buffer_idx], ssm_pages[buffer_idx])) + manager.mem_manager = SimpleNamespace(linear_att_big_page_buffers=pages) + expected_conv = manager.req_to_conv_state.buffer.clone() + expected_ssm = manager.req_to_ssm_state.buffer.clone() + expected_conv[:, req.req_idx, ..., :4] = conv_pages[2] + expected_ssm[:, req.req_idx * (mtp_step + 1)] = ssm_pages[2] + + if page_kind == "big": + manager.restore_big_page_state(big_page_buffer_idx=2, req=req) + else: + manager.restore_small_page_state(req=req, small_page_buffers=pages) + + torch.testing.assert_close(manager.req_to_conv_state.buffer, expected_conv, atol=0, rtol=0) + torch.testing.assert_close(manager.req_to_ssm_state.buffer, expected_ssm, atol=0, rtol=0) + if mtp_step: + torch.testing.assert_close( + manager.req_to_mtp_state_index, torch.tensor([mtp_step, 0, mtp_step], dtype=torch.int32) + ) + else: + assert manager.req_to_mtp_state_index is None + + @pytest.mark.parametrize("is_hybrid,radix_cache", [(False, object()), (True, None)]) def test_snapshot_without_hybrid_cache_returns_before_reading_requests(is_hybrid, radix_cache): context = InferenceContext(is_hybrid_att_mixed_model=is_hybrid, radix_cache=radix_cache) From 01016acd06134a041da8b432842618d8cbc59ae0 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Mon, 7 Sep 2026 09:16:31 +0000 Subject: [PATCH 09/14] refactor: address sliding window state without token tables --- .../triton_kernel/sliding_window_state.py | 82 ----- lightllm/common/req_manager/sliding_window.py | 49 +-- lightllm/models/gemma4/infer_struct.py | 16 - .../layer_infer/transformer_layer_infer.py | 26 +- .../context_attention_fwd_gemma4_mm.py | 37 +- .../triton_kernel/sliding_window_decode.py | 173 +++++++++ .../benchmark_sliding_window_attention.py | 348 ++++++++++++++++++ test/kernel/test_gemma4_hybrid_graph.py | 126 +++++++ ...test_sliding_window_cpu_cache_attention.py | 15 +- test/kernel/test_sliding_window_decode.py | 151 ++++++++ test/kernel/test_sliding_window_prefill.py | 162 ++++++++ test/kernel/test_sliding_window_state.py | 78 ++-- .../model_infer/test_hybrid_state_cache.py | 4 +- .../test_sliding_cpu_cache_loading.py | 6 +- .../test_sliding_cpu_cache_operator.py | 10 +- 15 files changed, 1101 insertions(+), 182 deletions(-) create mode 100644 lightllm/models/gemma4/triton_kernel/sliding_window_decode.py create mode 100644 test/kernel/benchmark_sliding_window_attention.py create mode 100644 test/kernel/test_gemma4_hybrid_graph.py create mode 100644 test/kernel/test_sliding_window_decode.py create mode 100644 test/kernel/test_sliding_window_prefill.py diff --git a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py index bd940950d7..9df3551a63 100644 --- a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py @@ -3,88 +3,6 @@ import triton.language as tl -@triton.jit -def _prepare_history_indexes( - ReqToSliding, - BReqIdx, - BSeqLen, - BQSeqLen, - stride_req, - stride_seq, - WINDOW: tl.constexpr, - BLOCK: tl.constexpr, -): - batch_idx = tl.program_id(0) - block_idx = tl.program_id(1) - req_idx = tl.load(BReqIdx + batch_idx) - seq_len = tl.load(BSeqLen + batch_idx) - q_len = tl.load(BQSeqLen + batch_idx) - history_end = seq_len - q_len - history_start = tl.maximum(0, history_end - WINDOW) - pos = history_start + block_idx * BLOCK + tl.arange(0, BLOCK) - mask = pos < history_end - physical = req_idx * WINDOW + pos % WINDOW - tl.store(ReqToSliding + req_idx * stride_req + pos * stride_seq, physical, mask=mask) - - -@triton.jit -def _prepare_current_indexes( - ReqToSliding, - BReqIdx, - BSeqLen, - BQSeqLen, - BQStartLoc, - stride_req, - stride_seq, - SCRATCH_START: tl.constexpr, - BLOCK: tl.constexpr, -): - batch_idx = tl.program_id(0) - block_idx = tl.program_id(1) - req_idx = tl.load(BReqIdx + batch_idx) - seq_len = tl.load(BSeqLen + batch_idx) - q_len = tl.load(BQSeqLen + batch_idx) - q_start = tl.load(BQStartLoc + batch_idx) - offset = block_idx * BLOCK + tl.arange(0, BLOCK) - mask = offset < q_len - pos = seq_len - q_len + offset - physical = SCRATCH_START + q_start + offset - tl.store(ReqToSliding + req_idx * stride_req + pos * stride_seq, physical, mask=mask) - - -@torch.no_grad() -def prepare_sliding_window_indexes( - req_to_sliding_window_indexs, - b_req_idx, - b_seq_len, - b_q_seq_len, - b_q_start_loc, - sliding_window, - scratch_start, - max_q_seq_len, -): - block = 256 - _prepare_history_indexes[(b_req_idx.shape[0], triton.cdiv(sliding_window, block))]( - req_to_sliding_window_indexs, - b_req_idx, - b_seq_len, - b_q_seq_len, - *req_to_sliding_window_indexs.stride(), - WINDOW=sliding_window, - BLOCK=block, - ) - _prepare_current_indexes[(b_req_idx.shape[0], triton.cdiv(max_q_seq_len, block))]( - req_to_sliding_window_indexs, - b_req_idx, - b_seq_len, - b_q_seq_len, - b_q_start_loc, - *req_to_sliding_window_indexs.stride(), - SCRATCH_START=scratch_start, - BLOCK=block, - ) - - @triton.jit def _commit_sliding_window_state( LayerBuffer, diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index 346c83ed07..541434d302 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -2,10 +2,7 @@ import torch -from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( - commit_sliding_window_state, - prepare_sliding_window_indexes, -) +from lightllm.common.basemodel.triton_kernel.sliding_window_state import commit_sliding_window_state from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager from .hybrid_att import HybridAttentionReqManager @@ -18,7 +15,7 @@ class ReqManagerForSlidingWindow(HybridAttentionReqManager): - """Token-granular virtual addresses plus request-granular sliding KV.""" + """按请求保存窗口运行态，并在大小页边界保存和恢复 checkpoint。""" def __init__( self, @@ -32,22 +29,25 @@ def __init__( self.sliding_config = sliding_config self.sliding_window = sliding_config.sliding_window self.scratch_token_num = scratch_token_num - self.runtime_token_num = (max_request_num + 1) * self.sliding_window - self.scratch_start = self.runtime_token_num - self.req_to_sliding_window = torch.zeros( + self.scratch_start = (max_request_num + 1) * self.sliding_window + # Attention reads history and current-chunk KV from one buffer. The + # request state is a view of its ring region, not a second allocation. + self.sliding_kv_buffer = torch.zeros( ( sliding_config.sliding_layer_num, - self.runtime_token_num + scratch_token_num, + self.scratch_start + scratch_token_num, 2 * sliding_config.sliding_head_num, sliding_config.sliding_head_dim, ), dtype=sliding_config.dtype, device="cuda", ) - self.req_to_sliding_window_indexs = torch.zeros( - (max_request_num + 1, max_sequence_length), - dtype=torch.int32, - device="cuda", + self.req_to_sliding_window = self.sliding_kv_buffer[:, : self.scratch_start].view( + sliding_config.sliding_layer_num, + max_request_num + 1, + self.sliding_window, + 2 * sliding_config.sliding_head_num, + sliding_config.sliding_head_dim, ) def create_state_cache_manager(self, size: int): @@ -55,8 +55,7 @@ def create_state_cache_manager(self, size: int): return self.mem_manager.sliding_small_page_buffers def init_hybrid_attention_state(self, req: "InferReq"): - start = req.req_idx * self.sliding_window - self.req_to_sliding_window[:, start : start + self.sliding_window].zero_() + self.req_to_sliding_window[:, req.req_idx].zero_() def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): self._restore_state(req.req_idx, self.mem_manager.linear_att_big_page_buffers, big_page_buffer_idx) @@ -65,8 +64,7 @@ def restore_small_page_state(self, req: "InferReq", small_page_buffers): self._restore_state(req.req_idx, small_page_buffers, req.shared_kv_node.small_page_buffer_idx) def _restore_state(self, req_idx: int, state_cache_manager, buffer_idx: int): - start = req_idx * self.sliding_window - self.req_to_sliding_window[:, start : start + self.sliding_window].copy_( + self.req_to_sliding_window[:, req_idx].copy_( state_cache_manager.get_state_cache(buffer_idx), non_blocking=True, ) @@ -79,25 +77,14 @@ def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], self.save_small_page_state(req_idx, buffer_idx, self.mem_manager.linear_att_big_page_buffers) def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffers: SlidingWindowStateCacheManager): - start = req_idx * self.sliding_window small_page_buffers.get_state_cache(buffer_idx).copy_( - self.req_to_sliding_window[:, start : start + self.sliding_window], + self.req_to_sliding_window[:, req_idx], non_blocking=True, ) def prepare_sliding_window(self, infer_state): q_token_num = infer_state.input_ids.shape[0] assert q_token_num <= self.scratch_token_num - prepare_sliding_window_indexes( - req_to_sliding_window_indexs=self.req_to_sliding_window_indexs, - b_req_idx=infer_state.b_req_idx, - b_seq_len=infer_state.b_seq_len, - b_q_seq_len=infer_state.b_q_seq_len, - b_q_start_loc=infer_state.b_q_start_loc, - sliding_window=self.sliding_window, - scratch_start=self.scratch_start, - max_q_seq_len=infer_state.max_q_seq_len, - ) infer_state.sliding_window_mem_index = torch.arange( self.scratch_start, self.scratch_start + q_token_num, @@ -107,14 +94,14 @@ def prepare_sliding_window(self, infer_state): def get_layer_kv(self, layer_index: int): local_layer = self.sliding_config.get_sliding_layer_index(layer_index) - layer_buffer = self.req_to_sliding_window[local_layer] + layer_buffer = self.sliding_kv_buffer[local_layer] head_num = self.sliding_config.sliding_head_num return layer_buffer[:, :head_num], layer_buffer[:, head_num:] def commit_layer_state(self, layer_index: int, infer_state): local_layer = self.sliding_config.get_sliding_layer_index(layer_index) commit_sliding_window_state( - layer_buffer=self.req_to_sliding_window[local_layer], + layer_buffer=self.sliding_kv_buffer[local_layer], b_req_idx=infer_state.b_req_idx, b_seq_len=infer_state.b_seq_len, b_q_seq_len=infer_state.b_q_seq_len, diff --git a/lightllm/models/gemma4/infer_struct.py b/lightllm/models/gemma4/infer_struct.py index cfbe745b52..d18a470dc5 100644 --- a/lightllm/models/gemma4/infer_struct.py +++ b/lightllm/models/gemma4/infer_struct.py @@ -1,5 +1,3 @@ -import copy - import torch from lightllm.common.basemodel import InferStateInfo from lightllm.models.gemma4.triton_kernel.build_b_image_token_end import build_b_image_token_end @@ -48,20 +46,6 @@ def init_some_extra_state(self, model): self.req_manager.prepare_sliding_window(self) return - def init_att_state(self): - if not self.is_prefill: - # Keep the common attention path unchanged: its decode kernels read - # req_to_token_indexs from infer_state.req_manager. The sliding - # state receives a shallow model-side view whose table addresses - # the request-window KV buffer; the full-attention state continues - # to use this infer state and the virtual token table. - sliding_infer_state = copy.copy(self) - sliding_req_manager = copy.copy(self.req_manager) - sliding_req_manager.req_to_token_indexs = self.req_manager.req_to_sliding_window_indexs - sliding_infer_state.req_manager = sliding_req_manager - self.decode_att_state.infer_state = sliding_infer_state - return super().init_att_state() - 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) diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 3206be89a8..043087f453 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -155,7 +155,7 @@ def _post_cache_kv(self, cache_kv, infer_state, layer_weight): destindex_copy_kv( cache_kv, infer_state.sliding_window_mem_index, - infer_state.req_manager.req_to_sliding_window[layer_idx], + infer_state.req_manager.sliding_kv_buffer[layer_idx], ) return super()._post_cache_kv(cache_kv, infer_state, layer_weight) @@ -200,9 +200,10 @@ def _context_attention_kernel( 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_indexs, + None, infer_state.b_image_token_end, sliding_window=sw, + scratch_start=infer_state.req_manager.scratch_start, ) if self.commit_sliding_state_: infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) @@ -224,8 +225,25 @@ def _token_attention_kernel( ) -> torch.Tensor: _k, _v = self._get_layer_kv(infer_state) _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: + from lightllm.models.gemma4.triton_kernel.sliding_window_decode import sliding_window_decode_attention + + o_tensor = sliding_window_decode_attention( + q=_q, + k=_k, + v=_v, + b_req_idx=infer_state.b_req_idx, + b_seq_len=infer_state.b_seq_len, + b_q_start_loc=infer_state.b_q_start_loc, + sliding_window=self.sliding_window_, + scratch_start=infer_state.req_manager.scratch_start, + out=out, + alloc_tensor_func=self.alloc_tensor, + ) + else: + o_tensor = infer_state.decode_att_state1.decode_att( + q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor + ) if self.commit_sliding_state_: infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) return o_tensor.view(q.shape) diff --git a/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py b/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py index dee10e96d3..80dd201dfe 100644 --- a/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py +++ b/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py @@ -59,12 +59,14 @@ def _fwd_kernel( stride_req_to_tokens_s, kv_group_num, b_prompt_cache_len, + scratch_start, H: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, USE_SLIDING_WINDOW: tl.constexpr, SLIDING_WINDOW_LEFT: tl.constexpr, + USE_RING_CACHE: tl.constexpr, ): start_m = tl.program_id(0) cur_bh = tl.program_id(1) @@ -126,11 +128,23 @@ def _fwd_kernel( k_pos = kv_start_index + start_n + offs_n # [N] k_valid = k_pos < block_end_loc - kv_loc = tl.load( - Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * k_pos, - mask=k_valid, - other=0, - ).to(tl.int64) + if USE_RING_CACHE: + history_loc = cur_batch_req_idx.to(tl.int64) * (SLIDING_WINDOW_LEFT + 1) + k_pos.to(tl.int64) % ( + SLIDING_WINDOW_LEFT + 1 + ) + current_loc = ( + tl.cast(scratch_start, tl.int64) + + cur_batch_in_all_start_index.to(tl.int64) + + k_pos.to(tl.int64) + - prompt_cache_len.to(tl.int64) + ) + kv_loc = tl.where(k_pos < prompt_cache_len, history_loc, current_loc) + else: + kv_loc = tl.load( + Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * k_pos, + mask=k_valid, + other=0, + ).to(tl.int64) off_k = kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd k = tl.load(K + off_k, mask=k_valid[None, :], other=0.0) @@ -188,6 +202,7 @@ def context_attention_fwd_gemma4_mm( req_to_token_indexs, b_image_token_end, sliding_window=(-1, -1), + scratch_start=None, ): """Prefill attention with image bidirectional masking on sliding layers. @@ -198,6 +213,8 @@ def context_attention_fwd_gemma4_mm( position (in the flattened new-token layout), value is the image span's end index (in absolute request position) if the token is inside an image span, else 0. + scratch_start: When set, use request rings plus current-token scratch + for sliding KV; ``req_to_token_indexs`` is unused and may be None. """ BLOCK_M = 128 if not is_tesla() else 64 Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] @@ -225,6 +242,10 @@ def context_attention_fwd_gemma4_mm( assert int(sliding_window[1]) == 0, "sliding_window right must be 0" sliding_window_left = int(sliding_window[0]) + use_ring_cache = scratch_start is not None + if use_ring_cache: + assert use_sliding_window and sliding_window_left >= 0, "ring KV requires a finite sliding window" + _fwd_kernel[grid]( q, k, @@ -248,16 +269,18 @@ def context_attention_fwd_gemma4_mm( o.stride(0), o.stride(1), o.stride(2), - req_to_token_indexs.stride(0), - req_to_token_indexs.stride(1), + 0 if use_ring_cache else req_to_token_indexs.stride(0), + 0 if use_ring_cache else req_to_token_indexs.stride(1), kv_group_num=kv_group_num, b_prompt_cache_len=b_prompt_cache_len, + scratch_start=scratch_start if use_ring_cache else 0, H=head, BLOCK_DMODEL=Lk, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, USE_SLIDING_WINDOW=use_sliding_window, SLIDING_WINDOW_LEFT=sliding_window_left, + USE_RING_CACHE=use_ring_cache, num_warps=num_warps, num_stages=num_stages, ) diff --git a/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py b/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py new file mode 100644 index 0000000000..9fb30084c0 --- /dev/null +++ b/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py @@ -0,0 +1,173 @@ +"""Gemma sliding decode over a request ring and the current token's scratch KV.""" + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage2 import ( + flash_decode_stage2, +) + + +@triton.jit +def _sliding_window_decode_stage1( + Q, + K, + V, + BReqIdx, + BSeqLen, + BQStartLoc, + MidO, + MidLogSumExp, + sm_scale, + stride_qb, + stride_qh, + stride_qd, + stride_kt, + stride_kh, + stride_kd, + stride_vt, + stride_vh, + stride_vd, + stride_ob, + stride_oh, + stride_os, + stride_od, + stride_lb, + stride_lh, + stride_ls, + gqa_group_size, + WINDOW: tl.constexpr, + SCRATCH_START: tl.constexpr, + Q_HEAD_NUM: tl.constexpr, + BLOCK_SEQ: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, +): + batch_idx = tl.program_id(0) + kv_head = tl.program_id(1) + block_idx = tl.program_id(2) + grid_block_num = tl.num_programs(2) + + seq_len = tl.load(BSeqLen + batch_idx).to(tl.int64) + kv_start = tl.maximum(seq_len - WINDOW, 0) + window_len = seq_len - kv_start + total_blocks = tl.cdiv(window_len, BLOCK_SEQ) + if block_idx >= total_blocks: + return + + req_idx = tl.load(BReqIdx + batch_idx).to(tl.int64) + q_start = tl.load(BQStartLoc + batch_idx).to(tl.int64) + scratch_token = tl.full((), SCRATCH_START, tl.int64) + q_start + head_offsets = tl.arange(0, Q_HEAD_NUM) + q_heads = kv_head * gqa_group_size + head_offsets + q_heads = tl.where(head_offsets < gqa_group_size, q_heads, kv_head * gqa_group_size) + offs_d = tl.arange(0, BLOCK_DMODEL) + q = tl.load(Q + batch_idx * stride_qb + q_heads[:, None] * stride_qh + offs_d[None, :] * stride_qd) + + # Match the common GQA stage1's tiling and online-softmax arithmetic. + sum_exp = tl.zeros([Q_HEAD_NUM], dtype=tl.float32) + max_logic = tl.zeros([Q_HEAD_NUM], dtype=tl.float32) - float("inf") + acc = tl.zeros([Q_HEAD_NUM, BLOCK_DMODEL], dtype=tl.float32) + for block in range(block_idx, total_blocks, grid_block_num): + block_start = block * BLOCK_SEQ + block_end = tl.minimum(window_len, block_start + BLOCK_SEQ) + offs_n = block_start + tl.arange(0, BLOCK_N) + for tile in range(0, tl.cdiv(block_end - block_start, BLOCK_N)): + positions = tile * BLOCK_N + offs_n + mask = positions < block_end + token_pos = kv_start + positions + k_loc = tl.where(token_pos < seq_len - 1, req_idx * WINDOW + token_pos % WINDOW, scratch_token) + k = tl.load( + K + k_loc[None, :] * stride_kt + kv_head * stride_kh + offs_d[:, None] * stride_kd, + mask=mask[None, :], + other=0.0, + ) + att_value = tl.dot(q, k.to(q.dtype)) + att_value *= sm_scale + att_value = tl.where(mask[None, :], att_value, float("-inf")) + v = tl.load( + V + k_loc[:, None] * stride_vt + kv_head * stride_vh + offs_d[None, :] * stride_vd, + mask=mask[:, None], + other=0.0, + ) + cur_max_logic = tl.max(att_value, axis=1) + new_max_logic = tl.maximum(cur_max_logic, max_logic) + exp_logic = tl.exp(att_value - new_max_logic[:, None]) + logic_scale = tl.exp(max_logic - new_max_logic) + acc *= logic_scale[:, None] + acc += tl.dot(exp_logic.to(v.dtype), v) + sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=1) + max_logic = new_max_logic + + out_offsets = ( + batch_idx * stride_ob + q_heads[:, None] * stride_oh + block_idx * stride_os + offs_d[None, :] * stride_od + ) + log_offsets = batch_idx * stride_lb + q_heads * stride_lh + block_idx * stride_ls + tl.store(MidO + out_offsets, acc / sum_exp[:, None], mask=(head_offsets < gqa_group_size)[:, None]) + tl.store(MidLogSumExp + log_offsets, max_logic + tl.log(sum_exp), mask=head_offsets < gqa_group_size) + + +@torch.no_grad() +def sliding_window_decode_attention( + q, + k, + v, + b_req_idx, + b_seq_len, + b_q_start_loc, + sliding_window: int, + scratch_start: int, + out=None, + alloc_tensor_func=torch.empty, +): + """Decode one token per request without a token-to-sliding-KV index table.""" + batch_size, q_head_num, head_dim = q.shape + assert k.shape == v.shape and k.shape[-1] == head_dim + assert head_dim in {16, 32, 64, 128, 256, 512} + assert q_head_num % k.shape[1] == 0 + assert b_req_idx.shape == b_seq_len.shape == b_q_start_loc.shape == (batch_size,) + assert sliding_window > 0 and scratch_start >= sliding_window + assert q.dtype == k.dtype == v.dtype + + # Keep the common GQA wrapper's launch and reduction schedule unchanged. + block_seq = 256 + block_num = 128 if batch_size <= 16 else (64 if batch_size <= 64 else 32) + mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device=q.device) + mid_logsumexp = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device=q.device) + out = alloc_tensor_func(q.shape, dtype=q.dtype, device=q.device) if out is None else out + group_size = q_head_num // k.shape[1] + _sliding_window_decode_stage1[(batch_size, k.shape[1], block_num)]( + q, + k, + v, + b_req_idx, + b_seq_len, + b_q_start_loc, + mid_o, + mid_logsumexp, + 1.0 / (head_dim ** 0.5), + *q.stride(), + *k.stride(), + *v.stride(), + *mid_o.stride(), + *mid_logsumexp.stride(), + group_size, + WINDOW=sliding_window, + SCRATCH_START=scratch_start, + Q_HEAD_NUM=max(16, triton.next_power_of_2(group_size)), + BLOCK_SEQ=block_seq, + BLOCK_DMODEL=head_dim, + BLOCK_N=16, + num_warps=4, + num_stages=2, + ) + flash_decode_stage2( + mid_out=mid_o, + mid_out_logexpsum=mid_logsumexp, + B_Seqlen=b_seq_len, + out=out, + block_seq=block_seq, + sliding_window=(sliding_window - 1, 0), + ) + return out diff --git a/test/kernel/benchmark_sliding_window_attention.py b/test/kernel/benchmark_sliding_window_attention.py new file mode 100644 index 0000000000..5c20d58b08 --- /dev/null +++ b/test/kernel/benchmark_sliding_window_attention.py @@ -0,0 +1,348 @@ +"""Compare paged lookup with direct sliding-window addressing on identical KV. + +This is a warm-cache, single-layer microbenchmark, not a serving benchmark. +Preparation runs once per model forward, NOT once per attention layer; its +separate timing must not be multiplied by the model's sliding-layer count. +KV writes and window commits are excluded from both paths. + +Example (run only on a GPU approved for benchmarking): + python test/kernel/benchmark_sliding_window_attention.py --family both --output /tmp/sliding-attention.json +""" + +import argparse +import gc +import json +import os +import statistics +import time +from pathlib import Path +from types import SimpleNamespace + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import ( + gqa_token_decode_attention_flash_decoding, +) +from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import context_attention_fwd_gemma4_mm + + +# Verified from the text_config of gemma-4-E4B-it (TP2) and gemma-4-31B-it (TP4). +MODEL_SHAPES = { + "e4b": {"tp": 2, "q_heads": 4, "kv_heads": 1, "head_dim": 256, "window": 512}, + "31b": {"tp": 4, "q_heads": 8, "kv_heads": 4, "head_dim": 256, "window": 1024}, +} + + +# Frozen pre-change preparation kernels. Keeping them here lets the benchmark +# remain usable after the second request-token table is removed from serving. +@triton.jit +def _legacy_prepare_history( + mapping, req_ids, seq_lens, q_lens, stride_req, stride_seq, WINDOW: tl.constexpr, BLOCK: tl.constexpr +): + batch, block = tl.program_id(0), tl.program_id(1) + req_idx = tl.load(req_ids + batch) + history_end = tl.load(seq_lens + batch) - tl.load(q_lens + batch) + history_start = tl.maximum(0, history_end - WINDOW) + positions = history_start + block * BLOCK + tl.arange(0, BLOCK) + tl.store( + mapping + req_idx * stride_req + positions * stride_seq, + req_idx * WINDOW + positions % WINDOW, + positions < history_end, + ) + + +@triton.jit +def _legacy_prepare_current( + mapping, + req_ids, + seq_lens, + q_lens, + q_starts, + stride_req, + stride_seq, + SCRATCH_START: tl.constexpr, + BLOCK: tl.constexpr, +): + batch, block = tl.program_id(0), tl.program_id(1) + req_idx = tl.load(req_ids + batch) + q_len = tl.load(q_lens + batch) + history_end = tl.load(seq_lens + batch) - q_len + q_start = tl.load(q_starts + batch) + offsets = block * BLOCK + tl.arange(0, BLOCK) + tl.store( + mapping + req_idx * stride_req + (history_end + offsets) * stride_seq, + SCRATCH_START + q_start + offsets, + offsets < q_len, + ) + + +def _graph_timing(fn, unroll, samples): + """Capture repeated calls so Python launch overhead is outside GPU timing.""" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + fn() + fn() + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for _ in range(unroll): + fn() + graph.replay() + stream.synchronize() + timings = [] + for _ in range(samples): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + graph.replay() + end.record() + end.synchronize() + timings.append(start.elapsed_time(end) * 1000 / unroll) + return {"median_us": statistics.median(timings), "min_us": min(timings), "max_us": max(timings)} + + +def _eager_timing(fn, iterations, samples): + """Report enqueue time separately from synchronized whole-call latency. + + Enqueue time includes Python, allocation and driver calls, and can include + queue backpressure. It is not an isolated measurement of CPU computation. + """ + enqueue, wall = [], [] + for _ in range(samples): + torch.cuda.synchronize() + start = time.perf_counter_ns() + for _ in range(iterations): + fn() + submitted = time.perf_counter_ns() + torch.cuda.synchronize() + completed = time.perf_counter_ns() + enqueue.append((submitted - start) / iterations / 1000) + wall.append((completed - start) / iterations / 1000) + return {"enqueue_median_us": statistics.median(enqueue), "wall_median_us": statistics.median(wall)} + + +def _accuracy(actual, expected, atol, rtol): + # Limit validation workspace so the 8192-token cases do not temporarily + # allocate several additional full-sized FP32 attention outputs. + max_abs, squared_error, bitwise_equal = 0.0, 0.0, True + for actual_chunk, expected_chunk in zip(actual.flatten().split(1 << 20), expected.flatten().split(1 << 20)): + torch.testing.assert_close(actual_chunk, expected_chunk, atol=atol, rtol=rtol) + difference = actual_chunk.float() - expected_chunk.float() + max_abs = max(max_abs, difference.abs().max().item()) + squared_error += difference.square().sum().item() + bitwise_equal = bitwise_equal and torch.equal(actual_chunk, expected_chunk) + return {"max_abs": max_abs, "rms": (squared_error / actual.numel()) ** 0.5, "bitwise_equal": bitwise_equal} + + +@torch.inference_mode() +def _benchmark_case(family, phase, batch, q_len, args): + shape = MODEL_SHAPES[family] + q_heads, kv_heads, dim, window = (shape[name] for name in ["q_heads", "kv_heads", "head_dim", "window"]) + request_slots = batch + 1 + scratch_start = request_slots * window + token_num = batch * q_len + max_seq_len = args.history + 7 * (batch - 1) + q_len + backing_bytes = (scratch_start + token_num) * 2 * kv_heads * dim * 2 + mapping_bytes = request_slots * max_seq_len * 4 + q_and_outputs_bytes = 3 * token_num * q_heads * dim * 2 + legacy_blocks = 128 if batch <= 16 else 64 if batch <= 64 else 32 + decode_workspace_bytes = batch * q_heads * legacy_blocks * (dim * 2 + 4) if phase == "decode" else 0 + validation_bytes = min(token_num * q_heads * dim, 1 << 20) * 16 + estimated_bytes = backing_bytes + mapping_bytes + q_and_outputs_bytes + decode_workspace_bytes + validation_bytes + if estimated_bytes > args.max_case_mib * 1024 ** 2: + raise ValueError( + f"case needs at least {estimated_bytes / 1024 ** 2:.1f} MiB; increase --max-case-mib explicitly" + ) + + req_ids = torch.arange(batch, 0, -1, dtype=torch.int32, device="cuda") + history_lens = args.history + torch.arange(batch, dtype=torch.int32, device="cuda") * 7 + seq_lens = history_lens + q_len + q_lens = torch.full((batch,), q_len, dtype=torch.int32, device="cuda") + q_starts = torch.arange(batch, dtype=torch.int32, device="cuda") * q_len + mapping = torch.full((request_slots, max_seq_len), -1, dtype=torch.int32, device="cuda") + # Both paths read precisely this tensor: only the address computation changes. + backing = torch.randn((scratch_start + token_num, 2 * kv_heads, dim), dtype=torch.bfloat16, device="cuda") + k, v = backing[:, :kv_heads], backing[:, kv_heads:] + q = torch.randn((token_num, q_heads, dim), dtype=torch.bfloat16, device="cuda") + old_out, new_out = torch.empty_like(q), torch.empty_like(q) + image_ends = torch.zeros((token_num,), dtype=torch.int32, device="cuda") + + def prepare(): + _legacy_prepare_history[(batch, triton.cdiv(window, 256))]( + mapping, req_ids, seq_lens, q_lens, *mapping.stride(), WINDOW=window, BLOCK=256 + ) + _legacy_prepare_current[(batch, triton.cdiv(q_len, 256))]( + mapping, req_ids, seq_lens, q_lens, q_starts, *mapping.stride(), SCRATCH_START=scratch_start, BLOCK=256 + ) + + if phase == "prefill": + + def legacy_attention(): + context_attention_fwd_gemma4_mm( + q, + k, + v, + old_out, + req_ids, + q_starts, + seq_lens, + history_lens, + q_len, + mapping, + image_ends, + sliding_window=(window - 1, 0), + ) + + def direct_attention(): + context_attention_fwd_gemma4_mm( + q, + k, + v, + new_out, + req_ids, + q_starts, + seq_lens, + history_lens, + q_len, + None, + image_ends, + sliding_window=(window - 1, 0), + scratch_start=scratch_start, + ) + + else: + from lightllm.models.gemma4.triton_kernel.sliding_window_decode import sliding_window_decode_attention + + infer_state = SimpleNamespace( + batch_size=batch, + req_manager=SimpleNamespace(req_to_token_indexs=mapping), + b_req_idx=req_ids, + b_seq_len=seq_lens, + max_kv_seq_len=max_seq_len, + ) + + def legacy_attention(): + gqa_token_decode_attention_flash_decoding(q, infer_state, k, v, out=old_out, sliding_window=(window - 1, 0)) + + def direct_attention(): + sliding_window_decode_attention( + q, + k, + v, + req_ids, + seq_lens, + q_starts, + sliding_window=window, + scratch_start=scratch_start, + out=new_out, + ) + + def legacy_step(): + prepare() + legacy_attention() + + legacy_step() + direct_attention() + torch.cuda.synchronize() + accuracy = _accuracy(new_out, old_out, atol=args.atol, rtol=args.rtol) + functions = { + "legacy_prepare_once": prepare, + "legacy_attention_only": legacy_attention, + "legacy_prepare_plus_one_attention": legacy_step, + "direct_attention_only": direct_attention, + } + graph = {name: _graph_timing(fn, args.graph_unroll, args.samples) for name, fn in functions.items()} + eager = {name: _eager_timing(fn, args.eager_iterations, args.samples) for name, fn in functions.items()} + return { + "family": family, + "phase": phase, + "batch_size": batch, + "q_len": q_len, + "history_min": args.history, + "history_max": max_seq_len - q_len, + **shape, + "backing_mib": backing_bytes / 1024 ** 2, + "removed_index_table_mib": mapping_bytes / 1024 ** 2, + "accuracy": accuracy, + "graph": graph, + "eager": eager, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--family", choices=["e4b", "31b", "both"], default="both") + parser.add_argument("--phase", choices=["prefill", "decode", "both"], default="both") + parser.add_argument("--device", type=int, default=0, help="logical CUDA device within CUDA_VISIBLE_DEVICES") + parser.add_argument("--history", type=int, default=32768) + parser.add_argument("--prefill-lengths", type=int, nargs="+", default=[512, 4096, 8192]) + parser.add_argument("--decode-batches", type=int, nargs="+", default=[1, 8, 32]) + parser.add_argument("--graph-unroll", type=int, default=16) + parser.add_argument("--eager-iterations", type=int, default=20) + parser.add_argument("--samples", type=int, default=5) + parser.add_argument("--max-case-mib", type=float, default=256) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--atol", type=float, default=0.005) + parser.add_argument("--rtol", type=float, default=0.02) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + if ( + args.history < 0 + or min(args.prefill_lengths + args.decode_batches + [args.graph_unroll, args.eager_iterations, args.samples]) + < 1 + ): + parser.error("history must be nonnegative; lengths, batches, unroll, iterations and samples must be positive") + torch.cuda.set_device(args.device) + torch.manual_seed(args.seed) + report = { + "scope": "Shared-GPU-capable warm-cache attention microbenchmark, not end-to-end serving throughput.", + "notes": [ + "Both paths use identical ring/scratch KV, Q and metadata; only attention addressing differs.", + "Legacy preparation is once per model forward, not once per attention layer.", + "KV writes, window commit, linear/full layers, scheduler and CPU-cache operations are excluded.", + "Eager enqueue time includes Python/allocator/driver calls and possible queue backpressure.", + "CUDA graph timing removes host enqueue overhead; shared GPU contention can still affect results.", + ], + "gpu": torch.cuda.get_device_name(args.device), + "torch": torch.__version__, + "triton": triton.__version__, + "cuda_visible_devices": os.getenv("CUDA_VISIBLE_DEVICES"), + "autotune_level": os.getenv("LIGHTLLM_TRITON_AUTOTUNE_LEVEL", "0"), + "arguments": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()}, + "cases": [], + } + print( + "family phase batch q old-att graph-us new-att graph-us prepare graph-us old-step graph-us", + flush=True, + ) + for family in MODEL_SHAPES if args.family == "both" else [args.family]: + cases = [] + if args.phase in ["prefill", "both"]: + cases.extend(("prefill", 1, length) for length in args.prefill_lengths) + if args.phase in ["decode", "both"]: + cases.extend(("decode", batch, 1) for batch in args.decode_batches) + for phase, batch, q_len in cases: + result = _benchmark_case(family, phase, batch, q_len, args) + report["cases"].append(result) + timings = result["graph"] + print( + f"{family:6} {phase:7} {batch:5} {q_len:5} " + f"{timings['legacy_attention_only']['median_us']:18.3f} " + f"{timings['direct_attention_only']['median_us']:17.3f} " + f"{timings['legacy_prepare_once']['median_us']:17.3f} " + f"{timings['legacy_prepare_plus_one_attention']['median_us']:18.3f}", + flush=True, + ) + gc.collect() + torch.cuda.empty_cache() + if args.output: + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(f"JSON report: {args.output}", flush=True) + else: + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/kernel/test_gemma4_hybrid_graph.py b/test/kernel/test_gemma4_hybrid_graph.py new file mode 100644 index 0000000000..0720af3411 --- /dev/null +++ b/test/kernel/test_gemma4_hybrid_graph.py @@ -0,0 +1,126 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState +from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager +from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow +from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig +from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo +from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _int_tensor(values): + return torch.tensor(values, device="cuda", dtype=torch.int32) + + +def _layer(is_sliding, shared): + layer = object.__new__(Gemma4TransformerLayerInfer) + layer.is_sliding, layer.is_kv_shared_ = is_sliding, shared + layer.layer_num_ = (4 if shared else 2) if is_sliding else (3 if shared else 1) + layer.kv_share_target_layer_ = (2 if is_sliding else 1) if shared else None + layer.tp_q_head_num_, layer.head_dim_ = (4, 256) if is_sliding else (8, 512) + layer.commit_sliding_state_ = False # Isolate graph metadata and KV reads from snapshot/commit tests. + layer.alloc_tensor = lambda shape, dtype, device="cuda": torch.empty(shape, dtype=dtype, device=device) + return layer + + +def _state(model, req_manager, req_ids, seq_lengths, q_starts): + state = Gemma4InferStateInfo() + state.req_manager, state.mem_manager = req_manager, req_manager.mem_manager + state.is_prefill = False + state.batch_size, state.max_q_seq_len = len(req_ids), 1 + state.max_kv_seq_len = req_manager.req_to_token_indexs.shape[1] + state.b_req_idx, state.b_seq_len = _int_tensor(req_ids), _int_tensor(seq_lengths) + state.input_ids = torch.zeros(len(req_ids), device="cuda", dtype=torch.int64) + backend = SimpleNamespace(model=model) + state.decode_att_state = TritonDecodeAttState(backend=backend, infer_state=state) + state.decode_att_state1 = TritonDecodeAttState(backend=backend, infer_state=state) + state.init_some_extra_state(model) + # Independently vary scratch locations to expose stale graph metadata. + state.b_q_start_loc = _int_tensor(q_starts) + state.init_att_state() + return state + + +@pytest.mark.parametrize("window", [512, 1024]) +@pytest.mark.parametrize("shared", [False, True]) +def test_gemma_hybrid_decode_graph_copies_state_without_replacing_full_token_table(monkeypatch, window, shared): + from lightllm.common.triton_utils import autotuner + + monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.CLOSE_AUTOTUNE) + torch.manual_seed(42) + req_slots, batch_size, max_seq_len = 6, 4, 3 * window + 32 + config = SlidingWindowCacheConfig({0: 0, 2: 1, 4: 1}, {1: 0, 3: 0}, window, 1, 256, 2, 512, torch.bfloat16) + # Skip launch-time distributed/profile setup, retaining the real cache access methods. + mem_manager = object.__new__(HybridSlidingMemoryManager) + mem_manager.sliding_config, mem_manager.head_num = config, config.full_head_num + mem_manager.kv_buffer = torch.randn((1, 1024, 4, 512), device="cuda", dtype=config.dtype) + manager = object.__new__(ReqManagerForSlidingWindow) + manager.mem_manager, manager.sliding_config = mem_manager, config + manager.sliding_window, manager.scratch_token_num = window, batch_size + manager.scratch_start = req_slots * window + manager.sliding_kv_buffer = torch.randn( + (2, manager.scratch_start + batch_size, 2, 256), device="cuda", dtype=config.dtype + ) + manager.req_to_sliding_window = manager.sliding_kv_buffer[:, : manager.scratch_start].view( + 2, req_slots, window, 2, 256 + ) + manager.req_to_token_indexs = torch.randint(1024, (req_slots, max_seq_len), device="cuda", dtype=torch.int32) + original_table = manager.req_to_token_indexs + expected_table = original_table.clone() + cos = torch.ones((max_seq_len, 128), device="cuda", dtype=config.dtype) + sin = torch.zeros_like(cos) + model = SimpleNamespace( + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda is_draft: 0), + is_mtp_draft_model=False, + _cos_cached_sliding=cos, + _sin_cached_sliding=sin, + _cos_cached_full=cos, + _sin_cached_full=sin, + ) + captured = _state(model, manager, [0, 2, 5, 5], [window + 5, 17, 2, 2], [0, 1, 2, 3]) + captured.is_cuda_graph = True + sliding, full = _layer(True, shared), _layer(False, shared) + sliding.sliding_window_, full.sliding_window_ = window, 0 + q_sliding = torch.randn((batch_size, 4, 256), device="cuda", dtype=config.dtype) + q_full = torch.randn((batch_size, 8, 512), device="cuda", dtype=config.dtype) + + def forward(state): + return ( + sliding._token_attention_kernel(q_sliding, state, None), + full._token_attention_kernel(q_full, state, None), + ) + + forward(captured) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_outputs = forward(captured) + copied_fields = ("b_req_idx", "b_seq_len", "b_q_start_loc", "position_ids", "sliding_window_mem_index") + captured_ptrs = {name: getattr(captured, name).data_ptr() for name in copied_fields} + new_state = _state(model, manager, [3, 1, 5, 5], [2 * window + 11, 1, 2, 2], [2, 0, 1, 3]) + captured.copy_for_cuda_graph(new_state) + q_sliding.mul_(0.5) + q_full.mul_(0.75) + manager.sliding_kv_buffer.mul_(0.75) + mem_manager.kv_buffer.mul_(0.5) + graph.replay() + eager_outputs = forward(new_state) + + for actual, expected in zip(graph_outputs, eager_outputs): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert torch.isfinite(actual).all() + for name in copied_fields: + assert getattr(captured, name).data_ptr() == captured_ptrs[name] + torch.testing.assert_close(getattr(captured, name), getattr(new_state, name), atol=0, rtol=0) + for state in (captured, new_state): + assert state.decode_att_state.infer_state is state + assert state.decode_att_state1.infer_state is state + assert state.req_manager is manager + assert state.req_manager.req_to_token_indexs is original_table + assert not hasattr(manager, "req_to_sliding_window_indexs") + torch.testing.assert_close(original_table, expected_table, atol=0, rtol=0) diff --git a/test/kernel/test_sliding_window_cpu_cache_attention.py b/test/kernel/test_sliding_window_cpu_cache_attention.py index f52084a648..1de02e8b7d 100644 --- a/test/kernel/test_sliding_window_cpu_cache_attention.py +++ b/test/kernel/test_sliding_window_cpu_cache_attention.py @@ -64,12 +64,14 @@ def test_cpu_window_checkpoint_resumes_attention_exactly(window, history_len, q_ manager.sliding_config, manager.sliding_window = config, window manager.scratch_token_num, manager.scratch_start = q_len, 3 * window manager.mem_manager = SimpleNamespace(linear_att_big_page_buffers=pages) - manager.req_to_sliding_window = torch.zeros( + manager.sliding_kv_buffer = torch.zeros( (2, manager.scratch_start + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16 ) - manager.req_to_sliding_window_indexs = torch.full((3, seq_len), -1, device="cuda", dtype=torch.int32) + manager.req_to_sliding_window = manager.sliding_kv_buffer[:, : manager.scratch_start].view( + 2, 3, window, 2, head_dim + ) manager.restore_big_page_state(len(endpoints) - 1, SimpleNamespace(req_idx=req_idx)) - manager.req_to_sliding_window[:, manager.scratch_start :] = reference[:, history_len:] + manager.sliding_kv_buffer[:, manager.scratch_start :] = reference[:, history_len:] state = SimpleNamespace( input_ids=int_tensor([0] * q_len), b_req_idx=int_tensor([req_idx]), @@ -85,9 +87,9 @@ def test_cpu_window_checkpoint_resumes_attention_exactly(window, history_len, q_ for layer_index in [0, 2, 3]: physical_layer = config.get_sliding_layer_index(layer_index) actual, expected = torch.empty_like(q), torch.empty_like(q) - for kv, indexes, output in [ - (manager.req_to_sliding_window[physical_layer], manager.req_to_sliding_window_indexs, actual), - (reference[physical_layer], reference_indexes, expected), + for kv, indexes, output, scratch in [ + (manager.sliding_kv_buffer[physical_layer], None, actual, manager.scratch_start), + (reference[physical_layer], reference_indexes, expected, None), ]: context_attention_fwd_gemma4_mm( q, @@ -102,5 +104,6 @@ def test_cpu_window_checkpoint_resumes_attention_exactly(window, history_len, q_ indexes, image_end, sliding_window=(window - 1, 0), + scratch_start=scratch, ) torch.testing.assert_close(actual, expected, atol=0, rtol=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..50504674b5 --- /dev/null +++ b/test/kernel/test_sliding_window_decode.py @@ -0,0 +1,151 @@ +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, +) +from lightllm.models.gemma4.triton_kernel.sliding_window_decode import sliding_window_decode_attention + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _int_tensor(values, dtype=torch.int32): + return torch.tensor(values, device="cuda", dtype=dtype) + + +def _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start): + # Only the visible suffix matters. Rebase very long sequences to avoid + # allocating a token table proportional to their virtual token positions. + indexes = torch.zeros((max(req_ids) + 1, window), device="cuda", dtype=torch.int64) + visible_lengths = [min(length, window) for length in seq_lengths] + for req_idx, length, q_start, visible_len in zip(req_ids, seq_lengths, q_starts, visible_lengths): + positions = torch.arange(length - visible_len, length, device="cuda", dtype=torch.int64) + indexes[req_idx, :visible_len] = req_idx * window + positions % window + indexes[req_idx, visible_len - 1] = scratch_start + q_start + state = SimpleNamespace( + batch_size=len(req_ids), + b_req_idx=_int_tensor(req_ids), + b_seq_len=_int_tensor(visible_lengths), + max_kv_seq_len=window, + req_manager=SimpleNamespace(req_to_token_indexs=indexes), + ) + kv_heads = kv.shape[1] // 2 + return gqa_token_decode_attention_flash_decoding( + q=q, + infer_state=state, + cache_k=kv[:, :kv_heads], + cache_v=kv[:, kv_heads:], + out=torch.empty_like(q), + sliding_window=(window - 1, 0), + ) + + +@pytest.fixture(autouse=True) +def _use_default_gqa_schedule(monkeypatch): + # Compare identical math and tiling rather than a machine-specific tune. + from lightllm.common.triton_utils import autotuner + + monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.CLOSE_AUTOTUNE) + + +@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_formula_decode_matches_table_gqa_exactly(dtype, window, q_heads, kv_heads, head_dim): + torch.manual_seed(42) + req_ids = [6, 0, 4, 2, 7, 1] + seq_lengths = [1, 2, window - 1, window, window + 1, 3 * window + 7] + q_starts = [7, 1, 11, 4, 9, 2] + scratch_start = 9 * window + kv = torch.randn((scratch_start + 12, 2 * kv_heads, head_dim), device="cuda", dtype=dtype) + q = torch.randn((len(req_ids), q_heads, head_dim), device="cuda", dtype=dtype) + output = torch.empty_like(q) + actual = sliding_window_decode_attention( + q, + kv[:, :kv_heads], + kv[:, kv_heads:], + _int_tensor(req_ids), + _int_tensor(seq_lengths), + _int_tensor(q_starts), + window, + scratch_start, + out=output, + ) + expected = _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start) + assert actual is output + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("batch_size", [1, 17, 65]) +def test_formula_decode_preserves_gqa_batch_schedule(batch_size): + window, head_dim, scratch_start = 512, 64, (batch_size + 1) * 512 + req_ids = list(reversed(range(batch_size))) + seq_lengths = [2 * window + i + 1 for i in range(batch_size)] + q_starts = list(range(batch_size)) + kv = torch.randn((scratch_start + batch_size, 2, head_dim), device="cuda", dtype=torch.bfloat16) + q = torch.randn((batch_size, 4, head_dim), device="cuda", dtype=torch.bfloat16) + actual = sliding_window_decode_attention( + q, + kv[:, :1], + kv[:, 1:], + _int_tensor(req_ids), + _int_tensor(seq_lengths), + _int_tensor(q_starts), + window, + scratch_start, + ) + expected = _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_formula_decode_supports_int64_virtual_token_positions(): + window, scratch_start = 512, 4 * 512 + req_ids, seq_lengths, q_starts = [2, 0], [2 ** 31 + 17, 2 ** 32 + 31], [0, 1] + kv = torch.randn((scratch_start + 2, 2, 64), device="cuda", dtype=torch.bfloat16) + q = torch.randn((2, 4, 64), device="cuda", dtype=torch.bfloat16) + actual = sliding_window_decode_attention( + q, + kv[:, :1], + kv[:, 1:], + _int_tensor(req_ids), + _int_tensor(seq_lengths, dtype=torch.int64), + _int_tensor(q_starts), + window, + scratch_start, + ) + expected = _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("window", [512, 1024]) +def test_formula_decode_cuda_graph_replay_with_changed_requests_and_padding(dtype, window): + scratch_start = 8 * window + kv = torch.randn((scratch_start + 4, 4, 256), device="cuda", dtype=dtype) + q = torch.randn((4, 8, 256), device="cuda", dtype=dtype) + b_req = _int_tensor([4, 1, 7, 7]) + b_seq = _int_tensor([window + 7, 3, 2, 2]) + b_start = _int_tensor([3, 0, 1, 2]) + out = torch.empty_like(q) + + def forward(): + sliding_window_decode_attention(q, kv[:, :2], kv[:, 2:], b_req, b_seq, b_start, window, scratch_start, out=out) + + forward() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + forward() + req_ids, seq_lengths, q_starts = [2, 5, 7, 7], [2 * window + 3, 1, 2, 2], [1, 3, 2, 0] + b_req.copy_(_int_tensor(req_ids)) + b_seq.copy_(_int_tensor(seq_lengths)) + b_start.copy_(_int_tensor(q_starts)) + q.mul_(0.5) + kv.mul_(0.75) + graph.replay() + expected = _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start) + # Padding may share the hold request ID; its outputs are intentionally discarded. + torch.testing.assert_close(out[:2], expected[:2], atol=0, rtol=0) + assert torch.isfinite(out).all() diff --git a/test/kernel/test_sliding_window_prefill.py b/test/kernel/test_sliding_window_prefill.py new file mode 100644 index 0000000000..16e0d6d726 --- /dev/null +++ b/test/kernel/test_sliding_window_prefill.py @@ -0,0 +1,162 @@ +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_ring_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 = [47 + q_lens[1], 11, 83 + q_lens[1] + q_lens[0]] + 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 + + scratch_start = req_slots * window + 13 + runtime = torch.full((scratch_start + query_num, 2 * kv_heads, head_dim), -3, device="cuda", dtype=dtype) + for req_id, history, length, start in zip(req_ids, histories, lengths, starts): + old_positions = torch.arange(max(0, history - window), history, device="cuda") + runtime[req_id * window + old_positions % window] = reference[mapping[req_id, old_positions].long()] + runtime[scratch_start + start : scratch_start + start + length - history] = reference[ + mapping[req_id, history:length].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=None, + scratch_start=scratch_start, + **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_ring_prefill_matches_paged_history_and_scratch(window, q_len, dtype): + _compare_ring_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_ring_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_ring_and_paged(window, 384, dtype, head_dim=256, image_span=image_span) + + +@pytest.mark.parametrize("window", [512, 1024]) +def test_ring_prefill_cuda_graph_replay_reads_updated_request_metadata(window): + torch.manual_seed(43) + req_slots, head_dim, max_q_len = 6, 64, window + 33 + scratch_start, query_num = req_slots * window + 11, 2 * (window + 64) + 97 + runtime = torch.randn((scratch_start + query_num, 4, head_dim), device="cuda", dtype=torch.bfloat16) + q = torch.randn((query_num, 4, head_dim), device="cuda", dtype=torch.bfloat16) + out = torch.full_like(q, -11) + int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + b_req = int_tensor([1, 4]) + b_history = int_tensor([window + 3, 0]) + b_seq = int_tensor([2 * window + 20, 31]) + b_start = int_tensor([13, max_q_len + 47]) + image_ends = torch.zeros(query_num, device="cuda", dtype=torch.int32) + kwargs = dict( + q=q, + k=runtime[:, :2], + v=runtime[:, 2:], + b_req_idx=b_req, + b_start_loc=b_start, + b_seq_len=b_seq, + b_prompt_cache_len=b_history, + max_input_len=max_q_len, + b_image_token_end=image_ends, + sliding_window=(window - 1, 0), + ) + + def forward(): + context_attention_fwd_gemma4_mm(o=out, req_to_token_indexs=None, scratch_start=scratch_start, **kwargs) + + forward() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + forward() + + req_ids, histories, q_lens, starts = [4, 2], [2 * window + 57, window - 5], [max_q_len, 17], [5, max_q_len + 59] + lengths = [history + count for history, count in zip(histories, q_lens)] + b_req.copy_(int_tensor(req_ids)) + b_history.copy_(int_tensor(histories)) + b_seq.copy_(int_tensor(lengths)) + b_start.copy_(int_tensor(starts)) + q.mul_(0.5) + runtime.mul_(0.75) + out.fill_(-11) + graph.replay() + + # Materialize a table only for the independent old-path reference, after + # changing every piece of GPU metadata used by the captured ring kernel. + mapping = torch.full((req_slots, max(lengths)), -1, device="cuda", dtype=torch.int32) + for req_id, history, length, start in zip(req_ids, histories, lengths, starts): + positions = torch.arange(length, device="cuda", dtype=torch.int32) + mapping[req_id, :length] = torch.where( + positions < history, + req_id * window + positions % window, + scratch_start + start + positions - history, + ) + expected = torch.full_like(q, -11) + context_attention_fwd_gemma4_mm(o=expected, req_to_token_indexs=mapping, **kwargs) + # Include gaps to verify the captured grid respects the new query lengths. + torch.testing.assert_close(out, expected, atol=0, rtol=0) diff --git a/test/kernel/test_sliding_window_state.py b/test/kernel/test_sliding_window_state.py index 4a8ab7599a..c01a00a09f 100644 --- a/test/kernel/test_sliding_window_state.py +++ b/test/kernel/test_sliding_window_state.py @@ -3,10 +3,7 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( - commit_sliding_window_state, - prepare_sliding_window_indexes, -) +from lightllm.common.basemodel.triton_kernel.sliding_window_state import commit_sliding_window_state from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow from lightllm.common.kv_cache_mem_manager.operator.hybrid_sliding import HybridSlidingMemOperator from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig, SlidingWindowStateCacheManager @@ -17,6 +14,36 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_request_window_view_shares_storage_and_has_no_second_token_table(monkeypatch): + monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.ReqSamplingParamsManager", lambda size: None) + config = SlidingWindowCacheConfig({0: 0, 2: 1}, {1: 0}, 32, 1, 64, 1, 64, torch.bfloat16) + manager = ReqManagerForSlidingWindow(2, 128, None, config, scratch_token_num=7) + assert manager.req_to_sliding_window.shape == (2, 3, 32, 2, 64) + assert manager.req_to_sliding_window.data_ptr() == manager.sliding_kv_buffer.data_ptr() + assert manager.req_to_sliding_window.stride(0) == manager.sliding_kv_buffer.stride(0) + assert not hasattr(manager, "req_to_sliding_window_indexs") + assert manager.req_to_token_indexs.shape == (3, 128) + manager.req_to_sliding_window[:, 1].fill_(7) + torch.testing.assert_close( + manager.sliding_kv_buffer[:, 32:64], torch.full_like(manager.sliding_kv_buffer[:, 32:64], 7) + ) + manager.sliding_kv_buffer[:, manager.scratch_start :].fill_(11) + pages = SlidingWindowStateCacheManager(1, config) + manager.save_small_page_state(1, 0, pages) + manager.init_hybrid_attention_state(SimpleNamespace(req_idx=1)) + assert torch.count_nonzero(manager.req_to_sliding_window).item() == 0 + manager._restore_state(2, pages, 0) # Include the reserved hold request slot. + assert torch.all(manager.req_to_sliding_window[:, 2] == 7) + assert torch.all(manager.sliding_kv_buffer[:, manager.scratch_start :] == 11) + state = SimpleNamespace(input_ids=torch.zeros(7, dtype=torch.int32, device="cuda")) + manager.prepare_sliding_window(state) + assert torch.count_nonzero(manager.req_to_token_indexs).item() == 0 + torch.testing.assert_close( + state.sliding_window_mem_index, + torch.arange(manager.scratch_start, manager.scratch_start + 7, device="cuda"), + ) + + @pytest.mark.parametrize("layer_index,is_shared", [(5, False), (11, False), (17, True)]) def test_full_kv_write_maps_logical_layer_once_and_skips_shared_readers(layer_index, is_shared): config = SlidingWindowCacheConfig({0: 0}, {5: 0, 11: 1, 17: 1}, 32, 1, 64, 1, 64, torch.bfloat16) @@ -47,18 +74,19 @@ def test_ring_attention_and_commit_match_token_cache(history_len, q_len): old_positions = torch.arange(max(0, history_len - window), history_len, device="cuda") runtime[req_idx * window + old_positions % window] = reference[old_positions] runtime[scratch_start:] = reference[history_len:] - mapping = torch.full((3, seq_len), -1, device="cuda", dtype=torch.int32) b_req = torch.tensor([req_idx], device="cuda", dtype=torch.int32) b_seq = torch.tensor([seq_len], device="cuda", dtype=torch.int32) b_q = torch.tensor([q_len], device="cuda", dtype=torch.int32) b_start = torch.tensor([0], device="cuda", dtype=torch.int32) b_history = torch.tensor([history_len], device="cuda", dtype=torch.int32) image_end = torch.zeros(q_len, device="cuda", dtype=torch.int32) - prepare_sliding_window_indexes(mapping, b_req, b_seq, b_q, b_start, window, scratch_start, q_len) q = torch.randn((q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) actual, expected = torch.empty_like(q), torch.empty_like(q) reference_mapping = torch.arange(seq_len, device="cuda", dtype=torch.int32).expand(3, -1) - for buffer, indexes, output in [(runtime, mapping, actual), (reference, reference_mapping, expected)]: + for buffer, indexes, output, scratch in [ + (runtime, None, actual, scratch_start), + (reference, reference_mapping, expected, None), + ]: context_attention_fwd_gemma4_mm( q, buffer[:, :1], @@ -72,6 +100,7 @@ def test_ring_attention_and_commit_match_token_cache(history_len, q_len): indexes, image_end, sliding_window=(window - 1, 0), + scratch_start=scratch, ) torch.testing.assert_close(actual, expected, atol=0, rtol=0) commit_sliding_window_state(runtime, b_req, b_seq, b_q, b_start, window, scratch_start, q_len) @@ -80,8 +109,9 @@ def test_ring_attention_and_commit_match_token_cache(history_len, q_len): assert torch.count_nonzero(runtime[:window]).item() == 0 -def test_shared_kv_is_committed_only_after_last_reader_and_snapshot_is_independent(): - window, history_len, q_len, head_dim = 512, 512, 256, 64 +@pytest.mark.parametrize("is_prefill", [True, False]) +def test_shared_kv_is_committed_only_after_last_reader_and_snapshot_is_independent(is_prefill): + window, history_len, q_len, head_dim = 512, 512, 256 if is_prefill else 1, 64 layout, owners, last_reader = get_kv_cache_layout( {"layer_types": ["sliding_attention", "full_attention", "sliding_attention"], "num_kv_shared_layers": 1} ) @@ -93,11 +123,11 @@ def test_shared_kv_is_committed_only_after_last_reader_and_snapshot_is_independe manager.sliding_window = window manager.scratch_token_num = q_len manager.scratch_start = 2 * window - manager.req_to_sliding_window = torch.zeros( - (1, 2 * window + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16 + manager.sliding_kv_buffer = torch.zeros((1, 2 * window + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) + manager.req_to_sliding_window = manager.sliding_kv_buffer[:, : manager.scratch_start].view( + 1, 2, window, 2, head_dim ) - manager.req_to_sliding_window[0, manager.scratch_start :, 1] = 1 - manager.req_to_sliding_window_indexs = torch.zeros((2, history_len + q_len), device="cuda", dtype=torch.int32) + manager.sliding_kv_buffer[0, manager.scratch_start :, 1] = 1 int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) state = SimpleNamespace( input_ids=int_tensor([0] * q_len), @@ -119,22 +149,25 @@ def test_shared_kv_is_committed_only_after_last_reader_and_snapshot_is_independe layer.is_kv_shared_, layer.kv_share_target_layer_ = index != owners[index], owners[index] layer.commit_sliding_state_ = last_reader[owners[index]] == index layer.tp_q_head_num_, layer.head_dim_ = 1, head_dim - layer.alloc_tensor = lambda shape, dtype: torch.empty(shape, dtype=dtype, device="cuda") - outputs.append(layer._context_attention_kernel(q, None, state, None)) + layer.alloc_tensor = lambda shape, dtype, device="cuda": torch.empty(shape, dtype=dtype, device=device) + if is_prefill: + outputs.append(layer._context_attention_kernel(q, None, state, None)) + else: + outputs.append(layer._token_attention_kernel(q, state, None)) if index == 0: - assert torch.count_nonzero(manager.req_to_sliding_window[:, :window]).item() == 0 + assert torch.count_nonzero(manager.req_to_sliding_window[:, 0]).item() == 0 torch.testing.assert_close(outputs[0], outputs[1], atol=0, rtol=0) assert outputs[1][0, 0, 0].item() == 1 / window - assert manager.req_to_sliding_window[0, 0, 1, 0].item() == 1 + assert manager.req_to_sliding_window[0, 0, 0, 1, 0].item() == 1 pages = SlidingWindowStateCacheManager(2, config) page = pages.alloc_one_state_cache() manager.save_small_page_state(0, page, pages) saved = pages.get_state_cache(page).clone() - manager.req_to_sliding_window[:, :window].fill_(7) + manager.req_to_sliding_window[:, 0].fill_(7) torch.testing.assert_close(pages.get_state_cache(page), saved, atol=0, rtol=0) manager._restore_state(1, pages, page) - torch.testing.assert_close(manager.req_to_sliding_window[:, window : 2 * window], saved, atol=0, rtol=0) + torch.testing.assert_close(manager.req_to_sliding_window[:, 1], saved, atol=0, rtol=0) pages.free_state_cache([page]) assert pages.get_free_cache_num() == 2 @@ -158,12 +191,10 @@ def test_batched_window_commit_with_hold_request_and_cuda_graph_replay(window, q starts = [0, q_lengths[0], q_lengths[0] + q_lengths[1]] int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) b_req, b_seq, b_q, b_start = map(int_tensor, [req_ids, lengths, q_lengths, starts]) - mapping = torch.full((4, max(replay_lengths)), -1, device="cuda", dtype=torch.int32) runtime = torch.zeros((scratch_start + sum(q_lengths), 2, head_dim), device="cuda", dtype=torch.bfloat16) runtime[scratch_start:] = torch.randn_like(runtime[scratch_start:]) def forward(): - prepare_sliding_window_indexes(mapping, b_req, b_seq, b_q, b_start, window, scratch_start, max(q_lengths)) commit_sliding_window_state(runtime, b_req, b_seq, b_q, b_start, window, scratch_start, max(q_lengths)) forward() @@ -182,11 +213,6 @@ def forward(): expected_ring = torch.zeros_like(runtime[req * window : (req + 1) * window]) expected_ring[positions % window] = runtime[scratch_start + start + tail_start : scratch_start + start + q_len] torch.testing.assert_close(runtime[req * window : (req + 1) * window], expected_ring, atol=0, rtol=0) - positions = torch.arange(seq - q_len, seq, device="cuda") - torch.testing.assert_close( - mapping[req, positions], - torch.arange(scratch_start + start, scratch_start + start + q_len, device="cuda", dtype=torch.int32), - ) assert torch.count_nonzero(runtime[window : 2 * window]).item() == 0 diff --git a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py index 0e4267ef5f..33ead8cb8d 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py @@ -119,8 +119,8 @@ def test_sliding_big_snapshot_skips_invalid_requests_and_copies_only_selected_pa manager = object.__new__(ReqManagerForSlidingWindow) manager.sliding_window = 4 manager.mem_manager = SimpleNamespace(linear_att_big_page_buffers=pages) - manager.req_to_sliding_window = torch.arange(2 * 12 * 2 * 4, dtype=torch.float32).reshape(2, 12, 2, 4) - expected = manager.req_to_sliding_window[:, 4:8].clone() + manager.req_to_sliding_window = torch.arange(2 * 12 * 2 * 4, dtype=torch.float32).reshape(2, 3, 4, 2, 4) + expected = manager.req_to_sliding_window[:, 1].clone() # Skipped request IDs are deliberately out of range; GPU request IDs must not be read. manager.save_big_page_states(b_req_idx=object(), req_indexes=[999, 1, 888], buffer_indexes=[-1, 2, -1]) diff --git a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py index 8a4247dac8..f8194a4705 100644 --- a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py +++ b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py @@ -51,7 +51,7 @@ def alloc(need_size): mem_manager.operator = operator_module.HybridSlidingMemOperator(mem_manager) req_manager = object.__new__(ReqManagerForSlidingWindow) req_manager.mem_manager, req_manager.sliding_window = mem_manager, 4 - req_manager.req_to_sliding_window = torch.full((1, 8, 2, 4), -1.0) + req_manager.req_to_sliding_window = torch.full((1, 2, 4, 2, 4), -1.0) req_manager.req_to_token_indexs = torch.full((2, 1024), -1, dtype=torch.int32) req_manager.req_to_token_indexs[1, :gpu_prefix] = torch.arange(10000, 10000 + gpu_prefix, dtype=torch.int32) original_mapping = req_manager.req_to_token_indexs.clone() @@ -116,8 +116,8 @@ def load(**kwargs): torch.testing.assert_close(req_manager.req_to_token_indexs[1, :gpu_prefix], original_mapping[1, :gpu_prefix]) torch.testing.assert_close(req_manager.req_to_token_indexs[1, gpu_prefix:cpu_prefix], new_indexes) assert torch.all(req_manager.req_to_token_indexs[1, cpu_prefix:] == -1) - assert torch.all(req_manager.req_to_sliding_window[:, :4] == -1) - assert torch.all(req_manager.req_to_sliding_window[:, 4:8] == 12) + assert torch.all(req_manager.req_to_sliding_window[:, 0] == -1) + assert torch.all(req_manager.req_to_sliding_window[:, 1] == 12) assert req.shm_req.cpu_prompt_cache_len == need_tokens assert req.shm_req.shm_cur_kv_len == cpu_prefix # Dereference all matched pages, including the page already covered by diff --git a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py index fedcccabca..85e56a074a 100644 --- a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py +++ b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py @@ -41,7 +41,7 @@ def sliding_operator(monkeypatch): req_manager = object.__new__(ReqManagerForSlidingWindow) req_manager.mem_manager = manager req_manager.sliding_window = 4 - req_manager.req_to_sliding_window = torch.full((1, 8, 2, 4), -1.0) + req_manager.req_to_sliding_window = torch.full((1, 2, 4, 2, 4), -1.0) small_pages = _cpu_pages(2) small_pages.get_state_cache(0).fill_(17) radix = SimpleNamespace( @@ -89,12 +89,12 @@ def load(**kwargs): if has_tail: assert captures[0]["big_page_buffer_ids"][-1].item() == 4 torch.testing.assert_close( - req_manager.req_to_sliding_window[:, 4:8], + req_manager.req_to_sliding_window[:, 1], torch.full((1, 4, 2, 4), 10.0 + cached_tokens // 8 + page_num - 1), atol=0, rtol=0, ) - assert torch.all(req_manager.req_to_sliding_window[:, :4] == -1) + assert torch.all(req_manager.req_to_sliding_window[:, 0] == -1) @pytest.mark.parametrize("token_num", [16, 22]) @@ -182,7 +182,7 @@ def window_page(page_id): ) req_manager = object.__new__(ReqManagerForSlidingWindow) req_manager.mem_manager, req_manager.sliding_window = manager, window - req_manager.req_to_sliding_window = torch.zeros((1, request_num * window, 2, 8), dtype=config.dtype, device="cuda") + req_manager.req_to_sliding_window = torch.zeros((1, request_num, window, 2, 8), dtype=config.dtype, device="cuda") monkeypatch.setattr(g_infer_context, "req_manager", req_manager) monkeypatch.setattr( g_infer_context, @@ -241,7 +241,7 @@ def window_page(page_id): rtol=0, ) torch.testing.assert_close( - req_manager.req_to_sliding_window[:, req_idx * window : (req_idx + 1) * window], + req_manager.req_to_sliding_window[:, req_idx], torch.full(config.get_state_shape(), 200 + req_idx, dtype=config.dtype, device="cuda"), atol=0, rtol=0, From 864182e4cf8e60cec28070cd65f0efdb77d90b6e Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Mon, 7 Sep 2026 11:52:10 +0000 Subject: [PATCH 10/14] refactor: unify sliding-window runtime KV --- .../triton_kernel/sliding_window_state.py | 146 +++++--- lightllm/common/req_manager/sliding_window.py | 83 +++-- lightllm/models/gemma4/kv_layout.py | 8 +- .../layer_infer/transformer_layer_infer.py | 25 +- lightllm/models/gemma4/model.py | 19 +- .../context_attention_fwd_gemma4_mm.py | 42 +-- .../triton_kernel/sliding_window_decode.py | 18 +- .../benchmark_sliding_window_attention.py | 348 ------------------ test/kernel/test_gemma4_hybrid_graph.py | 126 ------- ...test_sliding_window_cpu_cache_attention.py | 109 ------ test/kernel/test_sliding_window_decode.py | 85 +++-- test/kernel/test_sliding_window_prefill.py | 51 ++- test/kernel/test_sliding_window_state.py | 271 +++----------- test/utils/test_sliding_cpu_cache_meta.py | 173 --------- test/utils/test_sliding_window_cache.py | 249 ------------- .../mode_backend/test_multi_level_kv_cache.py | 12 - .../model_infer/test_hybrid_state_cache.py | 163 -------- .../test_sliding_cpu_cache_loading.py | 125 ------- .../test_sliding_cpu_cache_operator.py | 259 ------------- 19 files changed, 332 insertions(+), 1980 deletions(-) delete mode 100644 test/kernel/benchmark_sliding_window_attention.py delete mode 100644 test/kernel/test_gemma4_hybrid_graph.py delete mode 100644 test/kernel/test_sliding_window_cpu_cache_attention.py delete mode 100644 test/utils/test_sliding_cpu_cache_meta.py delete mode 100644 test/utils/test_sliding_window_cache.py delete mode 100644 unit_tests/server/router/model_infer/test_hybrid_state_cache.py delete mode 100644 unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py delete mode 100644 unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py diff --git a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py index 9df3551a63..2b57246682 100644 --- a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py @@ -4,65 +4,129 @@ @triton.jit -def _commit_sliding_window_state( - LayerBuffer, +def _move_sliding_window( + Pool, BReqIdx, BSeqLen, - BQSeqLen, + BReadyCacheLen, BQStartLoc, + runtime_token_start, + stride_layer, stride_token, - stride_head, - stride_dim, - HEAD_NUM: tl.constexpr, - HEAD_DIM: tl.constexpr, WINDOW: tl.constexpr, - SCRATCH_START: tl.constexpr, - BLOCK_D: tl.constexpr, + KV_DIM: tl.constexpr, + COMPACT: tl.constexpr, + BLOCK: tl.constexpr, +): + batch_idx = tl.program_id(0) + layer_idx = tl.program_id(1).to(tl.int64) + offsets = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + req_idx = tl.load(BReqIdx + batch_idx).to(tl.int64) + history_len = tl.load(BReadyCacheLen + batch_idx).to(tl.int64) + q_start = tl.load(BQStartLoc + batch_idx).to(tl.int64) + current_start = tl.cast(runtime_token_start, tl.int64) + q_start + (batch_idx + 1) * WINDOW + if COMPACT: + end = tl.load(BSeqLen + batch_idx).to(tl.int64) + else: + end = history_len + positions = end - WINDOW + offsets // KV_DIM + canonical = req_idx * WINDOW + positions % WINDOW + active = current_start + positions - history_len + if COMPACT: + src, dst = active, canonical + else: + src, dst = canonical, active + layer_offset = layer_idx * tl.cast(stride_layer, tl.int64) + mask = (offsets < WINDOW * KV_DIM) & (positions >= 0) + values = tl.load(Pool + layer_offset + src * stride_token + offsets % KV_DIM, mask=mask, other=0) + tl.store(Pool + layer_offset + dst * stride_token + offsets % KV_DIM, values, mask=mask) + + +@torch.no_grad() +def move_sliding_window( + pool, + b_req_idx, + b_seq_len, + b_ready_cache_len, + b_q_start_loc, + window, + runtime_token_start, + compact=False, +): + """Move every owner's window between canonical rings and the prefill region.""" + assert pool.ndim == 4 and pool.stride(-1) == 1 and pool.stride(-2) == pool.shape[-1] + if not b_req_idx.numel(): + return + kv_dim = pool.shape[2] * pool.shape[3] + grid = (b_req_idx.numel(), pool.shape[0], triton.cdiv(window * kv_dim, 1024)) + _move_sliding_window[grid]( + pool, + b_req_idx, + b_seq_len, + b_ready_cache_len, + b_q_start_loc, + runtime_token_start, + pool.stride(0), + pool.stride(1), + WINDOW=window, + KV_DIM=kv_dim, + COMPACT=compact, + BLOCK=1024, + num_warps=4, + ) + + +@triton.jit +def _get_sliding_window_mem_indexes( + Out, + BReqIdx, + BSeqLen, + BQSeqLen, + BQStartLoc, + runtime_token_start, + WINDOW: tl.constexpr, + IS_PREFILL: tl.constexpr, + BLOCK: tl.constexpr, ): batch_idx = tl.program_id(0) - tail_offset = tl.program_id(1) - head_idx = tl.program_id(2) - req_idx = tl.load(BReqIdx + batch_idx) - seq_len = tl.load(BSeqLen + batch_idx) - q_len = tl.load(BQSeqLen + batch_idx) - q_start = tl.load(BQStartLoc + batch_idx) - q_offset = tl.maximum(q_len - WINDOW, 0) + tail_offset - pos = seq_len - q_len + q_offset - mask_token = q_offset < q_len - src_token = SCRATCH_START + q_start + q_offset - dst_token = req_idx * WINDOW + pos % WINDOW - dims = tl.arange(0, BLOCK_D) - mask = mask_token & (head_idx < HEAD_NUM) & (dims < HEAD_DIM) - src = src_token * stride_token + head_idx * stride_head + dims * stride_dim - dst = dst_token * stride_token + head_idx * stride_head + dims * stride_dim - value = tl.load(LayerBuffer + src, mask=mask, other=0.0) - tl.store(LayerBuffer + dst, value, mask=mask) + if IS_PREFILL: + q_start = tl.load(BQStartLoc + batch_idx).to(tl.int64) + q_len = tl.load(BQSeqLen + batch_idx) + offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + current_start = tl.cast(runtime_token_start, tl.int64) + q_start + (batch_idx + 1) * WINDOW + tl.store(Out + q_start + offsets, current_start + offsets, mask=offsets < q_len) + else: + req_idx = tl.load(BReqIdx + batch_idx).to(tl.int64) + seq_len = tl.load(BSeqLen + batch_idx).to(tl.int64) + tl.store(Out + batch_idx, req_idx * WINDOW + (seq_len - 1) % WINDOW) @torch.no_grad() -def commit_sliding_window_state( - layer_buffer, +def get_sliding_window_mem_indexes( b_req_idx, b_seq_len, b_q_seq_len, b_q_start_loc, - sliding_window, - scratch_start, + window, + runtime_token_start, + total_token_num, max_q_seq_len, + is_prefill, ): - block_d = triton.next_power_of_2(layer_buffer.shape[-1]) - # Earlier chunk tokens are only used by attention, never by the next step. - grid = (b_req_idx.shape[0], min(max_q_seq_len, sliding_window), layer_buffer.shape[1]) - _commit_sliding_window_state[grid]( - layer_buffer, + indexes = torch.empty(total_token_num, dtype=torch.int32, device=b_req_idx.device) + if not b_req_idx.numel(): + return indexes + grid = (b_req_idx.numel(), triton.cdiv(max_q_seq_len, 256) if is_prefill else 1) + _get_sliding_window_mem_indexes[grid]( + indexes, b_req_idx, b_seq_len, b_q_seq_len, b_q_start_loc, - *layer_buffer.stride(), - HEAD_NUM=layer_buffer.shape[1], - HEAD_DIM=layer_buffer.shape[2], - WINDOW=sliding_window, - SCRATCH_START=scratch_start, - BLOCK_D=block_d, + runtime_token_start, + WINDOW=window, + IS_PREFILL=is_prefill, + BLOCK=256, + num_warps=4 if is_prefill else 1, ) + return indexes diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index 541434d302..fb633144aa 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -2,7 +2,10 @@ import torch -from lightllm.common.basemodel.triton_kernel.sliding_window_state import commit_sliding_window_state +from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( + get_sliding_window_mem_indexes, + move_sliding_window, +) from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager from .hybrid_att import HybridAttentionReqManager @@ -23,26 +26,30 @@ def __init__( max_sequence_length: int, mem_manager: Optional["HybridSlidingMemoryManager"], sliding_config: "SlidingWindowCacheConfig", - scratch_token_num: int, + max_prefill_token_num: int, + prefill_microbatch_num: int = 1, ): super().__init__(max_request_num, max_sequence_length, mem_manager) self.sliding_config = sliding_config self.sliding_window = sliding_config.sliding_window - self.scratch_token_num = scratch_token_num - self.scratch_start = (max_request_num + 1) * self.sliding_window - # Attention reads history and current-chunk KV from one buffer. The - # request state is a view of its ring region, not a second allocation. + self.max_prefill_token_num = max_prefill_token_num + self.runtime_token_start = (max_request_num + 1) * self.sliding_window + # One KV pool. Prefill expands each active window to history + chunk; + # decode uses its compact request window. Both attention paths read this pool. + self.prefill_capacity = ( + max_prefill_token_num + min(max_request_num + 1, max_prefill_token_num) * self.sliding_window + ) self.sliding_kv_buffer = torch.zeros( ( sliding_config.sliding_layer_num, - self.scratch_start + scratch_token_num, + self.runtime_token_start + prefill_microbatch_num * self.prefill_capacity, 2 * sliding_config.sliding_head_num, sliding_config.sliding_head_dim, ), dtype=sliding_config.dtype, device="cuda", ) - self.req_to_sliding_window = self.sliding_kv_buffer[:, : self.scratch_start].view( + self.req_to_sliding_window = self.sliding_kv_buffer[:, : self.runtime_token_start].view( sliding_config.sliding_layer_num, max_request_num + 1, self.sliding_window, @@ -83,30 +90,48 @@ def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffer ) def prepare_sliding_window(self, infer_state): - q_token_num = infer_state.input_ids.shape[0] - assert q_token_num <= self.scratch_token_num - infer_state.sliding_window_mem_index = torch.arange( - self.scratch_start, - self.scratch_start + q_token_num, - dtype=torch.int64, - device="cuda", + token_num = infer_state.input_ids.shape[0] + infer_state.sliding_window_runtime_start = ( + self.runtime_token_start + infer_state.microbatch_index * self.prefill_capacity + ) + if infer_state.is_prefill: + assert token_num <= self.max_prefill_token_num + move_sliding_window( + self.sliding_kv_buffer, + infer_state.b_req_idx, + infer_state.b_seq_len, + infer_state.b_ready_cache_len, + infer_state.b_q_start_loc, + self.sliding_window, + infer_state.sliding_window_runtime_start, + ) + infer_state.sliding_window_mem_index = get_sliding_window_mem_indexes( + infer_state.b_req_idx, + infer_state.b_seq_len, + infer_state.b_q_seq_len, + infer_state.b_q_start_loc, + self.sliding_window, + infer_state.sliding_window_runtime_start, + token_num, + infer_state.max_q_seq_len, + infer_state.is_prefill, + ) + + def finish_prefill(self, infer_state): + # Compact all physical layers together, after every shared reader. + move_sliding_window( + self.sliding_kv_buffer, + infer_state.b_req_idx, + infer_state.b_seq_len, + infer_state.b_ready_cache_len, + infer_state.b_q_start_loc, + self.sliding_window, + infer_state.sliding_window_runtime_start, + compact=True, ) def get_layer_kv(self, layer_index: int): local_layer = self.sliding_config.get_sliding_layer_index(layer_index) - layer_buffer = self.sliding_kv_buffer[local_layer] head_num = self.sliding_config.sliding_head_num + layer_buffer = self.sliding_kv_buffer[local_layer] return layer_buffer[:, :head_num], layer_buffer[:, head_num:] - - def commit_layer_state(self, layer_index: int, infer_state): - local_layer = self.sliding_config.get_sliding_layer_index(layer_index) - commit_sliding_window_state( - layer_buffer=self.sliding_kv_buffer[local_layer], - b_req_idx=infer_state.b_req_idx, - b_seq_len=infer_state.b_seq_len, - b_q_seq_len=infer_state.b_q_seq_len, - b_q_start_loc=infer_state.b_q_start_loc, - sliding_window=self.sliding_window, - scratch_start=self.scratch_start, - max_q_seq_len=infer_state.max_q_seq_len, - ) diff --git a/lightllm/models/gemma4/kv_layout.py b/lightllm/models/gemma4/kv_layout.py index d363cb8bd0..c637fea777 100644 --- a/lightllm/models/gemma4/kv_layout.py +++ b/lightllm/models/gemma4/kv_layout.py @@ -2,13 +2,12 @@ def get_kv_cache_layout(config): - """Map Gemma's shared tail layers to physical owners and their last readers.""" + """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 = {} - last_reader = {} owners = [] for layer_index, layer_type in enumerate(layer_types): cache_map = layer_maps[layer_type] @@ -19,8 +18,7 @@ def get_kv_cache_layout(config): cache_map[layer_index] = cache_map[last_owner[layer_type]] owner = last_owner[layer_type] owners.append(owner) - last_reader[owner] = layer_index - return layer_maps, owners, last_reader + return layer_maps, owners def build_sliding_cache_config(config, tp_world_size, dtype): @@ -30,7 +28,7 @@ def build_sliding_cache_config(config, tp_world_size, dtype): assert tp_world_size > 0 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) + 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"], diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 043087f453..1e482e5665 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -18,8 +18,7 @@ class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): """ Gemma-4 decoder block. Full-attention KV stays token granular, while - sliding-attention KV is written to request-window state plus per-forward - scratch storage. + sliding attention reads one runtime KV pool with an adjustable window. """ def __init__(self, layer_num, network_config): @@ -70,13 +69,13 @@ 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_owners, last_reader = get_kv_cache_layout(network_config) + _, kv_owners = get_kv_cache_layout(network_config) kv_owner = kv_owners[layer_num] self.is_kv_shared_ = kv_owner != layer_num self.kv_share_target_layer_ = kv_owner if self.is_kv_shared_ else None - # A chunk must not overwrite the history ring until every shared-KV - # consumer has read it. This also keeps graph capture/replay layer-local. - self.commit_sliding_state_ = self.is_sliding and last_reader[kv_owner] == layer_num + self.finish_sliding_prefill_ = self.is_sliding and not any( + kind == "sliding_attention" for kind in network_config["layer_types"][layer_num + 1 :] + ) # 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 @@ -112,7 +111,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. @@ -203,10 +202,12 @@ def _context_attention_kernel( None, infer_state.b_image_token_end, sliding_window=sw, - scratch_start=infer_state.req_manager.scratch_start, + runtime_token_start=infer_state.sliding_window_runtime_start, ) - if self.commit_sliding_state_: - infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) + # The final sliding reader compacts all physical windows together. + # Graph shape probing must not mutate state; replay uses fresh metadata. + if self.finish_sliding_prefill_ and not torch.cuda.is_current_stream_capturing(): + infer_state.req_manager.finish_prefill(infer_state) return o_tensor.view(q.shape) # Full-attn layers: head_dim=512, no SWA, no image bidi — standard @@ -234,9 +235,7 @@ def _token_attention_kernel( v=_v, b_req_idx=infer_state.b_req_idx, b_seq_len=infer_state.b_seq_len, - b_q_start_loc=infer_state.b_q_start_loc, sliding_window=self.sliding_window_, - scratch_start=infer_state.req_manager.scratch_start, out=out, alloc_tensor_func=self.alloc_tensor, ) @@ -244,8 +243,6 @@ def _token_attention_kernel( o_tensor = infer_state.decode_att_state1.decode_att( q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor ) - if self.commit_sliding_state_: - infer_state.req_manager.commit_layer_state(self.layer_num_, infer_state) 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 603bec119b..eeacae6ce1 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -84,8 +84,7 @@ def _verify_params(self): f"num_hidden_layers={self.config['num_hidden_layers']}" ) if kv_shared: - # Shared layers retain the owner's scratch KV across layers. Two - # interleaved microbatches would overwrite the same scratch slots. + # Shared-KV microbatch overlap needs separate lifecycle validation. assert not ( args.enable_prefill_microbatch_overlap or args.enable_decode_microbatch_overlap ), "Gemma-4 shared sliding-window KV does not support microbatch overlap yet" @@ -108,17 +107,21 @@ def _get_sliding_cache_config(self): def _init_req_manager(self): args = get_env_start_args() create_max_seq_len = max(int(self.batch_max_tokens or 0), int(self.max_seq_length or 0)) - scratch_token_num = max( + max_prefill_token_num = max( int(self.batch_max_tokens or 0), - int(self.graph_max_batch_size or 0), int(args.prefill_cudagraph_max_handle_token or 0) if args.enable_prefill_cudagraph else 0, ) + if args.enable_tpsp_mix_mode: + max_prefill_token_num = ( + (max(1, max_prefill_token_num) + self.tp_world_size_ - 1) // self.tp_world_size_ * self.tp_world_size_ + ) self.req_manager = ReqManagerForSlidingWindow( max_request_num=self.max_req_num, max_sequence_length=create_max_seq_len, mem_manager=None, sliding_config=self._get_sliding_cache_config(), - scratch_token_num=scratch_token_num, + max_prefill_token_num=max_prefill_token_num, + prefill_microbatch_num=2 if args.enable_prefill_microbatch_overlap else 1, ) def _init_mem_manager(self): @@ -136,10 +139,8 @@ def _init_att_backend(self): # 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. + # Sliding layers read their runtime KV pool through model-local kernels. + # The framework still requires primary attention states. self.prefill_att_backend = TritonAttBackend(model=self) self.decode_att_backend = TritonAttBackend(model=self) diff --git a/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py b/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py index 80dd201dfe..bf159b42ba 100644 --- a/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py +++ b/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py @@ -59,14 +59,13 @@ def _fwd_kernel( stride_req_to_tokens_s, kv_group_num, b_prompt_cache_len, - scratch_start, H: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, USE_SLIDING_WINDOW: tl.constexpr, SLIDING_WINDOW_LEFT: tl.constexpr, - USE_RING_CACHE: tl.constexpr, + RUNTIME_TOKEN_START: tl.constexpr, ): start_m = tl.program_id(0) cur_bh = tl.program_id(1) @@ -128,26 +127,22 @@ def _fwd_kernel( k_pos = kv_start_index + start_n + offs_n # [N] k_valid = k_pos < block_end_loc - if USE_RING_CACHE: - history_loc = cur_batch_req_idx.to(tl.int64) * (SLIDING_WINDOW_LEFT + 1) + k_pos.to(tl.int64) % ( - SLIDING_WINDOW_LEFT + 1 - ) - current_loc = ( - tl.cast(scratch_start, tl.int64) + if RUNTIME_TOKEN_START is not None: + kv_loc = ( + RUNTIME_TOKEN_START + cur_batch_in_all_start_index.to(tl.int64) + + (cur_batch + 1) * (SLIDING_WINDOW_LEFT + 1) + k_pos.to(tl.int64) - prompt_cache_len.to(tl.int64) ) - kv_loc = tl.where(k_pos < prompt_cache_len, history_loc, current_loc) else: kv_loc = tl.load( Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * k_pos, mask=k_valid, other=0, ).to(tl.int64) - - off_k = kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd - k = tl.load(K + off_k, mask=k_valid[None, :], other=0.0) + k_ptr = K + kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd + k = tl.load(k_ptr, mask=k_valid[None, :], other=0.0) qk = tl.dot(q, k) if USE_SLIDING_WINDOW: @@ -172,8 +167,8 @@ def _fwd_kernel( l_i = l_i * alpha + l_ij acc = acc * alpha[:, None] - off_v = kv_loc[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd - v = tl.load(V + off_v, mask=k_valid[:, None], other=0.0) + v_ptr = V + kv_loc[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd + v = tl.load(v_ptr, mask=k_valid[:, None], other=0.0) p = p.to(v.dtype) acc = tl.dot(p, v, acc) @@ -202,7 +197,7 @@ def context_attention_fwd_gemma4_mm( req_to_token_indexs, b_image_token_end, sliding_window=(-1, -1), - scratch_start=None, + runtime_token_start=None, ): """Prefill attention with image bidirectional masking on sliding layers. @@ -213,8 +208,9 @@ def context_attention_fwd_gemma4_mm( position (in the flattened new-token layout), value is the image span's end index (in absolute request position) if the token is inside an image span, else 0. - scratch_start: When set, use request rings plus current-token scratch - for sliding KV; ``req_to_token_indexs`` is unused and may be None. + runtime_token_start: Start of the single prefill KV region, where each + request's W history tokens precede its current tokens. The token + index table is unused and may be None. """ BLOCK_M = 128 if not is_tesla() else 64 Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] @@ -242,9 +238,8 @@ def context_attention_fwd_gemma4_mm( assert int(sliding_window[1]) == 0, "sliding_window right must be 0" sliding_window_left = int(sliding_window[0]) - use_ring_cache = scratch_start is not None - if use_ring_cache: - assert use_sliding_window and sliding_window_left >= 0, "ring KV requires a finite sliding window" + if runtime_token_start is not None: + assert use_sliding_window and sliding_window_left >= 0, "runtime KV requires a finite sliding window" _fwd_kernel[grid]( q, @@ -269,18 +264,17 @@ def context_attention_fwd_gemma4_mm( o.stride(0), o.stride(1), o.stride(2), - 0 if use_ring_cache else req_to_token_indexs.stride(0), - 0 if use_ring_cache else req_to_token_indexs.stride(1), + 0 if runtime_token_start is not None else req_to_token_indexs.stride(0), + 0 if runtime_token_start is not None else req_to_token_indexs.stride(1), kv_group_num=kv_group_num, b_prompt_cache_len=b_prompt_cache_len, - scratch_start=scratch_start if use_ring_cache else 0, H=head, BLOCK_DMODEL=Lk, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, USE_SLIDING_WINDOW=use_sliding_window, SLIDING_WINDOW_LEFT=sliding_window_left, - USE_RING_CACHE=use_ring_cache, + RUNTIME_TOKEN_START=runtime_token_start, num_warps=num_warps, num_stages=num_stages, ) diff --git a/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py b/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py index 9fb30084c0..f96365b580 100644 --- a/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py +++ b/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py @@ -1,4 +1,4 @@ -"""Gemma sliding decode over a request ring and the current token's scratch KV.""" +"""Gemma sliding decode over each request's canonical KV ring.""" import torch import triton @@ -16,7 +16,6 @@ def _sliding_window_decode_stage1( V, BReqIdx, BSeqLen, - BQStartLoc, MidO, MidLogSumExp, sm_scale, @@ -38,7 +37,6 @@ def _sliding_window_decode_stage1( stride_ls, gqa_group_size, WINDOW: tl.constexpr, - SCRATCH_START: tl.constexpr, Q_HEAD_NUM: tl.constexpr, BLOCK_SEQ: tl.constexpr, BLOCK_DMODEL: tl.constexpr, @@ -57,8 +55,6 @@ def _sliding_window_decode_stage1( return req_idx = tl.load(BReqIdx + batch_idx).to(tl.int64) - q_start = tl.load(BQStartLoc + batch_idx).to(tl.int64) - scratch_token = tl.full((), SCRATCH_START, tl.int64) + q_start head_offsets = tl.arange(0, Q_HEAD_NUM) q_heads = kv_head * gqa_group_size + head_offsets q_heads = tl.where(head_offsets < gqa_group_size, q_heads, kv_head * gqa_group_size) @@ -77,7 +73,7 @@ def _sliding_window_decode_stage1( positions = tile * BLOCK_N + offs_n mask = positions < block_end token_pos = kv_start + positions - k_loc = tl.where(token_pos < seq_len - 1, req_idx * WINDOW + token_pos % WINDOW, scratch_token) + k_loc = req_idx * WINDOW + token_pos % WINDOW k = tl.load( K + k_loc[None, :] * stride_kt + kv_head * stride_kh + offs_d[:, None] * stride_kd, mask=mask[None, :], @@ -115,19 +111,17 @@ def sliding_window_decode_attention( v, b_req_idx, b_seq_len, - b_q_start_loc, sliding_window: int, - scratch_start: int, out=None, alloc_tensor_func=torch.empty, ): - """Decode one token per request without a token-to-sliding-KV index table.""" + """Decode one token per request after its current KV has been written to the ring.""" batch_size, q_head_num, head_dim = q.shape assert k.shape == v.shape and k.shape[-1] == head_dim assert head_dim in {16, 32, 64, 128, 256, 512} assert q_head_num % k.shape[1] == 0 - assert b_req_idx.shape == b_seq_len.shape == b_q_start_loc.shape == (batch_size,) - assert sliding_window > 0 and scratch_start >= sliding_window + assert b_req_idx.shape == b_seq_len.shape == (batch_size,) + assert sliding_window > 0 assert q.dtype == k.dtype == v.dtype # Keep the common GQA wrapper's launch and reduction schedule unchanged. @@ -143,7 +137,6 @@ def sliding_window_decode_attention( v, b_req_idx, b_seq_len, - b_q_start_loc, mid_o, mid_logsumexp, 1.0 / (head_dim ** 0.5), @@ -154,7 +147,6 @@ def sliding_window_decode_attention( *mid_logsumexp.stride(), group_size, WINDOW=sliding_window, - SCRATCH_START=scratch_start, Q_HEAD_NUM=max(16, triton.next_power_of_2(group_size)), BLOCK_SEQ=block_seq, BLOCK_DMODEL=head_dim, diff --git a/test/kernel/benchmark_sliding_window_attention.py b/test/kernel/benchmark_sliding_window_attention.py deleted file mode 100644 index 5c20d58b08..0000000000 --- a/test/kernel/benchmark_sliding_window_attention.py +++ /dev/null @@ -1,348 +0,0 @@ -"""Compare paged lookup with direct sliding-window addressing on identical KV. - -This is a warm-cache, single-layer microbenchmark, not a serving benchmark. -Preparation runs once per model forward, NOT once per attention layer; its -separate timing must not be multiplied by the model's sliding-layer count. -KV writes and window commits are excluded from both paths. - -Example (run only on a GPU approved for benchmarking): - python test/kernel/benchmark_sliding_window_attention.py --family both --output /tmp/sliding-attention.json -""" - -import argparse -import gc -import json -import os -import statistics -import time -from pathlib import Path -from types import SimpleNamespace - -import torch -import triton -import triton.language as tl - -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import ( - gqa_token_decode_attention_flash_decoding, -) -from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import context_attention_fwd_gemma4_mm - - -# Verified from the text_config of gemma-4-E4B-it (TP2) and gemma-4-31B-it (TP4). -MODEL_SHAPES = { - "e4b": {"tp": 2, "q_heads": 4, "kv_heads": 1, "head_dim": 256, "window": 512}, - "31b": {"tp": 4, "q_heads": 8, "kv_heads": 4, "head_dim": 256, "window": 1024}, -} - - -# Frozen pre-change preparation kernels. Keeping them here lets the benchmark -# remain usable after the second request-token table is removed from serving. -@triton.jit -def _legacy_prepare_history( - mapping, req_ids, seq_lens, q_lens, stride_req, stride_seq, WINDOW: tl.constexpr, BLOCK: tl.constexpr -): - batch, block = tl.program_id(0), tl.program_id(1) - req_idx = tl.load(req_ids + batch) - history_end = tl.load(seq_lens + batch) - tl.load(q_lens + batch) - history_start = tl.maximum(0, history_end - WINDOW) - positions = history_start + block * BLOCK + tl.arange(0, BLOCK) - tl.store( - mapping + req_idx * stride_req + positions * stride_seq, - req_idx * WINDOW + positions % WINDOW, - positions < history_end, - ) - - -@triton.jit -def _legacy_prepare_current( - mapping, - req_ids, - seq_lens, - q_lens, - q_starts, - stride_req, - stride_seq, - SCRATCH_START: tl.constexpr, - BLOCK: tl.constexpr, -): - batch, block = tl.program_id(0), tl.program_id(1) - req_idx = tl.load(req_ids + batch) - q_len = tl.load(q_lens + batch) - history_end = tl.load(seq_lens + batch) - q_len - q_start = tl.load(q_starts + batch) - offsets = block * BLOCK + tl.arange(0, BLOCK) - tl.store( - mapping + req_idx * stride_req + (history_end + offsets) * stride_seq, - SCRATCH_START + q_start + offsets, - offsets < q_len, - ) - - -def _graph_timing(fn, unroll, samples): - """Capture repeated calls so Python launch overhead is outside GPU timing.""" - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - fn() - fn() - stream.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - for _ in range(unroll): - fn() - graph.replay() - stream.synchronize() - timings = [] - for _ in range(samples): - start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) - start.record() - graph.replay() - end.record() - end.synchronize() - timings.append(start.elapsed_time(end) * 1000 / unroll) - return {"median_us": statistics.median(timings), "min_us": min(timings), "max_us": max(timings)} - - -def _eager_timing(fn, iterations, samples): - """Report enqueue time separately from synchronized whole-call latency. - - Enqueue time includes Python, allocation and driver calls, and can include - queue backpressure. It is not an isolated measurement of CPU computation. - """ - enqueue, wall = [], [] - for _ in range(samples): - torch.cuda.synchronize() - start = time.perf_counter_ns() - for _ in range(iterations): - fn() - submitted = time.perf_counter_ns() - torch.cuda.synchronize() - completed = time.perf_counter_ns() - enqueue.append((submitted - start) / iterations / 1000) - wall.append((completed - start) / iterations / 1000) - return {"enqueue_median_us": statistics.median(enqueue), "wall_median_us": statistics.median(wall)} - - -def _accuracy(actual, expected, atol, rtol): - # Limit validation workspace so the 8192-token cases do not temporarily - # allocate several additional full-sized FP32 attention outputs. - max_abs, squared_error, bitwise_equal = 0.0, 0.0, True - for actual_chunk, expected_chunk in zip(actual.flatten().split(1 << 20), expected.flatten().split(1 << 20)): - torch.testing.assert_close(actual_chunk, expected_chunk, atol=atol, rtol=rtol) - difference = actual_chunk.float() - expected_chunk.float() - max_abs = max(max_abs, difference.abs().max().item()) - squared_error += difference.square().sum().item() - bitwise_equal = bitwise_equal and torch.equal(actual_chunk, expected_chunk) - return {"max_abs": max_abs, "rms": (squared_error / actual.numel()) ** 0.5, "bitwise_equal": bitwise_equal} - - -@torch.inference_mode() -def _benchmark_case(family, phase, batch, q_len, args): - shape = MODEL_SHAPES[family] - q_heads, kv_heads, dim, window = (shape[name] for name in ["q_heads", "kv_heads", "head_dim", "window"]) - request_slots = batch + 1 - scratch_start = request_slots * window - token_num = batch * q_len - max_seq_len = args.history + 7 * (batch - 1) + q_len - backing_bytes = (scratch_start + token_num) * 2 * kv_heads * dim * 2 - mapping_bytes = request_slots * max_seq_len * 4 - q_and_outputs_bytes = 3 * token_num * q_heads * dim * 2 - legacy_blocks = 128 if batch <= 16 else 64 if batch <= 64 else 32 - decode_workspace_bytes = batch * q_heads * legacy_blocks * (dim * 2 + 4) if phase == "decode" else 0 - validation_bytes = min(token_num * q_heads * dim, 1 << 20) * 16 - estimated_bytes = backing_bytes + mapping_bytes + q_and_outputs_bytes + decode_workspace_bytes + validation_bytes - if estimated_bytes > args.max_case_mib * 1024 ** 2: - raise ValueError( - f"case needs at least {estimated_bytes / 1024 ** 2:.1f} MiB; increase --max-case-mib explicitly" - ) - - req_ids = torch.arange(batch, 0, -1, dtype=torch.int32, device="cuda") - history_lens = args.history + torch.arange(batch, dtype=torch.int32, device="cuda") * 7 - seq_lens = history_lens + q_len - q_lens = torch.full((batch,), q_len, dtype=torch.int32, device="cuda") - q_starts = torch.arange(batch, dtype=torch.int32, device="cuda") * q_len - mapping = torch.full((request_slots, max_seq_len), -1, dtype=torch.int32, device="cuda") - # Both paths read precisely this tensor: only the address computation changes. - backing = torch.randn((scratch_start + token_num, 2 * kv_heads, dim), dtype=torch.bfloat16, device="cuda") - k, v = backing[:, :kv_heads], backing[:, kv_heads:] - q = torch.randn((token_num, q_heads, dim), dtype=torch.bfloat16, device="cuda") - old_out, new_out = torch.empty_like(q), torch.empty_like(q) - image_ends = torch.zeros((token_num,), dtype=torch.int32, device="cuda") - - def prepare(): - _legacy_prepare_history[(batch, triton.cdiv(window, 256))]( - mapping, req_ids, seq_lens, q_lens, *mapping.stride(), WINDOW=window, BLOCK=256 - ) - _legacy_prepare_current[(batch, triton.cdiv(q_len, 256))]( - mapping, req_ids, seq_lens, q_lens, q_starts, *mapping.stride(), SCRATCH_START=scratch_start, BLOCK=256 - ) - - if phase == "prefill": - - def legacy_attention(): - context_attention_fwd_gemma4_mm( - q, - k, - v, - old_out, - req_ids, - q_starts, - seq_lens, - history_lens, - q_len, - mapping, - image_ends, - sliding_window=(window - 1, 0), - ) - - def direct_attention(): - context_attention_fwd_gemma4_mm( - q, - k, - v, - new_out, - req_ids, - q_starts, - seq_lens, - history_lens, - q_len, - None, - image_ends, - sliding_window=(window - 1, 0), - scratch_start=scratch_start, - ) - - else: - from lightllm.models.gemma4.triton_kernel.sliding_window_decode import sliding_window_decode_attention - - infer_state = SimpleNamespace( - batch_size=batch, - req_manager=SimpleNamespace(req_to_token_indexs=mapping), - b_req_idx=req_ids, - b_seq_len=seq_lens, - max_kv_seq_len=max_seq_len, - ) - - def legacy_attention(): - gqa_token_decode_attention_flash_decoding(q, infer_state, k, v, out=old_out, sliding_window=(window - 1, 0)) - - def direct_attention(): - sliding_window_decode_attention( - q, - k, - v, - req_ids, - seq_lens, - q_starts, - sliding_window=window, - scratch_start=scratch_start, - out=new_out, - ) - - def legacy_step(): - prepare() - legacy_attention() - - legacy_step() - direct_attention() - torch.cuda.synchronize() - accuracy = _accuracy(new_out, old_out, atol=args.atol, rtol=args.rtol) - functions = { - "legacy_prepare_once": prepare, - "legacy_attention_only": legacy_attention, - "legacy_prepare_plus_one_attention": legacy_step, - "direct_attention_only": direct_attention, - } - graph = {name: _graph_timing(fn, args.graph_unroll, args.samples) for name, fn in functions.items()} - eager = {name: _eager_timing(fn, args.eager_iterations, args.samples) for name, fn in functions.items()} - return { - "family": family, - "phase": phase, - "batch_size": batch, - "q_len": q_len, - "history_min": args.history, - "history_max": max_seq_len - q_len, - **shape, - "backing_mib": backing_bytes / 1024 ** 2, - "removed_index_table_mib": mapping_bytes / 1024 ** 2, - "accuracy": accuracy, - "graph": graph, - "eager": eager, - } - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--family", choices=["e4b", "31b", "both"], default="both") - parser.add_argument("--phase", choices=["prefill", "decode", "both"], default="both") - parser.add_argument("--device", type=int, default=0, help="logical CUDA device within CUDA_VISIBLE_DEVICES") - parser.add_argument("--history", type=int, default=32768) - parser.add_argument("--prefill-lengths", type=int, nargs="+", default=[512, 4096, 8192]) - parser.add_argument("--decode-batches", type=int, nargs="+", default=[1, 8, 32]) - parser.add_argument("--graph-unroll", type=int, default=16) - parser.add_argument("--eager-iterations", type=int, default=20) - parser.add_argument("--samples", type=int, default=5) - parser.add_argument("--max-case-mib", type=float, default=256) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--atol", type=float, default=0.005) - parser.add_argument("--rtol", type=float, default=0.02) - parser.add_argument("--output", type=Path) - args = parser.parse_args() - if ( - args.history < 0 - or min(args.prefill_lengths + args.decode_batches + [args.graph_unroll, args.eager_iterations, args.samples]) - < 1 - ): - parser.error("history must be nonnegative; lengths, batches, unroll, iterations and samples must be positive") - torch.cuda.set_device(args.device) - torch.manual_seed(args.seed) - report = { - "scope": "Shared-GPU-capable warm-cache attention microbenchmark, not end-to-end serving throughput.", - "notes": [ - "Both paths use identical ring/scratch KV, Q and metadata; only attention addressing differs.", - "Legacy preparation is once per model forward, not once per attention layer.", - "KV writes, window commit, linear/full layers, scheduler and CPU-cache operations are excluded.", - "Eager enqueue time includes Python/allocator/driver calls and possible queue backpressure.", - "CUDA graph timing removes host enqueue overhead; shared GPU contention can still affect results.", - ], - "gpu": torch.cuda.get_device_name(args.device), - "torch": torch.__version__, - "triton": triton.__version__, - "cuda_visible_devices": os.getenv("CUDA_VISIBLE_DEVICES"), - "autotune_level": os.getenv("LIGHTLLM_TRITON_AUTOTUNE_LEVEL", "0"), - "arguments": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()}, - "cases": [], - } - print( - "family phase batch q old-att graph-us new-att graph-us prepare graph-us old-step graph-us", - flush=True, - ) - for family in MODEL_SHAPES if args.family == "both" else [args.family]: - cases = [] - if args.phase in ["prefill", "both"]: - cases.extend(("prefill", 1, length) for length in args.prefill_lengths) - if args.phase in ["decode", "both"]: - cases.extend(("decode", batch, 1) for batch in args.decode_batches) - for phase, batch, q_len in cases: - result = _benchmark_case(family, phase, batch, q_len, args) - report["cases"].append(result) - timings = result["graph"] - print( - f"{family:6} {phase:7} {batch:5} {q_len:5} " - f"{timings['legacy_attention_only']['median_us']:18.3f} " - f"{timings['direct_attention_only']['median_us']:17.3f} " - f"{timings['legacy_prepare_once']['median_us']:17.3f} " - f"{timings['legacy_prepare_plus_one_attention']['median_us']:18.3f}", - flush=True, - ) - gc.collect() - torch.cuda.empty_cache() - if args.output: - args.output.write_text(json.dumps(report, indent=2) + "\n") - print(f"JSON report: {args.output}", flush=True) - else: - print(json.dumps(report, indent=2)) - - -if __name__ == "__main__": - main() diff --git a/test/kernel/test_gemma4_hybrid_graph.py b/test/kernel/test_gemma4_hybrid_graph.py deleted file mode 100644 index 0720af3411..0000000000 --- a/test/kernel/test_gemma4_hybrid_graph.py +++ /dev/null @@ -1,126 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState -from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager -from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig -from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo -from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - -def _int_tensor(values): - return torch.tensor(values, device="cuda", dtype=torch.int32) - - -def _layer(is_sliding, shared): - layer = object.__new__(Gemma4TransformerLayerInfer) - layer.is_sliding, layer.is_kv_shared_ = is_sliding, shared - layer.layer_num_ = (4 if shared else 2) if is_sliding else (3 if shared else 1) - layer.kv_share_target_layer_ = (2 if is_sliding else 1) if shared else None - layer.tp_q_head_num_, layer.head_dim_ = (4, 256) if is_sliding else (8, 512) - layer.commit_sliding_state_ = False # Isolate graph metadata and KV reads from snapshot/commit tests. - layer.alloc_tensor = lambda shape, dtype, device="cuda": torch.empty(shape, dtype=dtype, device=device) - return layer - - -def _state(model, req_manager, req_ids, seq_lengths, q_starts): - state = Gemma4InferStateInfo() - state.req_manager, state.mem_manager = req_manager, req_manager.mem_manager - state.is_prefill = False - state.batch_size, state.max_q_seq_len = len(req_ids), 1 - state.max_kv_seq_len = req_manager.req_to_token_indexs.shape[1] - state.b_req_idx, state.b_seq_len = _int_tensor(req_ids), _int_tensor(seq_lengths) - state.input_ids = torch.zeros(len(req_ids), device="cuda", dtype=torch.int64) - backend = SimpleNamespace(model=model) - state.decode_att_state = TritonDecodeAttState(backend=backend, infer_state=state) - state.decode_att_state1 = TritonDecodeAttState(backend=backend, infer_state=state) - state.init_some_extra_state(model) - # Independently vary scratch locations to expose stale graph metadata. - state.b_q_start_loc = _int_tensor(q_starts) - state.init_att_state() - return state - - -@pytest.mark.parametrize("window", [512, 1024]) -@pytest.mark.parametrize("shared", [False, True]) -def test_gemma_hybrid_decode_graph_copies_state_without_replacing_full_token_table(monkeypatch, window, shared): - from lightllm.common.triton_utils import autotuner - - monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.CLOSE_AUTOTUNE) - torch.manual_seed(42) - req_slots, batch_size, max_seq_len = 6, 4, 3 * window + 32 - config = SlidingWindowCacheConfig({0: 0, 2: 1, 4: 1}, {1: 0, 3: 0}, window, 1, 256, 2, 512, torch.bfloat16) - # Skip launch-time distributed/profile setup, retaining the real cache access methods. - mem_manager = object.__new__(HybridSlidingMemoryManager) - mem_manager.sliding_config, mem_manager.head_num = config, config.full_head_num - mem_manager.kv_buffer = torch.randn((1, 1024, 4, 512), device="cuda", dtype=config.dtype) - manager = object.__new__(ReqManagerForSlidingWindow) - manager.mem_manager, manager.sliding_config = mem_manager, config - manager.sliding_window, manager.scratch_token_num = window, batch_size - manager.scratch_start = req_slots * window - manager.sliding_kv_buffer = torch.randn( - (2, manager.scratch_start + batch_size, 2, 256), device="cuda", dtype=config.dtype - ) - manager.req_to_sliding_window = manager.sliding_kv_buffer[:, : manager.scratch_start].view( - 2, req_slots, window, 2, 256 - ) - manager.req_to_token_indexs = torch.randint(1024, (req_slots, max_seq_len), device="cuda", dtype=torch.int32) - original_table = manager.req_to_token_indexs - expected_table = original_table.clone() - cos = torch.ones((max_seq_len, 128), device="cuda", dtype=config.dtype) - sin = torch.zeros_like(cos) - model = SimpleNamespace( - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda is_draft: 0), - is_mtp_draft_model=False, - _cos_cached_sliding=cos, - _sin_cached_sliding=sin, - _cos_cached_full=cos, - _sin_cached_full=sin, - ) - captured = _state(model, manager, [0, 2, 5, 5], [window + 5, 17, 2, 2], [0, 1, 2, 3]) - captured.is_cuda_graph = True - sliding, full = _layer(True, shared), _layer(False, shared) - sliding.sliding_window_, full.sliding_window_ = window, 0 - q_sliding = torch.randn((batch_size, 4, 256), device="cuda", dtype=config.dtype) - q_full = torch.randn((batch_size, 8, 512), device="cuda", dtype=config.dtype) - - def forward(state): - return ( - sliding._token_attention_kernel(q_sliding, state, None), - full._token_attention_kernel(q_full, state, None), - ) - - forward(captured) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_outputs = forward(captured) - copied_fields = ("b_req_idx", "b_seq_len", "b_q_start_loc", "position_ids", "sliding_window_mem_index") - captured_ptrs = {name: getattr(captured, name).data_ptr() for name in copied_fields} - new_state = _state(model, manager, [3, 1, 5, 5], [2 * window + 11, 1, 2, 2], [2, 0, 1, 3]) - captured.copy_for_cuda_graph(new_state) - q_sliding.mul_(0.5) - q_full.mul_(0.75) - manager.sliding_kv_buffer.mul_(0.75) - mem_manager.kv_buffer.mul_(0.5) - graph.replay() - eager_outputs = forward(new_state) - - for actual, expected in zip(graph_outputs, eager_outputs): - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - assert torch.isfinite(actual).all() - for name in copied_fields: - assert getattr(captured, name).data_ptr() == captured_ptrs[name] - torch.testing.assert_close(getattr(captured, name), getattr(new_state, name), atol=0, rtol=0) - for state in (captured, new_state): - assert state.decode_att_state.infer_state is state - assert state.decode_att_state1.infer_state is state - assert state.req_manager is manager - assert state.req_manager.req_to_token_indexs is original_table - assert not hasattr(manager, "req_to_sliding_window_indexs") - torch.testing.assert_close(original_table, expected_table, atol=0, rtol=0) diff --git a/test/kernel/test_sliding_window_cpu_cache_attention.py b/test/kernel/test_sliding_window_cpu_cache_attention.py deleted file mode 100644 index 1de02e8b7d..0000000000 --- a/test/kernel/test_sliding_window_cpu_cache_attention.py +++ /dev/null @@ -1,109 +0,0 @@ -from types import SimpleNamespace - -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, -) -from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig, SlidingWindowStateCacheManager -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") - - -@pytest.mark.parametrize("window,history_len", [(512, 256), (512, 512), (512, 544), (1024, 1056)]) -@pytest.mark.parametrize("q_len", [1, 31]) -def test_cpu_window_checkpoint_resumes_attention_exactly(window, history_len, q_len): - torch.manual_seed(42) - page_size, head_dim, req_idx = 512, 64, 1 - seq_len = history_len + q_len - # Logical layer 3 shares layer 2's KV: only physical owners are stored. - config = SlidingWindowCacheConfig({0: 0, 2: 1, 3: 1}, {1: 0}, window, 1, head_dim, 1, head_dim, torch.bfloat16) - reference = torch.randn((2, seq_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) - full_kv = torch.randn((1, history_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) - expected_full_kv = full_kv.clone() - endpoints = list(range(page_size, history_len + 1, page_size)) - if history_len % page_size: - endpoints.append(history_len) - pages = SlidingWindowStateCacheManager(len(endpoints), config) - for page_id, endpoint in enumerate(endpoints): - positions = torch.arange(max(0, endpoint - window), endpoint, device="cuda") - pages.state_cache[page_id, :, positions % window] = reference[:, positions] - - int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - mem_indexes = torch.full((len(endpoints) * page_size,), -1, device="cuda", dtype=torch.int32) - mem_indexes[:history_len] = torch.arange(history_len, device="cuda", dtype=torch.int32) - page_ids = int_tensor(list(range(len(endpoints)))) - cpu_cache = torch.zeros( - (len(endpoints), 1, 1, 1, config.get_cpu_cache_big_page_bytes(page_size, 1)), - dtype=torch.uint8, - pin_memory=True, - ) - copy_args = dict( - mem_indexes=mem_indexes, - page_indexes=page_ids, - big_page_buffer_ids=page_ids, - gpu_full_att_kv_state=full_kv, - gpu_sliding_state=pages.state_cache, - cpu_cache_tensor=cpu_cache, - tp_rank=0, - tp_world_size=1, - big_page_token_num=page_size, - sliding_config=config, - ) - copy_kv_buffer_to_cpu_cache(page_readies=torch.zeros_like(page_ids, dtype=torch.bool), **copy_args) - full_kv.fill_(-7) - pages.state_cache.fill_(-9) - copy_cpu_cache_to_kv_buffer(**copy_args) - torch.testing.assert_close(full_kv, expected_full_kv, atol=0, rtol=0) - - manager = object.__new__(ReqManagerForSlidingWindow) - manager.sliding_config, manager.sliding_window = config, window - manager.scratch_token_num, manager.scratch_start = q_len, 3 * window - manager.mem_manager = SimpleNamespace(linear_att_big_page_buffers=pages) - manager.sliding_kv_buffer = torch.zeros( - (2, manager.scratch_start + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16 - ) - manager.req_to_sliding_window = manager.sliding_kv_buffer[:, : manager.scratch_start].view( - 2, 3, window, 2, head_dim - ) - manager.restore_big_page_state(len(endpoints) - 1, SimpleNamespace(req_idx=req_idx)) - manager.sliding_kv_buffer[:, manager.scratch_start :] = reference[:, history_len:] - state = SimpleNamespace( - input_ids=int_tensor([0] * q_len), - b_req_idx=int_tensor([req_idx]), - b_seq_len=int_tensor([seq_len]), - b_q_seq_len=int_tensor([q_len]), - b_q_start_loc=int_tensor([0]), - max_q_seq_len=q_len, - ) - manager.prepare_sliding_window(state) - q = torch.randn((q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) - reference_indexes = torch.arange(seq_len, device="cuda", dtype=torch.int32).expand(3, -1) - image_end = int_tensor([0] * q_len) - for layer_index in [0, 2, 3]: - physical_layer = config.get_sliding_layer_index(layer_index) - actual, expected = torch.empty_like(q), torch.empty_like(q) - for kv, indexes, output, scratch in [ - (manager.sliding_kv_buffer[physical_layer], None, actual, manager.scratch_start), - (reference[physical_layer], reference_indexes, expected, None), - ]: - context_attention_fwd_gemma4_mm( - q, - kv[:, :1], - kv[:, 1:], - output, - state.b_req_idx, - state.b_q_start_loc, - state.b_seq_len, - int_tensor([history_len]), - q_len, - indexes, - image_end, - sliding_window=(window - 1, 0), - scratch_start=scratch, - ) - torch.testing.assert_close(actual, expected, atol=0, rtol=0) diff --git a/test/kernel/test_sliding_window_decode.py b/test/kernel/test_sliding_window_decode.py index 50504674b5..774b6494f8 100644 --- a/test/kernel/test_sliding_window_decode.py +++ b/test/kernel/test_sliding_window_decode.py @@ -15,15 +15,14 @@ def _int_tensor(values, dtype=torch.int32): return torch.tensor(values, device="cuda", dtype=dtype) -def _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start): +def _table_reference(q, k, v, req_ids, seq_lengths, window): # Only the visible suffix matters. Rebase very long sequences to avoid # allocating a token table proportional to their virtual token positions. indexes = torch.zeros((max(req_ids) + 1, window), device="cuda", dtype=torch.int64) visible_lengths = [min(length, window) for length in seq_lengths] - for req_idx, length, q_start, visible_len in zip(req_ids, seq_lengths, q_starts, visible_lengths): + for req_idx, length, visible_len in zip(req_ids, seq_lengths, visible_lengths): positions = torch.arange(length - visible_len, length, device="cuda", dtype=torch.int64) indexes[req_idx, :visible_len] = req_idx * window + positions % window - indexes[req_idx, visible_len - 1] = scratch_start + q_start state = SimpleNamespace( batch_size=len(req_ids), b_req_idx=_int_tensor(req_ids), @@ -31,12 +30,12 @@ def _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_star max_kv_seq_len=window, req_manager=SimpleNamespace(req_to_token_indexs=indexes), ) - kv_heads = kv.shape[1] // 2 return gqa_token_decode_attention_flash_decoding( q=q, infer_state=state, - cache_k=kv[:, :kv_heads], - cache_v=kv[:, kv_heads:], + # The common kernel requires equal K/V strides and contiguous head_dim. + cache_k=k.contiguous(), + cache_v=v.contiguous(), out=torch.empty_like(q), sliding_window=(window - 1, 0), ) @@ -57,95 +56,105 @@ def test_formula_decode_matches_table_gqa_exactly(dtype, window, q_heads, kv_hea torch.manual_seed(42) req_ids = [6, 0, 4, 2, 7, 1] seq_lengths = [1, 2, window - 1, window, window + 1, 3 * window + 7] - q_starts = [7, 1, 11, 4, 9, 2] - scratch_start = 9 * window - kv = torch.randn((scratch_start + 12, 2 * kv_heads, head_dim), device="cuda", dtype=dtype) + runtime = torch.randn((9 * window, 2 * kv_heads, head_dim), device="cuda", dtype=dtype) q = torch.randn((len(req_ids), q_heads, head_dim), device="cuda", dtype=dtype) output = torch.empty_like(q) actual = sliding_window_decode_attention( q, - kv[:, :kv_heads], - kv[:, kv_heads:], + runtime[:, :kv_heads], + runtime[:, kv_heads:], _int_tensor(req_ids), _int_tensor(seq_lengths), - _int_tensor(q_starts), window, - scratch_start, out=output, ) - expected = _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start) + expected = _table_reference(q, runtime[:, :kv_heads], runtime[:, kv_heads:], req_ids, seq_lengths, window) assert actual is output torch.testing.assert_close(actual, expected, atol=0, rtol=0) @pytest.mark.parametrize("batch_size", [1, 17, 65]) def test_formula_decode_preserves_gqa_batch_schedule(batch_size): - window, head_dim, scratch_start = 512, 64, (batch_size + 1) * 512 + window, head_dim = 512, 64 req_ids = list(reversed(range(batch_size))) seq_lengths = [2 * window + i + 1 for i in range(batch_size)] - q_starts = list(range(batch_size)) - kv = torch.randn((scratch_start + batch_size, 2, head_dim), device="cuda", dtype=torch.bfloat16) + runtime = torch.randn(((batch_size + 1) * window, 2, head_dim), device="cuda", dtype=torch.bfloat16) q = torch.randn((batch_size, 4, head_dim), device="cuda", dtype=torch.bfloat16) actual = sliding_window_decode_attention( q, - kv[:, :1], - kv[:, 1:], + runtime[:, :1], + runtime[:, 1:], _int_tensor(req_ids), _int_tensor(seq_lengths), - _int_tensor(q_starts), window, - scratch_start, ) - expected = _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start) + expected = _table_reference(q, runtime[:, :1], runtime[:, 1:], req_ids, seq_lengths, window) torch.testing.assert_close(actual, expected, atol=0, rtol=0) def test_formula_decode_supports_int64_virtual_token_positions(): - window, scratch_start = 512, 4 * 512 - req_ids, seq_lengths, q_starts = [2, 0], [2 ** 31 + 17, 2 ** 32 + 31], [0, 1] - kv = torch.randn((scratch_start + 2, 2, 64), device="cuda", dtype=torch.bfloat16) + window = 512 + req_ids, seq_lengths = [2, 0], [2 ** 31 + 17, 2 ** 32 + 31] + runtime = torch.randn((4 * window, 2, 64), device="cuda", dtype=torch.bfloat16) q = torch.randn((2, 4, 64), device="cuda", dtype=torch.bfloat16) actual = sliding_window_decode_attention( q, - kv[:, :1], - kv[:, 1:], + runtime[:, :1], + runtime[:, 1:], _int_tensor(req_ids), _int_tensor(seq_lengths, dtype=torch.int64), - _int_tensor(q_starts), window, - scratch_start, ) - expected = _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start) + expected = _table_reference(q, runtime[:, :1], runtime[:, 1:], req_ids, seq_lengths, window) torch.testing.assert_close(actual, expected, atol=0, rtol=0) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("window", [512, 1024]) def test_formula_decode_cuda_graph_replay_with_changed_requests_and_padding(dtype, window): - scratch_start = 8 * window - kv = torch.randn((scratch_start + 4, 4, 256), device="cuda", dtype=dtype) + runtime = torch.randn((8 * window, 4, 256), device="cuda", dtype=dtype) q = torch.randn((4, 8, 256), device="cuda", dtype=dtype) b_req = _int_tensor([4, 1, 7, 7]) b_seq = _int_tensor([window + 7, 3, 2, 2]) - b_start = _int_tensor([3, 0, 1, 2]) out = torch.empty_like(q) def forward(): - sliding_window_decode_attention(q, kv[:, :2], kv[:, 2:], b_req, b_seq, b_start, window, scratch_start, out=out) + sliding_window_decode_attention(q, runtime[:, :2], runtime[:, 2:], b_req, b_seq, window, out=out) forward() torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): forward() - req_ids, seq_lengths, q_starts = [2, 5, 7, 7], [2 * window + 3, 1, 2, 2], [1, 3, 2, 0] + req_ids, seq_lengths = [2, 5, 7, 7], [2 * window + 3, 1, 2, 2] b_req.copy_(_int_tensor(req_ids)) b_seq.copy_(_int_tensor(seq_lengths)) - b_start.copy_(_int_tensor(q_starts)) q.mul_(0.5) - kv.mul_(0.75) + runtime.mul_(0.75) graph.replay() - expected = _table_reference(q, kv, req_ids, seq_lengths, q_starts, window, scratch_start) + expected = _table_reference(q, runtime[:, :2], runtime[:, 2:], req_ids, seq_lengths, window) # Padding may share the hold request ID; its outputs are intentionally discarded. torch.testing.assert_close(out[:2], expected[:2], atol=0, rtol=0) assert torch.isfinite(out).all() + + +def test_formula_decode_reads_independently_strided_kv_without_updating_runtime(): + window, kv_heads, head_dim = 512, 2, 64 + req_ids, seq_lengths = [2, 0], [window + 17, 1] + k = torch.randn((4 * window, kv_heads, head_dim * 2), device="cuda", dtype=torch.bfloat16)[..., ::2] + v = torch.randn((kv_heads, 4 * window, head_dim), device="cuda", dtype=torch.bfloat16).transpose(0, 1) + original_k, original_v = k.clone(), v.clone() + assert k.stride() != v.stride() + q = torch.randn((2, 8, head_dim), device="cuda", dtype=torch.bfloat16) + actual = sliding_window_decode_attention( + q, + k, + v, + _int_tensor(req_ids), + _int_tensor(seq_lengths), + window, + ) + expected = _table_reference(q, k, v, req_ids, seq_lengths, window) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(k, original_k, atol=0, rtol=0) + torch.testing.assert_close(v, original_v, atol=0, rtol=0) diff --git a/test/kernel/test_sliding_window_prefill.py b/test/kernel/test_sliding_window_prefill.py index 16e0d6d726..3d72864502 100644 --- a/test/kernel/test_sliding_window_prefill.py +++ b/test/kernel/test_sliding_window_prefill.py @@ -6,14 +6,14 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def _compare_ring_and_paged(window, q_len, dtype, head_dim=64, image_span=None): +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 = [47 + q_lens[1], 11, 83 + q_lens[1] + q_lens[0]] + 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 @@ -25,14 +25,14 @@ def _compare_ring_and_paged(window, q_len, dtype, head_dim=64, image_span=None): mapping[req_id, :length] = shuffled_indexes[offset : offset + length].to(torch.int32) offset += length - scratch_start = req_slots * window + 13 - runtime = torch.full((scratch_start + query_num, 2 * kv_heads, head_dim), -3, device="cuda", dtype=dtype) - for req_id, history, length, start in zip(req_ids, histories, lengths, starts): - old_positions = torch.arange(max(0, history - window), history, device="cuda") - runtime[req_id * window + old_positions % window] = reference[mapping[req_id, old_positions].long()] - runtime[scratch_start + start : scratch_start + start + length - history] = reference[ - mapping[req_id, history:length].long() - ] + runtime_start = req_slots * window + runtime = torch.full( + (runtime_start + query_num + len(req_ids) * window, 2 * kv_heads, head_dim), -3, device="cuda", dtype=dtype + ) + for batch, (req_id, history, length, start) in enumerate(zip(req_ids, histories, lengths, starts)): + positions = torch.arange(max(0, history - window), length, device="cuda") + current_start = runtime_start + start + (batch + 1) * window + runtime[current_start + positions - history] = reference[mapping[req_id, positions].long()] image_ends = torch.zeros(query_num, device="cuda", dtype=torch.int32) if image_span is not None: @@ -63,7 +63,7 @@ def _compare_ring_and_paged(window, q_len, dtype, head_dim=64, image_span=None): v=runtime[:, kv_heads:], o=actual, req_to_token_indexs=None, - scratch_start=scratch_start, + runtime_token_start=runtime_start, **kwargs, ) torch.testing.assert_close(actual, expected, atol=0, rtol=0) @@ -86,25 +86,26 @@ def _compare_ring_and_paged(window, q_len, dtype, head_dim=64, image_span=None): @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_ring_prefill_matches_paged_history_and_scratch(window, q_len, dtype): - _compare_ring_and_paged(window, q_len, dtype) +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_ring_prefill_preserves_image_bidirectional_mask(window, dtype, image_span): +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_ring_and_paged(window, 384, dtype, head_dim=256, image_span=image_span) + _compare_runtime_and_paged(window, 384, dtype, head_dim=256, image_span=image_span) @pytest.mark.parametrize("window", [512, 1024]) -def test_ring_prefill_cuda_graph_replay_reads_updated_request_metadata(window): +def test_runtime_prefill_cuda_graph_replay_reads_updated_request_metadata(window): torch.manual_seed(43) req_slots, head_dim, max_q_len = 6, 64, window + 33 - scratch_start, query_num = req_slots * window + 11, 2 * (window + 64) + 97 - runtime = torch.randn((scratch_start + query_num, 4, head_dim), device="cuda", dtype=torch.bfloat16) + query_num = 2 * (window + 64) + 97 + runtime_start = req_slots * window + runtime = torch.randn((runtime_start + query_num + 2 * window, 4, head_dim), device="cuda", dtype=torch.bfloat16) q = torch.randn((query_num, 4, head_dim), device="cuda", dtype=torch.bfloat16) out = torch.full_like(q, -11) int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) @@ -127,7 +128,7 @@ def test_ring_prefill_cuda_graph_replay_reads_updated_request_metadata(window): ) def forward(): - context_attention_fwd_gemma4_mm(o=out, req_to_token_indexs=None, scratch_start=scratch_start, **kwargs) + context_attention_fwd_gemma4_mm(o=out, req_to_token_indexs=None, runtime_token_start=runtime_start, **kwargs) forward() torch.cuda.synchronize() @@ -147,15 +148,11 @@ def forward(): graph.replay() # Materialize a table only for the independent old-path reference, after - # changing every piece of GPU metadata used by the captured ring kernel. + # changing every piece of GPU metadata used by the captured runtime kernel. mapping = torch.full((req_slots, max(lengths)), -1, device="cuda", dtype=torch.int32) - for req_id, history, length, start in zip(req_ids, histories, lengths, starts): - positions = torch.arange(length, device="cuda", dtype=torch.int32) - mapping[req_id, :length] = torch.where( - positions < history, - req_id * window + positions % window, - scratch_start + start + positions - history, - ) + for batch, (req_id, history, length, start) in enumerate(zip(req_ids, histories, lengths, starts)): + positions = torch.arange(max(0, history - window), length, device="cuda", dtype=torch.int32) + mapping[req_id, positions.long()] = runtime_start + start + (batch + 1) * window + positions - history expected = torch.full_like(q, -11) context_attention_fwd_gemma4_mm(o=expected, req_to_token_indexs=mapping, **kwargs) # Include gaps to verify the captured grid respects the new query lengths. diff --git a/test/kernel/test_sliding_window_state.py b/test/kernel/test_sliding_window_state.py index c01a00a09f..a99fb4567b 100644 --- a/test/kernel/test_sliding_window_state.py +++ b/test/kernel/test_sliding_window_state.py @@ -1,234 +1,73 @@ -from types import SimpleNamespace - import pytest import torch -from lightllm.common.basemodel.triton_kernel.sliding_window_state import commit_sliding_window_state -from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow -from lightllm.common.kv_cache_mem_manager.operator.hybrid_sliding import HybridSlidingMemOperator -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig, SlidingWindowStateCacheManager -from lightllm.models.gemma4.kv_layout import get_kv_cache_layout -from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer -from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import context_attention_fwd_gemma4_mm +from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( + get_sliding_window_mem_indexes, + move_sliding_window, +) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_request_window_view_shares_storage_and_has_no_second_token_table(monkeypatch): - monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.ReqSamplingParamsManager", lambda size: None) - config = SlidingWindowCacheConfig({0: 0, 2: 1}, {1: 0}, 32, 1, 64, 1, 64, torch.bfloat16) - manager = ReqManagerForSlidingWindow(2, 128, None, config, scratch_token_num=7) - assert manager.req_to_sliding_window.shape == (2, 3, 32, 2, 64) - assert manager.req_to_sliding_window.data_ptr() == manager.sliding_kv_buffer.data_ptr() - assert manager.req_to_sliding_window.stride(0) == manager.sliding_kv_buffer.stride(0) - assert not hasattr(manager, "req_to_sliding_window_indexs") - assert manager.req_to_token_indexs.shape == (3, 128) - manager.req_to_sliding_window[:, 1].fill_(7) - torch.testing.assert_close( - manager.sliding_kv_buffer[:, 32:64], torch.full_like(manager.sliding_kv_buffer[:, 32:64], 7) - ) - manager.sliding_kv_buffer[:, manager.scratch_start :].fill_(11) - pages = SlidingWindowStateCacheManager(1, config) - manager.save_small_page_state(1, 0, pages) - manager.init_hybrid_attention_state(SimpleNamespace(req_idx=1)) - assert torch.count_nonzero(manager.req_to_sliding_window).item() == 0 - manager._restore_state(2, pages, 0) # Include the reserved hold request slot. - assert torch.all(manager.req_to_sliding_window[:, 2] == 7) - assert torch.all(manager.sliding_kv_buffer[:, manager.scratch_start :] == 11) - state = SimpleNamespace(input_ids=torch.zeros(7, dtype=torch.int32, device="cuda")) - manager.prepare_sliding_window(state) - assert torch.count_nonzero(manager.req_to_token_indexs).item() == 0 - torch.testing.assert_close( - state.sliding_window_mem_index, - torch.arange(manager.scratch_start, manager.scratch_start + 7, device="cuda"), - ) - - -@pytest.mark.parametrize("layer_index,is_shared", [(5, False), (11, False), (17, True)]) -def test_full_kv_write_maps_logical_layer_once_and_skips_shared_readers(layer_index, is_shared): - config = SlidingWindowCacheConfig({0: 0}, {5: 0, 11: 1, 17: 1}, 32, 1, 64, 1, 64, torch.bfloat16) - mem_manager = SimpleNamespace( - sliding_config=config, kv_buffer=torch.zeros((2, 8, 2, 64), dtype=torch.bfloat16, device="cuda") - ) - mem_manager.operator = HybridSlidingMemOperator(mem_manager) - layer = object.__new__(Gemma4TransformerLayerInfer) - layer.layer_num_, layer.is_sliding, layer.is_kv_shared_ = layer_index, False, is_shared - indexes = torch.tensor([1, 3], dtype=torch.int32, device="cuda") - kv = torch.randn((2, 2, 64), dtype=torch.bfloat16, device="cuda") - layer._post_cache_kv(kv, SimpleNamespace(mem_manager=mem_manager, mem_index=indexes), None) - expected = torch.zeros_like(mem_manager.kv_buffer) - if not is_shared: - expected[config.get_full_layer_index(layer_index), indexes] = kv - torch.testing.assert_close(mem_manager.kv_buffer, expected, atol=0, rtol=0) - - -@pytest.mark.parametrize("history_len", [0, 511, 512, 513, 1024]) -@pytest.mark.parametrize("q_len", [1, 31, 256, 768]) -def test_ring_attention_and_commit_match_token_cache(history_len, q_len): +@pytest.mark.parametrize("window,q_lengths", [(32, [1, 31, 65]), (512, [4096, 1, 513]), (1024, [8192, 7, 1023])]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_prefill_moves_all_layers_and_preserves_canonical_snapshot(window, q_lengths, dtype): torch.manual_seed(42) - window, head_dim, req_idx = 512, 64, 1 - seq_len = history_len + q_len - scratch_start = 3 * window - reference = torch.randn((seq_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) - runtime = torch.zeros((scratch_start + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) - old_positions = torch.arange(max(0, history_len - window), history_len, device="cuda") - runtime[req_idx * window + old_positions % window] = reference[old_positions] - runtime[scratch_start:] = reference[history_len:] - b_req = torch.tensor([req_idx], device="cuda", dtype=torch.int32) - b_seq = torch.tensor([seq_len], device="cuda", dtype=torch.int32) - b_q = torch.tensor([q_len], device="cuda", dtype=torch.int32) - b_start = torch.tensor([0], device="cuda", dtype=torch.int32) - b_history = torch.tensor([history_len], device="cuda", dtype=torch.int32) - image_end = torch.zeros(q_len, device="cuda", dtype=torch.int32) - q = torch.randn((q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) - actual, expected = torch.empty_like(q), torch.empty_like(q) - reference_mapping = torch.arange(seq_len, device="cuda", dtype=torch.int32).expand(3, -1) - for buffer, indexes, output, scratch in [ - (runtime, None, actual, scratch_start), - (reference, reference_mapping, expected, None), - ]: - context_attention_fwd_gemma4_mm( - q, - buffer[:, :1], - buffer[:, 1:], - output, - b_req, - b_start, - b_seq, - b_history, - q_len, - indexes, - image_end, - sliding_window=(window - 1, 0), - scratch_start=scratch, - ) - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - commit_sliding_window_state(runtime, b_req, b_seq, b_q, b_start, window, scratch_start, q_len) - positions = torch.arange(max(0, seq_len - window), seq_len, device="cuda") - torch.testing.assert_close(runtime[req_idx * window + positions % window], reference[positions], atol=0, rtol=0) - assert torch.count_nonzero(runtime[:window]).item() == 0 - - -@pytest.mark.parametrize("is_prefill", [True, False]) -def test_shared_kv_is_committed_only_after_last_reader_and_snapshot_is_independent(is_prefill): - window, history_len, q_len, head_dim = 512, 512, 256 if is_prefill else 1, 64 - layout, owners, last_reader = get_kv_cache_layout( - {"layer_types": ["sliding_attention", "full_attention", "sliding_attention"], "num_kv_shared_layers": 1} - ) - config = SlidingWindowCacheConfig( - layout["sliding_attention"], layout["full_attention"], window, 1, head_dim, 1, 64, torch.bfloat16 - ) - manager = object.__new__(ReqManagerForSlidingWindow) - manager.sliding_config = config - manager.sliding_window = window - manager.scratch_token_num = q_len - manager.scratch_start = 2 * window - manager.sliding_kv_buffer = torch.zeros((1, 2 * window + q_len, 2, head_dim), device="cuda", dtype=torch.bfloat16) - manager.req_to_sliding_window = manager.sliding_kv_buffer[:, : manager.scratch_start].view( - 1, 2, window, 2, head_dim - ) - manager.sliding_kv_buffer[0, manager.scratch_start :, 1] = 1 + req_ids, histories = [3, 0, 5], [0, window - 1, 2 * window + 3] + starts = [11, 18 + q_lengths[0], 31 + q_lengths[0] + q_lengths[1]] + lengths = [history + q_len for history, q_len in zip(histories, q_lengths)] + total_tokens = starts[-1] + q_lengths[-1] + 7 + runtime_start = 6 * window + 17 + pool = torch.full((3, runtime_start + total_tokens + 3 * window, 4, 32), -11, device="cuda", dtype=dtype) + references = [torch.randn((3, length, 4, 32), device="cuda", dtype=dtype) for length in lengths] int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - state = SimpleNamespace( - input_ids=int_tensor([0] * q_len), - b_req_idx=int_tensor([0]), - b_seq_len=int_tensor([history_len + q_len]), - b_q_seq_len=int_tensor([q_len]), - b_q_start_loc=int_tensor([0]), - b_ready_cache_len=int_tensor([history_len]), - max_q_seq_len=q_len, - b_image_token_end=int_tensor([0] * q_len), - req_manager=manager, + b_req, b_seq, b_history, b_q, b_start = map(int_tensor, [req_ids, lengths, histories, q_lengths, starts]) + for req, history, reference in zip(req_ids, histories, references): + positions = torch.arange(max(0, history - window), history, device="cuda") + pool[:, req * window + positions % window] = reference[:, positions] + + expected = pool.clone() + for batch, (history, start, reference) in enumerate(zip(histories, starts, references)): + positions = torch.arange(max(0, history - window), history, device="cuda") + current_start = runtime_start + start + (batch + 1) * window + expected[:, current_start + positions - history] = reference[:, positions] + move_sliding_window(pool, b_req, b_seq, b_history, b_start, window, runtime_start) + torch.testing.assert_close(pool, expected, atol=0, rtol=0) + + indexes = get_sliding_window_mem_indexes( + b_req, b_seq, b_q, b_start, window, runtime_start, total_tokens, max(q_lengths), is_prefill=True ) - manager.prepare_sliding_window(state) - q = torch.zeros((q_len, 1, head_dim), device="cuda", dtype=torch.bfloat16) - outputs = [] - for index in [0, 2]: - layer = object.__new__(Gemma4TransformerLayerInfer) - layer.layer_num_, layer.is_sliding, layer.sliding_window_ = index, True, window - layer.is_kv_shared_, layer.kv_share_target_layer_ = index != owners[index], owners[index] - layer.commit_sliding_state_ = last_reader[owners[index]] == index - layer.tp_q_head_num_, layer.head_dim_ = 1, head_dim - layer.alloc_tensor = lambda shape, dtype, device="cuda": torch.empty(shape, dtype=dtype, device=device) - if is_prefill: - outputs.append(layer._context_attention_kernel(q, None, state, None)) - else: - outputs.append(layer._token_attention_kernel(q, state, None)) - if index == 0: - assert torch.count_nonzero(manager.req_to_sliding_window[:, 0]).item() == 0 - torch.testing.assert_close(outputs[0], outputs[1], atol=0, rtol=0) - assert outputs[1][0, 0, 0].item() == 1 / window - assert manager.req_to_sliding_window[0, 0, 0, 1, 0].item() == 1 - - pages = SlidingWindowStateCacheManager(2, config) - page = pages.alloc_one_state_cache() - manager.save_small_page_state(0, page, pages) - saved = pages.get_state_cache(page).clone() - manager.req_to_sliding_window[:, 0].fill_(7) - torch.testing.assert_close(pages.get_state_cache(page), saved, atol=0, rtol=0) - manager._restore_state(1, pages, page) - torch.testing.assert_close(manager.req_to_sliding_window[:, 1], saved, atol=0, rtol=0) - pages.free_state_cache([page]) - assert pages.get_free_cache_num() == 2 - - -def test_empty_snapshot_does_not_read_gpu_request_ids(): - manager = object.__new__(ReqManagerForSlidingWindow) - # No runtime or page pool is needed for a no-op. In particular, no .tolist() - # or other GPU operation should be performed on b_req_idx. - manager.save_big_page_states(object(), [0, 1], [-1, -1]) - - -@pytest.mark.parametrize( - "window,q_lengths", - [(32, [6, 5, 32]), (512, [4096, 1, 513]), (512, [1, 8192, 511]), (1024, [8192, 4096, 1])], -) -def test_batched_window_commit_with_hold_request_and_cuda_graph_replay(window, q_lengths): - scratch_start, head_dim = 4 * window, 64 - req_ids = [2, 0, 3] - lengths = [q_lengths[0] + 2 * window + 3, q_lengths[1], q_lengths[2] + window - 1] - replay_lengths = [length + window + 7 for length in lengths] - starts = [0, q_lengths[0], q_lengths[0] + q_lengths[1]] + for batch, (start, q_len, history, reference) in enumerate(zip(starts, q_lengths, histories, references)): + expected_indexes = runtime_start + start + (batch + 1) * window + torch.arange(q_len, device="cuda") + torch.testing.assert_close(indexes[start : start + q_len].long(), expected_indexes, atol=0, rtol=0) + pool[:, indexes[start : start + q_len].long()] = reference[:, history:] + + expected = pool.clone() + for req, length, reference in zip(req_ids, lengths, references): + positions = torch.arange(max(0, length - window), length, device="cuda") + expected[:, req * window + positions % window] = reference[:, positions] + move_sliding_window(pool, b_req, b_seq, b_history, b_start, window, runtime_start, compact=True) + torch.testing.assert_close(pool, expected, atol=0, rtol=0) + # The page format remains a raw W-slot ring, independent of the active area. + snapshot = pool[:, 5 * window : 6 * window].clone() + pool[:, runtime_start:].zero_() + torch.testing.assert_close(snapshot, expected[:, 5 * window : 6 * window], atol=0, rtol=0) + + +@pytest.mark.parametrize("window", [32, 512, 1024]) +def test_decode_indexes_cuda_graph_replay(window): int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - b_req, b_seq, b_q, b_start = map(int_tensor, [req_ids, lengths, q_lengths, starts]) - runtime = torch.zeros((scratch_start + sum(q_lengths), 2, head_dim), device="cuda", dtype=torch.bfloat16) - runtime[scratch_start:] = torch.randn_like(runtime[scratch_start:]) + b_req, b_seq = int_tensor([2, 0, 5]), int_tensor([1, window, 2 * window + 1]) def forward(): - commit_sliding_window_state(runtime, b_req, b_seq, b_q, b_start, window, scratch_start, max(q_lengths)) + return get_sliding_window_mem_indexes(b_req, b_seq, None, None, window, 6 * window, 3, 1, False) - forward() + torch.testing.assert_close(forward(), int_tensor([2 * window, window - 1, 5 * window]), atol=0, rtol=0) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - forward() - # Replay with changed GPU state, including a padded/hold request ID. - runtime[:scratch_start].zero_() - runtime[scratch_start:].mul_(2) - b_seq.copy_(int_tensor(replay_lengths)) + indexes = forward() + b_req.copy_(int_tensor([5, 1, 3])) + b_seq.copy_(int_tensor([window + 3, 7, 4 * window])) graph.replay() - for req, seq, q_len, start in zip(req_ids, replay_lengths, q_lengths, starts): - tail_start = max(q_len - window, 0) - positions = torch.arange(seq - q_len + tail_start, seq, device="cuda") - expected_ring = torch.zeros_like(runtime[req * window : (req + 1) * window]) - expected_ring[positions % window] = runtime[scratch_start + start + tail_start : scratch_start + start + q_len] - torch.testing.assert_close(runtime[req * window : (req + 1) * window], expected_ring, atol=0, rtol=0) - assert torch.count_nonzero(runtime[window : 2 * window]).item() == 0 - - -@pytest.mark.parametrize("max_q_seq_len", [1, 511, 512, 513, 4096, 8192]) -def test_commit_grid_is_bounded_by_window(monkeypatch, max_q_seq_len): - import lightllm.common.basemodel.triton_kernel.sliding_window_state as state_kernel - - grids = [] - - class RecordingKernel: - def __getitem__(self, grid): - grids.append(grid) - return lambda *args, **kwargs: None - - monkeypatch.setattr(state_kernel, "_commit_sliding_window_state", RecordingKernel()) - layer_buffer = SimpleNamespace(shape=(8192, 2, 64), stride=lambda: (128, 64, 1)) - req_ids = SimpleNamespace(shape=(3,)) - state_kernel.commit_sliding_window_state(layer_buffer, req_ids, None, None, None, 512, 2048, max_q_seq_len) - assert grids == [(3, min(max_q_seq_len, 512), 2)] + torch.testing.assert_close(indexes, int_tensor([5 * window + 2, window + 6, 4 * window - 1]), atol=0, rtol=0) diff --git a/test/utils/test_sliding_cpu_cache_meta.py b/test/utils/test_sliding_cpu_cache_meta.py deleted file mode 100644 index 8f98fa5e3f..0000000000 --- a/test/utils/test_sliding_cpu_cache_meta.py +++ /dev/null @@ -1,173 +0,0 @@ -from types import SimpleNamespace - -import numpy as np -import pytest -import torch - -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig -from lightllm.models.gemma4.kv_layout import build_sliding_cache_config - - -@pytest.mark.parametrize( - "model_config,expected", - [ - ({"model_type": "gemma4"}, True), - ({"model_type": "gemma4_text"}, True), - ({"text_config": {"model_type": "gemma4_text"}}, True), - ({"model_type": "gemma4", "layer_types": ["full_attention"]}, True), - ({"model_type": "gemma3", "layer_types": ["sliding_attention", "full_attention"]}, False), - ({"model_type": "qwen3_5"}, False), - ], -) -def test_sliding_cache_architecture_is_selected_by_model_type(monkeypatch, model_config, expected): - import lightllm.utils.config_utils as config_utils - - monkeypatch.setattr(config_utils, "get_config_json", lambda _: model_config) - assert config_utils.is_sliding_att_mixed_model.__wrapped__("test-model") is expected - - -def _gemma_config(shared): - layer_num = 42 if shared else 60 - return { - "model_type": "gemma4_text", - "num_hidden_layers": layer_num, - "num_attention_heads": 8, - "num_key_value_heads": 2 if shared else 16, - "num_global_key_value_heads": None if shared else 4, - "head_dim": 256, - "global_head_dim": 512, - "sliding_window": 512 if shared else 1024, - "layer_types": (["sliding_attention"] * 5 + ["full_attention"]) * (layer_num // 6), - "num_kv_shared_layers": 18 if shared else 0, - } - - -@pytest.mark.parametrize("shared,tp_world_size", [(True, 1), (True, 2), (False, 1), (False, 2), (False, 4)]) -def test_cpu_page_layout_uses_physical_owners_and_all_tp_shards(shared, tp_world_size): - config = _gemma_config(shared) - layout = build_sliding_cache_config(config, tp_world_size, torch.bfloat16) - full_layers, sliding_layers = (4, 20) if shared else (10, 50) - full_heads = config["num_global_key_value_heads"] or config["num_key_value_heads"] - big_page_tokens = 2048 - full_bytes = full_layers * big_page_tokens * 2 * full_heads * config["global_head_dim"] * 2 - state_bytes = sliding_layers * config["sliding_window"] * 2 * config["num_key_value_heads"] * config["head_dim"] * 2 - assert layout.get_cpu_cache_full_att_bytes(big_page_tokens, tp_world_size) == full_bytes - assert layout.get_cpu_cache_state_bytes(tp_world_size) == state_bytes - assert layout.get_cpu_cache_big_page_bytes(big_page_tokens, tp_world_size) == full_bytes + state_bytes - - -def test_cpu_page_payload_is_aligned_without_changing_section_sizes(): - layout = SlidingWindowCacheConfig({0: 0}, {1: 0}, 1, 1, 3, 1, 5, torch.bfloat16) - assert layout.get_cpu_cache_full_att_bytes(3, 1) == 60 - assert layout.get_cpu_cache_state_bytes(1) == 12 - assert layout.get_cpu_cache_big_page_bytes(3, 1) == 80 - - -@pytest.mark.parametrize("shared,tp_world_size", [(True, 4), (False, 8)]) -def test_cpu_layout_rejects_replicated_kv_heads(shared, tp_world_size): - with pytest.raises(AssertionError, match="KV heads must be divisible"): - build_sliding_cache_config(_gemma_config(shared), tp_world_size, torch.bfloat16) - - -@pytest.mark.parametrize("wrapped", [False, True]) -def test_sliding_cpu_meta_is_a_flat_global_payload(monkeypatch, wrapped): - import lightllm.utils.kv_cache_utils as cache_utils - - config = _gemma_config(True) - layout = build_sliding_cache_config(config, 2, torch.bfloat16) - page_bytes = layout.get_cpu_cache_big_page_bytes(2048, 2) - args = SimpleNamespace( - model_dir="gemma-test", - enable_cpu_cache=True, - tp=4, - dp=2, - linear_att_hash_page_size=128, - linear_att_page_block_num=16, - cpu_cache_token_page_size=2048, - cpu_cache_storage_size=3 * page_bytes / 1024 ** 3, - mtp_mode=None, - ) - monkeypatch.setattr(cache_utils, "get_env_start_args", lambda: args) - monkeypatch.setattr(cache_utils, "is_linear_att_mixed_model", lambda _: False) - monkeypatch.setattr(cache_utils, "is_sliding_att_mixed_model", lambda _: True) - monkeypatch.setattr(cache_utils, "get_llm_data_type", lambda: torch.bfloat16) - monkeypatch.setattr(cache_utils, "get_config_json", lambda _: {"text_config": config} if wrapped else config) - meta = cache_utils.calcu_cpu_cache_meta.__wrapped__() - assert meta.data_type == torch.uint8 - assert (meta.layer_num, meta.token_page_size, meta.num_heads) == (1, 1, 1) - assert meta.head_dim == meta.calcu_one_page_size() == page_bytes - assert meta.page_num == 3 - assert args.cpu_cache_token_page_size == 2048 - - -def test_hybrid_request_initializes_cpu_hashes_without_linear_state(monkeypatch): - import lightllm.server.core.objs.req as req_module - - prompt = list(range(18)) - args = SimpleNamespace( - model_dir="gemma-test", - mtp_step=0, - enable_cpu_cache=True, - linear_att_hash_page_size=4, - linear_att_page_block_num=3, - cpu_cache_token_page_size=12, - ) - monkeypatch.setattr(req_module, "get_env_start_args", lambda: args) - monkeypatch.setattr(req_module, "is_hybrid_att_mixed_model", lambda _: True) - req = SimpleNamespace(index_in_shm_mem=0, ref_count=0) - req.create_logprobs_shm_array = lambda: None - req.create_prompt_ids_shm_array = lambda: setattr( - req, "shm_prompt_ids", SimpleNamespace(arr=np.empty(2048, dtype=np.int64)) - ) - req.post_init = lambda: None - req.get_prompt_ids = lambda: prompt - req._fill_linear_att_token_hash = lambda: req_module.Req._fill_linear_att_token_hash(req) - req._calcu_linear_att_cpu_cache_page_len_list = lambda: req_module.Req._calcu_linear_att_cpu_cache_page_len_list( - req - ) - req_module.Req.init(req, 0, prompt, req_module.SamplingParams(), tokenizer=None, chunked_prefill_size=16) - hashes = req.linear_att_token_hash_list.get_all() - assert len(hashes) == 4 - assert req.token_hash_list.get_all() == [hashes[2], hashes[3]] - assert req.token_hash_page_len_list.get_all() == [12, 16] - assert req.cpu_cache_match_page_indexes.get_all() == [] - - -@pytest.mark.parametrize("disable_tail,tail_buffer", [(False, None), (False, 3), (True, 3)]) -def test_sliding_offload_uses_existing_hybrid_tail_policy(monkeypatch, disable_tail, tail_buffer): - import lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache as cache_module - - module = object.__new__(cache_module.MultiLevelKvCacheModule) - module.args = SimpleNamespace(cpu_cache_token_page_size=16, disable_linear_att_small_page_cpu_cache=disable_tail) - monkeypatch.setattr(cache_module.g_infer_context, "is_linear_att_mixed_model", False) - monkeypatch.setattr(cache_module.g_infer_context, "is_hybrid_att_mixed_model", True) - req = SimpleNamespace(tail_linear_att_small_page_buffer_id=tail_buffer) - expected_pages = 2 if not disable_tail and tail_buffer is not None else 1 - assert module._handle_linear_att_last_page(req, 2, [16, 20]) == expected_pages - - -@pytest.mark.parametrize("enable_cpu_cache,disable_gpu_cache", [(True, False), (True, True), (False, True)]) -def test_cpu_state_transfer_requires_gpu_prefix_cache(monkeypatch, enable_cpu_cache, disable_gpu_cache): - import lightllm.models.gemma4.model as gemma_model - - model = object.__new__(gemma_model.Gemma4TpPartModel) - model.load_way, model.tp_world_size_ = "HF", 2 - model.config = _gemma_config(True) - args = SimpleNamespace( - mtp_step=0, - enable_cpu_cache=enable_cpu_cache, - disable_dynamic_prompt_cache=disable_gpu_cache, - disable_chunked_prefill=False, - run_mode="normal", - llm_kv_type="None", - enable_dp_prompt_cache_fetch=False, - diverse_mode=False, - enable_prefill_microbatch_overlap=False, - enable_decode_microbatch_overlap=False, - ) - monkeypatch.setattr(gemma_model, "get_env_start_args", lambda: args) - if enable_cpu_cache and disable_gpu_cache: - with pytest.raises(AssertionError, match="CPU cache requires GPU prefix cache"): - model._verify_params() - else: - model._verify_params() diff --git a/test/utils/test_sliding_window_cache.py b/test/utils/test_sliding_window_cache.py deleted file mode 100644 index fc066bb3a4..0000000000 --- a/test/utils/test_sliding_window_cache.py +++ /dev/null @@ -1,249 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager -from lightllm.common.req_manager.linear_att import ReqManagerForMamba -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig -from lightllm.models.gemma4.kv_layout import get_kv_cache_layout - - -@pytest.mark.parametrize("layer_num,shared,sliding_num,full_num", [(42, 18, 20, 4), (60, 0, 50, 10)]) -def test_gemma_physical_owners_and_last_readers(layer_num, shared, sliding_num, full_num): - layer_types = ["sliding_attention"] * 5 + ["full_attention"] - maps, owners, last_readers = get_kv_cache_layout( - {"layer_types": layer_types * (layer_num // 6), "num_kv_shared_layers": shared} - ) - assert len(set(maps["sliding_attention"].values())) == sliding_num - assert len(set(maps["full_attention"].values())) == full_num - for index, owner in enumerate(owners): - assert owner <= index <= last_readers[owner] - if shared: - assert owners[40] == 22 and last_readers[22] == 40 - assert owners[41] == 23 and last_readers[23] == 41 - - -def _memory_manager(big_page_tokens=2048, small_pages=8, enabled=True, cpu_cache=False): - manager = object.__new__(HybridSlidingMemoryManager) - manager.size = None - manager.head_num, manager.head_dim, manager.layer_num, manager.dtype = 1, 512, 10, torch.bfloat16 - manager.sliding_config = SlidingWindowCacheConfig( - {i: i for i in range(50)}, {50 + i: i for i in range(10)}, 1024, 4, 256, 1, 512, torch.bfloat16 - ) - manager.big_page_token_num, manager.small_page_num, manager.enable_prompt_cache = ( - big_page_tokens, - small_pages, - enabled, - ) - manager.cpu_cache_temp_page_num = 2 if cpu_cache else 0 - return manager - - -def _required_bytes(manager, token_num): - big_pages = (token_num + manager.big_page_token_num - 1) // manager.big_page_token_num - state_pages = manager.small_page_num + manager.cpu_cache_temp_page_num - if manager.enable_prompt_cache: - state_pages += big_pages - return (token_num + 1) * manager.get_cell_size() + state_pages * manager.sliding_config.get_state_nbytes() - - -def _profile_with_budget(monkeypatch, manager, available_bytes, mem_fraction=1.0, total_bytes=16 * 1024 ** 3): - import lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager as memory_module - - monkeypatch.setattr(memory_module.dist, "get_world_size", lambda: 1) - monkeypatch.setattr(memory_module, "get_available_gpu_memory", lambda world_size: available_bytes / 1024 ** 3) - monkeypatch.setattr(memory_module, "get_total_gpu_memory", lambda: total_bytes / 1024 ** 3) - monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) - manager.profile_size(mem_fraction) - - -@pytest.mark.parametrize("token_num", [1, 2047, 2048, 2049, 8192]) -def test_profile_accounts_for_small_big_partial_page_and_hold_token(monkeypatch, token_num): - manager = _memory_manager() - assert manager.sliding_config.get_state_nbytes() == 200 * 1024 ** 2 - expected = (token_num + 1) * 20480 + (8 + (token_num + 2047) // 2048) * 200 * 1024 ** 2 - assert _required_bytes(manager, token_num) == expected - _profile_with_budget(monkeypatch, manager, expected) - assert manager.size == token_num - manager.size = None - if token_num > 1: - _profile_with_budget(monkeypatch, manager, expected - 1) - assert manager.size == token_num - 1 - else: - with pytest.raises(ValueError, match="Insufficient GPU memory"): - _profile_with_budget(monkeypatch, manager, expected - 1) - - -@pytest.mark.parametrize("leftover", [1, 200 * 1024 ** 2 - 1, 200 * 1024 ** 2]) -def test_profile_cannot_start_next_page_without_checkpoint_and_token_budget(monkeypatch, leftover): - manager = _memory_manager() - _profile_with_budget(monkeypatch, manager, _required_bytes(manager, 2048) + leftover) - assert manager.size == 2048 - - -@pytest.mark.parametrize("available_bytes", [-1, 0, 80 * 1024 ** 3]) -def test_profile_reports_impossible_checkpoint_budget(monkeypatch, available_bytes): - manager = _memory_manager(small_pages=512) - with pytest.raises(ValueError, match="linear_att_cache_size"): - _profile_with_budget(monkeypatch, manager, available_bytes) - - -def test_disabled_prompt_cache_does_not_reserve_pages(monkeypatch): - manager = _memory_manager(small_pages=0, enabled=False) - _profile_with_budget(monkeypatch, manager, 4097 * manager.get_cell_size()) - assert manager.size == 4096 - - -def test_cpu_cache_reserves_two_additional_window_checkpoints(monkeypatch): - gpu_only = _memory_manager() - cpu_cache = _memory_manager(cpu_cache=True) - gpu_budget = _required_bytes(gpu_only, 4096) - cpu_budget = gpu_budget + 2 * cpu_cache.sliding_config.get_state_nbytes() - _profile_with_budget(monkeypatch, gpu_only, gpu_budget) - _profile_with_budget(monkeypatch, cpu_cache, cpu_budget) - assert gpu_only.size == cpu_cache.size == 4096 - - -def test_explicit_size_is_checked_against_complete_cache_budget(monkeypatch): - manager = _memory_manager(cpu_cache=True) - manager.size = 2049 - budget = _required_bytes(manager, manager.size) - _profile_with_budget(monkeypatch, manager, budget) - assert manager.size == 2049 - with pytest.raises(ValueError, match="exceed available GPU memory"): - _profile_with_budget(monkeypatch, manager, budget - 1) - - -@pytest.mark.parametrize("explicit_size", [None, 2049]) -def test_mem_fraction_reserves_headroom_only_for_automatic_size(monkeypatch, explicit_size): - manager = _memory_manager() - manager.size = explicit_size - total_bytes, mem_fraction = 8 * 1024 ** 3, 0.5 - budget = _required_bytes(manager, 2048) + total_bytes // 2 - _profile_with_budget(monkeypatch, manager, budget, mem_fraction=mem_fraction, total_bytes=total_bytes) - assert manager.size == (2048 if explicit_size is None else explicit_size) - - -@pytest.mark.parametrize("disabled", [False, True]) -def test_page_pools_follow_active_prompt_cache_flag(monkeypatch, disabled): - import lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager as memory_module - - args = SimpleNamespace( - use_dynamic_prompt_cache=False, - disable_dynamic_prompt_cache=disabled, - enable_cpu_cache=False, - linear_att_cache_size=3, - linear_att_hash_page_size=32, - linear_att_page_block_num=8, - ) - monkeypatch.setattr(memory_module, "get_env_start_args", lambda: args) - monkeypatch.setattr(memory_module.MemoryManager, "__init__", lambda self, **kwargs: None) - config = SlidingWindowCacheConfig({0: 0}, {1: 0}, 32, 1, 64, 1, 64, torch.bfloat16) - manager = HybridSlidingMemoryManager(size=256, sliding_config=config) - assert manager.enable_prompt_cache is not disabled - assert manager.small_page_num == (0 if disabled else 3) - assert manager.big_page_token_num == 256 - assert manager.cpu_cache_temp_page_num == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("enabled,cpu_cache", [(False, False), (True, False), (True, True)]) -def test_profiled_gpu_pools_match_reserved_bytes_and_are_reused(monkeypatch, enabled, cpu_cache): - from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow - - manager = _memory_manager(big_page_tokens=32, small_pages=2 if enabled else 0, enabled=enabled, cpu_cache=cpu_cache) - manager.head_num, manager.head_dim, manager.layer_num = 1, 64, 1 - manager.sliding_config = SlidingWindowCacheConfig({0: 0}, {1: 0}, 32, 1, 64, 1, 64, torch.bfloat16) - budget = _required_bytes(manager, 65) - _profile_with_budget(monkeypatch, manager, budget) - assert manager.size == 65 - manager._init_buffers(manager.size, manager.dtype, manager.head_num, manager.head_dim, manager.layer_num) - assert manager.linear_att_big_page_buffers.size == (3 if enabled else 0) + (2 if cpu_cache else 0) - assert manager.linear_att_big_page_buffers.get_free_cache_num() == (3 if enabled else 0) - assert manager.sliding_small_page_buffers.size == (2 if enabled else 0) - allocated = sum( - t.numel() * t.element_size() - for t in [ - manager.kv_buffer, - manager.linear_att_big_page_buffers.state_cache, - manager.sliding_small_page_buffers.state_cache, - ] - ) - assert allocated == budget - req_manager = object.__new__(ReqManagerForSlidingWindow) - req_manager.mem_manager = manager - assert req_manager.create_state_cache_manager(2) is manager.sliding_small_page_buffers - - -@pytest.mark.parametrize("mtp_step", [0, 2]) -def test_linear_small_page_preserves_main_copy_and_mtp_crop(mtp_step): - manager = object.__new__(ReqManagerForMamba) - manager.mtp_step = mtp_step - manager.linear_config = SimpleNamespace(get_conv_state_shape=lambda: (3, 4)) - conv = torch.arange(2 * 3 * 3 * (4 + mtp_step)).reshape(2, 3, 3, 4 + mtp_step) - ssm = torch.arange(2 * 3 * (mtp_step + 1) * 5).reshape(2, 3 * (mtp_step + 1), 5) - manager.req_to_conv_state, manager.req_to_ssm_state = SimpleNamespace(buffer=conv), SimpleNamespace(buffer=ssm) - dst_conv, dst_ssm = torch.empty((2, 3, 4), dtype=conv.dtype), torch.empty((2, 5), dtype=ssm.dtype) - pages = SimpleNamespace(get_state_cache=lambda buffer_idx: (dst_conv, dst_ssm)) - manager.save_small_page_state(1, 0, pages) - torch.testing.assert_close(dst_conv, conv[:, 1, :, :4]) - torch.testing.assert_close(dst_ssm, ssm[:, mtp_step + 1]) - - -@pytest.mark.parametrize("unsupported_mode", ["enable_dp_prompt_cache_fetch", "diverse_mode"]) -def test_unsupported_sliding_state_transfer_modes_fail_before_loading_weights(monkeypatch, unsupported_mode): - import lightllm.models.gemma4.model as gemma_model - - model = object.__new__(gemma_model.Gemma4TpPartModel) - model.load_way, model.tp_world_size_ = "HF", 2 - model.config = {"num_attention_heads": 8, "num_key_value_heads": 2, "num_hidden_layers": 42} - args = SimpleNamespace( - mtp_step=0, - enable_cpu_cache=False, - disable_chunked_prefill=False, - run_mode="normal", - llm_kv_type="None", - enable_dp_prompt_cache_fetch=False, - diverse_mode=False, - ) - setattr(args, unsupported_mode, True) - monkeypatch.setattr(gemma_model, "get_env_start_args", lambda: args) - with pytest.raises(AssertionError, match="does not support"): - model._verify_params() - - -@pytest.mark.parametrize("shared_layers", [0, 18]) -@pytest.mark.parametrize( - "overlap_mode", [None, "enable_prefill_microbatch_overlap", "enable_decode_microbatch_overlap"] -) -def test_shared_kv_rejects_interleaved_microbatches(monkeypatch, shared_layers, overlap_mode): - import lightllm.models.gemma4.model as gemma_model - - model = object.__new__(gemma_model.Gemma4TpPartModel) - model.load_way, model.tp_world_size_ = "HF", 2 - model.config = { - "num_attention_heads": 8, - "num_key_value_heads": 2, - "num_hidden_layers": 42, - "num_kv_shared_layers": shared_layers, - } - args = SimpleNamespace( - mtp_step=0, - enable_cpu_cache=False, - disable_chunked_prefill=False, - run_mode="normal", - llm_kv_type="None", - enable_dp_prompt_cache_fetch=False, - diverse_mode=False, - enable_prefill_microbatch_overlap=False, - enable_decode_microbatch_overlap=False, - ) - if overlap_mode is not None: - setattr(args, overlap_mode, True) - monkeypatch.setattr(gemma_model, "get_env_start_args", lambda: args) - if shared_layers and overlap_mode is not None: - with pytest.raises(AssertionError, match="shared sliding-window KV does not support microbatch overlap"): - model._verify_params() - else: - model._verify_params() diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index f9aebf83bc..efd37a986b 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -159,15 +159,3 @@ def test_non_gpu_hybrid_cache_tiers_release_pending_state_pages(is_linear): assert freed_big_pages == [8, 9] assert req.tail_linear_att_small_page_buffer_id is None assert req.linear_att_len_to_big_page_id == {} - - -def test_hybrid_snapshot_outside_page_boundaries_does_not_touch_runtime(): - context = InferenceContext() - context.is_hybrid_att_mixed_model = True - context.args = SimpleNamespace( - linear_att_hash_page_size=32, linear_att_page_block_num=8, disable_chunked_prefill=False - ) - context.radix_cache = SimpleNamespace() - context.req_manager = None # Any attempted snapshot would fail. - reqs = [SimpleNamespace(req_idx=0, get_chuncked_input_token_len=lambda: 17, linear_att_cache_len=32)] - context.copy_linear_att_state_to_cache_buffer(b_req_idx=[0], reqs=reqs) diff --git a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_state_cache.py deleted file mode 100644 index 33ead8cb8d..0000000000 --- a/unit_tests/server/router/model_infer/test_hybrid_state_cache.py +++ /dev/null @@ -1,163 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.req_manager.linear_att import ReqManagerForMamba -from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow -from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager -from lightllm.server.router.model_infer.infer_batch import InferenceContext - - -def _cpu_linear_req_manager(mtp_step): - manager = object.__new__(ReqManagerForMamba) - manager.mtp_step = mtp_step - manager.req_to_conv_state = SimpleNamespace( - buffer=torch.arange(2 * 3 * 3 * (4 + mtp_step), dtype=torch.float32).reshape(2, 3, 3, 4 + mtp_step) - ) - manager.req_to_ssm_state = SimpleNamespace( - buffer=torch.arange(2 * 3 * (mtp_step + 1) * 5, dtype=torch.float32).reshape(2, 3 * (mtp_step + 1), 5) - ) - manager.req_to_mtp_state_index = torch.full((3,), mtp_step, dtype=torch.int32) if mtp_step else None - return manager - - -@pytest.mark.parametrize("mtp_step", [0, 2]) -def test_linear_init_clears_entire_request_state_and_resets_mtp_index(mtp_step): - manager = _cpu_linear_req_manager(mtp_step) - req = SimpleNamespace(req_idx=1) - conv = manager.req_to_conv_state.buffer - ssm = manager.req_to_ssm_state.buffer - ssm_start = req.req_idx * (mtp_step + 1) - conv[:, req.req_idx] = float("nan") - ssm[:, ssm_start : ssm_start + mtp_step + 1] = float("nan") - expected_conv, expected_ssm = conv.clone(), ssm.clone() - expected_conv[:, req.req_idx].zero_() - expected_ssm[:, ssm_start : ssm_start + mtp_step + 1].zero_() - - manager.init_hybrid_attention_state(req=req) - - torch.testing.assert_close(conv, expected_conv, atol=0, rtol=0) - torch.testing.assert_close(ssm, expected_ssm, atol=0, rtol=0) - if mtp_step: - torch.testing.assert_close( - manager.req_to_mtp_state_index, torch.tensor([mtp_step, 0, mtp_step], dtype=torch.int32) - ) - else: - assert manager.req_to_mtp_state_index is None - - -@pytest.mark.parametrize("mtp_step", [0, 2]) -@pytest.mark.parametrize("page_kind", ["big", "small"]) -def test_linear_restore_preserves_mtp_conv_tail_and_noncanonical_ssm_rows(mtp_step, page_kind): - manager = _cpu_linear_req_manager(mtp_step) - req = SimpleNamespace(req_idx=1, shared_kv_node=SimpleNamespace(small_page_buffer_idx=2)) - conv_pages = torch.arange(3 * 2 * 3 * 4, dtype=torch.float32).reshape(3, 2, 3, 4) + 1000 - ssm_pages = torch.arange(3 * 2 * 5, dtype=torch.float32).reshape(3, 2, 5) + 2000 - pages = SimpleNamespace(get_state_cache=lambda buffer_idx: (conv_pages[buffer_idx], ssm_pages[buffer_idx])) - manager.mem_manager = SimpleNamespace(linear_att_big_page_buffers=pages) - expected_conv = manager.req_to_conv_state.buffer.clone() - expected_ssm = manager.req_to_ssm_state.buffer.clone() - expected_conv[:, req.req_idx, ..., :4] = conv_pages[2] - expected_ssm[:, req.req_idx * (mtp_step + 1)] = ssm_pages[2] - - if page_kind == "big": - manager.restore_big_page_state(big_page_buffer_idx=2, req=req) - else: - manager.restore_small_page_state(req=req, small_page_buffers=pages) - - torch.testing.assert_close(manager.req_to_conv_state.buffer, expected_conv, atol=0, rtol=0) - torch.testing.assert_close(manager.req_to_ssm_state.buffer, expected_ssm, atol=0, rtol=0) - if mtp_step: - torch.testing.assert_close( - manager.req_to_mtp_state_index, torch.tensor([mtp_step, 0, mtp_step], dtype=torch.int32) - ) - else: - assert manager.req_to_mtp_state_index is None - - -@pytest.mark.parametrize("is_hybrid,radix_cache", [(False, object()), (True, None)]) -def test_snapshot_without_hybrid_cache_returns_before_reading_requests(is_hybrid, radix_cache): - context = InferenceContext(is_hybrid_att_mixed_model=is_hybrid, radix_cache=radix_cache) - - # Neither argument supports iteration or len: an early return must not inspect them. - context.copy_linear_att_state_to_cache_buffer(b_req_idx=object(), reqs=object()) - - -@pytest.mark.parametrize("chunk_end,cache_len", [(17, 32), (768, 544)]) -def test_snapshot_outside_cacheable_boundaries_does_not_allocate_or_copy(chunk_end, cache_len): - context = InferenceContext(is_hybrid_att_mixed_model=True, radix_cache=object()) - context.args = SimpleNamespace( - linear_att_hash_page_size=32, linear_att_page_block_num=8, disable_chunked_prefill=False - ) - req = SimpleNamespace( - req_idx=0, - get_chuncked_input_token_len=lambda: chunk_end, - linear_att_cache_len=cache_len, - linear_att_len_to_big_page_id={}, - tail_linear_att_small_page_buffer_id=None, - ) - - # The radix object has no allocator and req_manager is None, so either access fails. - context.copy_linear_att_state_to_cache_buffer(b_req_idx=[0], reqs=[req]) - - assert req.linear_att_len_to_big_page_id == {} - assert req.tail_linear_att_small_page_buffer_id is None - - -def _cpu_state_cache(size, keep_num=0): - pages = object.__new__(SlidingWindowStateCacheManager) - pages.size = size - pages.keep_num = keep_num - pages.state_cache = torch.empty((size, 2, 4, 2, 4), dtype=torch.float32) - pages.clear_to_init_state() - return pages - - -def test_sliding_big_snapshot_skips_invalid_requests_and_copies_only_selected_page(): - pages = _cpu_state_cache(3) - manager = object.__new__(ReqManagerForSlidingWindow) - manager.sliding_window = 4 - manager.mem_manager = SimpleNamespace(linear_att_big_page_buffers=pages) - manager.req_to_sliding_window = torch.arange(2 * 12 * 2 * 4, dtype=torch.float32).reshape(2, 3, 4, 2, 4) - expected = manager.req_to_sliding_window[:, 1].clone() - - # Skipped request IDs are deliberately out of range; GPU request IDs must not be read. - manager.save_big_page_states(b_req_idx=object(), req_indexes=[999, 1, 888], buffer_indexes=[-1, 2, -1]) - - torch.testing.assert_close(pages.get_state_cache(2), expected, atol=0, rtol=0) - assert torch.count_nonzero(pages.state_cache[:2]).item() == 0 - manager.req_to_sliding_window.fill_(-1) - torch.testing.assert_close(pages.get_state_cache(2), expected, atol=0, rtol=0) - - -def test_sliding_state_pool_exhaustion_and_released_slot_reuse(): - pages = _cpu_state_cache(2) - assert pages.alloc_state_cache(3) is None - assert pages.get_free_cache_num() == 2 - assert pages.alloc_state_cache(2) == [0, 1] - assert pages.get_used_cache_num() == 2 - assert pages.alloc_one_state_cache() is None - - pages.free_state_cache([1]) - assert pages.get_free_cache_num() == 1 - assert pages.alloc_one_state_cache() == 1 - assert pages.alloc_one_state_cache() is None - - pages.free_state_cache([0, 1]) - assert pages.get_free_cache_num() == 2 - assert pages.get_used_cache_num() == 0 - - -def test_sliding_state_pool_preserves_cpu_transfer_slots(): - pages = _cpu_state_cache(5, keep_num=2) - assert pages.alloc_state_cache(3) == [0, 1, 2] - assert pages.alloc_one_state_cache() is None - for reserved_id in [3, 4]: - with pytest.raises(AssertionError): - pages.free_state_cache([reserved_id]) - pages.free_state_cache([0, 1, 2]) - assert pages.get_free_cache_num() == 3 - pages.clear_to_init_state() - assert pages.alloc_state_cache(3) == [0, 1, 2] - assert pages.alloc_one_state_cache() is None diff --git a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py deleted file mode 100644 index f8194a4705..0000000000 --- a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_loading.py +++ /dev/null @@ -1,125 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel import sliding_window_cpu_cache_copy as copy_kernels -from lightllm.common.kv_cache_mem_manager.operator import hybrid_sliding as operator_module -from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow -from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as cache_module - - -@pytest.mark.parametrize( - "gpu_prefix,cpu_prefix,expected_endpoints", - [(288, 736, {512: 0}), (288, 768, {512: 1, 768: 0}), (544, 736, {})], -) -def test_cpu_load_prepends_partial_gpu_page_and_restores_absolute_checkpoint( - monkeypatch, gpu_prefix, cpu_prefix, expected_endpoints -): - # Run the real public loader, sliding operator and runtime restore with - # CPU tensors. Only the transfer kernel and CUDA/distributed calls are - # replaced; GPU kernel and stream behavior have separate coverage. - monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, non_blocking=False: self) - monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None)) - monkeypatch.setattr(cache_module.dist, "barrier", lambda group: None) - args = SimpleNamespace(cpu_cache_token_page_size=256, linear_att_hash_page_size=32, linear_att_page_block_num=8) - monkeypatch.setattr(operator_module, "get_env_start_args", lambda: args) - monkeypatch.setattr(operator_module, "get_current_rank_in_dp", lambda: 0) - monkeypatch.setattr(operator_module, "get_dp_world_size", lambda: 1) - - allocated_tokens, evicted_tokens, dereferenced_pages, transfers = [], [], [], [] - states = torch.zeros((6, 1, 4, 2, 4)) - free_ids = iter(range(4)) - state_pool = SimpleNamespace( - state_cache=states, - alloc_one_state_cache=lambda: next(free_ids), - get_state_cache=lambda index: states[index], - ) - - def alloc(need_size): - allocated_tokens.append(need_size) - return torch.arange(1000, 1000 + need_size, dtype=torch.int32) - - mem_manager = SimpleNamespace( - alloc=alloc, - sliding_config=object(), - kv_buffer=object(), - linear_att_big_page_buffers=state_pool, - CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID=4, - ) - mem_manager.operator = operator_module.HybridSlidingMemOperator(mem_manager) - req_manager = object.__new__(ReqManagerForSlidingWindow) - req_manager.mem_manager, req_manager.sliding_window = mem_manager, 4 - req_manager.req_to_sliding_window = torch.full((1, 2, 4, 2, 4), -1.0) - req_manager.req_to_token_indexs = torch.full((2, 1024), -1, dtype=torch.int32) - req_manager.req_to_token_indexs[1, :gpu_prefix] = torch.arange(10000, 10000 + gpu_prefix, dtype=torch.int32) - original_mapping = req_manager.req_to_token_indexs.clone() - radix_cache = SimpleNamespace( - free_radix_cache_to_get_enough_token=lambda need_token_num: evicted_tokens.append(need_token_num) - ) - monkeypatch.setattr(cache_module.g_infer_context, "req_manager", req_manager) - monkeypatch.setattr(cache_module.g_infer_context, "radix_cache", radix_cache) - monkeypatch.setattr(cache_module.g_infer_context, "get_can_alloc_token_num", lambda: 2048) - - req = SimpleNamespace( - req_idx=1, - cur_kv_len=gpu_prefix, - linear_att_len_to_big_page_id={}, - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - input_len=cpu_prefix + 1, - disk_prompt_cache_len=0, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: [4, 8, 12]), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: [256, 512, cpu_prefix]), - ), - ) - - def load(**kwargs): - # cur_kv_len must already be the absolute CPU endpoint when the - # operator assigns full-page checkpoints, not the old GPU hit length. - assert req.cur_kv_len == cpu_prefix - transfers.append(kwargs) - for state_id, cpu_page in zip(kwargs["big_page_buffer_ids"], kwargs["page_indexes"]): - kwargs["gpu_sliding_state"][state_id].fill_(cpu_page.item()) - - monkeypatch.setattr(copy_kernels, "copy_cpu_cache_to_kv_buffer", load) - module = object.__new__(cache_module.MultiLevelKvCacheModule) - module.backend = SimpleNamespace( - is_master_in_dp=True, - radix_cache=radix_cache, - model=SimpleNamespace(mem_manager=mem_manager, req_manager=req_manager), - ) - module.need_sync_compute_stream = lambda: False - module.init_sync_group = object() - module.cpu_cache_client = SimpleNamespace( - cpu_kv_cache_tensor=object(), - lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), - deref_pages=lambda page_list: dereferenced_pages.extend(page_list), - ) - - module.load_cpu_cache_to_reqs([req]) - - need_tokens = cpu_prefix - gpu_prefix - page_start = gpu_prefix // 256 * 256 - new_indexes = torch.arange(1000, 1000 + need_tokens, dtype=torch.int32) - expected_transfer = torch.cat([original_mapping[1, page_start:gpu_prefix], new_indexes]) - padding = (-len(expected_transfer)) % 256 - assert allocated_tokens == evicted_tokens == [need_tokens] - assert len(transfers) == 1 - assert transfers[0]["mem_indexes"].tolist() == expected_transfer.tolist() + [-1] * padding - assert transfers[0]["page_indexes"].tolist() == [4, 8, 12][gpu_prefix // 256 :] - assert req.linear_att_len_to_big_page_id == expected_endpoints - assert 4 not in req.linear_att_len_to_big_page_id.values() - if cpu_prefix % 256: - assert transfers[0]["big_page_buffer_ids"][-1].item() == 4 - torch.testing.assert_close(req_manager.req_to_token_indexs[1, :gpu_prefix], original_mapping[1, :gpu_prefix]) - torch.testing.assert_close(req_manager.req_to_token_indexs[1, gpu_prefix:cpu_prefix], new_indexes) - assert torch.all(req_manager.req_to_token_indexs[1, cpu_prefix:] == -1) - assert torch.all(req_manager.req_to_sliding_window[:, 0] == -1) - assert torch.all(req_manager.req_to_sliding_window[:, 1] == 12) - assert req.shm_req.cpu_prompt_cache_len == need_tokens - assert req.shm_req.shm_cur_kv_len == cpu_prefix - # Dereference all matched pages, including the page already covered by - # the GPU prefix and omitted from the actual transfer. - assert dereferenced_pages == [4, 8, 12] diff --git a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py b/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py deleted file mode 100644 index 85e56a074a..0000000000 --- a/unit_tests/server/router/model_infer/test_sliding_cpu_cache_operator.py +++ /dev/null @@ -1,259 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel import sliding_window_cpu_cache_copy as copy_kernels -from lightllm.common.kv_cache_mem_manager.operator import hybrid_sliding as operator_module -from lightllm.common.req_manager.sliding_window import ReqManagerForSlidingWindow -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig, SlidingWindowStateCacheManager -from lightllm.server.router.model_infer.infer_batch import g_infer_context - - -def _cpu_pages(size, keep_num=0): - pages = object.__new__(SlidingWindowStateCacheManager) - pages.size, pages.keep_num = size, keep_num - pages.state_cache = torch.empty((size, 1, 4, 2, 4), dtype=torch.float32) - pages.clear_to_init_state() - return pages - - -@pytest.fixture -def sliding_operator(monkeypatch): - # Exercise page ownership and transfer orchestration on CPU; kernel tests - # separately cover real CUDA pointers, byte layout, and stream ordering. - monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, non_blocking=False: self) - monkeypatch.setattr( - operator_module, - "get_env_start_args", - lambda: SimpleNamespace(cpu_cache_token_page_size=8, linear_att_hash_page_size=2, linear_att_page_block_num=4), - ) - monkeypatch.setattr(operator_module, "get_current_rank_in_dp", lambda: 0) - monkeypatch.setattr(operator_module, "get_dp_world_size", lambda: 1) - manager = SimpleNamespace( - sliding_config=object(), - kv_buffer=torch.zeros((1, 33, 2, 4)), - linear_att_big_page_buffers=_cpu_pages(6, keep_num=2), - CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID=4, - CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID=5, - ) - req_manager = object.__new__(ReqManagerForSlidingWindow) - req_manager.mem_manager = manager - req_manager.sliding_window = 4 - req_manager.req_to_sliding_window = torch.full((1, 2, 4, 2, 4), -1.0) - small_pages = _cpu_pages(2) - small_pages.get_state_cache(0).fill_(17) - radix = SimpleNamespace( - linear_att_small_page_buffers=small_pages, - get_big_page_ids_by_node=lambda node: [] if node is None else node.big_page_ids.copy(), - ) - monkeypatch.setattr(g_infer_context, "req_manager", req_manager) - monkeypatch.setattr(g_infer_context, "radix_cache", radix) - return operator_module.HybridSlidingMemOperator(manager), req_manager, small_pages - - -@pytest.mark.parametrize( - "cached_tokens,token_num,expected_endpoints,has_tail", - [(0, 6, {}, True), (0, 16, {8: 1, 16: 0}, False), (0, 22, {8: 1, 16: 0}, True), (8, 22, {16: 1, 24: 0}, True)], -) -def test_load_restores_last_checkpoint_without_owning_tail_staging_slot( - monkeypatch, sliding_operator, cached_tokens, token_num, expected_endpoints, has_tail -): - operator, req_manager, _ = sliding_operator - page_num = (token_num + 7) // 8 - captures = [] - - def load(**kwargs): - captures.append(kwargs) - for buffer_id, page_id in zip(kwargs["big_page_buffer_ids"].tolist(), kwargs["page_indexes"].tolist()): - kwargs["gpu_sliding_state"][buffer_id].fill_(10 + page_id) - - monkeypatch.setattr(copy_kernels, "copy_cpu_cache_to_kv_buffer", load) - req = SimpleNamespace(req_idx=1, cur_kv_len=cached_tokens + token_num, linear_att_len_to_big_page_id={}) - operator.load_cpu_cache_to_gpu( - torch.arange(cached_tokens, cached_tokens + token_num, dtype=torch.int32), - torch.arange(cached_tokens // 8, cached_tokens // 8 + page_num, dtype=torch.int32), - SimpleNamespace(cpu_kv_cache_tensor=object()), - req, - ) - - assert req.linear_att_len_to_big_page_id == expected_endpoints - assert len(captures) == 1 - padded_indexes = captures[0]["mem_indexes"] - assert padded_indexes.tolist() == list(range(cached_tokens, cached_tokens + token_num)) + [-1] * ( - page_num * 8 - token_num - ) - assert 4 not in req.linear_att_len_to_big_page_id.values() - assert operator.mem_manager.linear_att_big_page_buffers.get_free_cache_num() == 4 - token_num // 8 - if has_tail: - assert captures[0]["big_page_buffer_ids"][-1].item() == 4 - torch.testing.assert_close( - req_manager.req_to_sliding_window[:, 1], - torch.full((1, 4, 2, 4), 10.0 + cached_tokens // 8 + page_num - 1), - atol=0, - rtol=0, - ) - assert torch.all(req_manager.req_to_sliding_window[:, 0] == -1) - - -@pytest.mark.parametrize("token_num", [16, 22]) -def test_offload_combines_shared_owned_and_tail_checkpoints(monkeypatch, sliding_operator, token_num): - operator, _, small_pages = sliding_operator - captures = [] - monkeypatch.setattr(copy_kernels, "copy_kv_buffer_to_cpu_cache", lambda **kwargs: captures.append(kwargs)) - big_pages = operator.mem_manager.linear_att_big_page_buffers - big_pages.get_state_cache(1).fill_(11) - big_pages.get_state_cache(3).fill_(13) - req = SimpleNamespace( - shared_kv_node=SimpleNamespace(big_page_ids=[1]), - linear_att_len_to_big_page_id={16: 3}, - tail_linear_att_small_page_buffer_id=0 if token_num % 8 else None, - ) - page_num = (token_num + 7) // 8 - ready = torch.tensor([True] + [False] * (page_num - 1)) - operator.offload_gpu_kv_to_cpu_cache( - torch.arange(token_num, dtype=torch.int32), - torch.arange(page_num, dtype=torch.int32), - ready, - SimpleNamespace(cpu_kv_cache_tensor=object()), - req, - ) - - assert len(captures) == 1 - assert captures[0]["big_page_buffer_ids"].tolist() == ([1, 3, 5] if token_num % 8 else [1, 3]) - assert captures[0]["mem_indexes"].tolist() == list(range(token_num)) + [-1] * (page_num * 8 - token_num) - assert captures[0]["page_readies"] is ready - assert req.linear_att_len_to_big_page_id == {16: 3} - assert torch.count_nonzero(big_pages.get_state_cache(4)) == 0 - if token_num % 8: - torch.testing.assert_close(big_pages.get_state_cache(5), small_pages.get_state_cache(0), atol=0, rtol=0) - small_pages.get_state_cache(0).zero_() - assert torch.all(big_pages.get_state_cache(5) == 17) - - -def test_reserved_state_slots_are_never_allocated_or_freed(): - pages = _cpu_pages(4, keep_num=2) - assert pages.alloc_state_cache(2) == [0, 1] - assert pages.alloc_one_state_cache() is None - for reserved_id in [2, 3]: - with pytest.raises(AssertionError): - pages.free_state_cache([reserved_id]) - pages.free_state_cache([0, 1]) - assert pages.get_free_cache_num() == 2 - assert pages.get_used_cache_num() == 2 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_real_cpu_transfers_reuse_separate_load_and_offload_slots_across_streams(monkeypatch): - request_num, page_size, tail_len, window = 8, 8, 6, 4 - config = SlidingWindowCacheConfig({0: 0}, {1: 0}, window, 1, 8, 1, 8, torch.bfloat16) - full_bytes = config.get_cpu_cache_full_att_bytes(page_size, 1) - state_bytes = config.get_cpu_cache_state_bytes(1) - cpu_cache = torch.zeros( - (request_num * 2, config.get_cpu_cache_big_page_bytes(page_size, 1)), - dtype=torch.uint8, - device="cpu", - pin_memory=True, - ) - - def full_page(page_id): - return cpu_cache[page_id, :full_bytes].view(config.dtype).view(page_size, 1, 2, 8) - - def window_page(page_id): - return ( - cpu_cache[page_id, full_bytes : full_bytes + state_bytes].view(config.dtype).view(config.get_state_shape()) - ) - - # The load stream reads already-ready pages, while offload writes disjoint - # CPU pages. Both directions share the same big-state pool, as in serving. - for req_idx in range(request_num): - full_page(req_idx).fill_(10 + req_idx) - window_page(req_idx).fill_(200 + req_idx) - - big_pages = SlidingWindowStateCacheManager(2, config, keep_num=2) - small_pages = SlidingWindowStateCacheManager(request_num, config) - manager = SimpleNamespace( - sliding_config=config, - kv_buffer=torch.zeros((1, request_num * tail_len * 2, 2, 8), dtype=config.dtype, device="cuda"), - linear_att_big_page_buffers=big_pages, - CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID=0, - CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID=1, - ) - req_manager = object.__new__(ReqManagerForSlidingWindow) - req_manager.mem_manager, req_manager.sliding_window = manager, window - req_manager.req_to_sliding_window = torch.zeros((1, request_num, window, 2, 8), dtype=config.dtype, device="cuda") - monkeypatch.setattr(g_infer_context, "req_manager", req_manager) - monkeypatch.setattr( - g_infer_context, - "radix_cache", - SimpleNamespace(linear_att_small_page_buffers=small_pages, get_big_page_ids_by_node=lambda node: []), - ) - monkeypatch.setattr( - operator_module, - "get_env_start_args", - lambda: SimpleNamespace( - cpu_cache_token_page_size=page_size, linear_att_hash_page_size=2, linear_att_page_block_num=4 - ), - ) - monkeypatch.setattr(operator_module, "get_current_rank_in_dp", lambda: 0) - monkeypatch.setattr(operator_module, "get_dp_world_size", lambda: 1) - operator = operator_module.HybridSlidingMemOperator(manager) - client = SimpleNamespace(cpu_kv_cache_tensor=cpu_cache) - transfers = [] - for req_idx in range(request_num): - start = req_idx * tail_len - manager.kv_buffer[:, start : start + tail_len].fill_(30 + req_idx) - small_page_id = small_pages.alloc_one_state_cache() - small_pages.get_state_cache(small_page_id).fill_(100 + req_idx) - req = SimpleNamespace( - req_idx=req_idx, - cur_kv_len=tail_len, - shared_kv_node=None, - linear_att_len_to_big_page_id={}, - tail_linear_att_small_page_buffer_id=small_page_id, - ) - source_indexes = torch.arange(start, start + tail_len, dtype=torch.int32, device="cuda") - load_indexes = source_indexes + request_num * tail_len - load_page = torch.tensor([req_idx], dtype=torch.int32, device="cuda") - offload_page = torch.tensor([request_num + req_idx], dtype=torch.int32, device="cuda") - ready = torch.tensor([False], dtype=torch.bool, device="cuda") - transfers.append((req, source_indexes, load_indexes, load_page, offload_page, ready)) - - offload_stream, load_stream = torch.cuda.Stream(), torch.cuda.Stream() - offload_stream.wait_stream(torch.cuda.current_stream()) - load_stream.wait_stream(torch.cuda.current_stream()) - for req, source_indexes, load_indexes, load_page, offload_page, ready in transfers: - with torch.cuda.stream(offload_stream): - operator.offload_gpu_kv_to_cpu_cache(source_indexes, offload_page, ready, client, req) - with torch.cuda.stream(load_stream): - operator.load_cpu_cache_to_gpu(load_indexes, load_page, client, req) - # No per-request wait: each reserved slot has been reused eight times. - offload_stream.synchronize() - load_stream.synchronize() - - for req_idx, (req, source_indexes, load_indexes, _, _, _) in enumerate(transfers): - assert req.linear_att_len_to_big_page_id == {} - torch.testing.assert_close( - manager.kv_buffer[:, load_indexes], - torch.full((1, tail_len, 2, 8), 10 + req_idx, dtype=config.dtype, device="cuda"), - atol=0, - rtol=0, - ) - torch.testing.assert_close( - req_manager.req_to_sliding_window[:, req_idx], - torch.full(config.get_state_shape(), 200 + req_idx, dtype=config.dtype, device="cuda"), - atol=0, - rtol=0, - ) - assert torch.all(full_page(request_num + req_idx)[:tail_len] == 30 + req_idx) - assert torch.count_nonzero(full_page(request_num + req_idx)[tail_len:]) == 0 - assert torch.all(window_page(request_num + req_idx) == 100 + req_idx) - assert torch.all(full_page(req_idx) == 10 + req_idx) - assert torch.all(window_page(req_idx) == 200 + req_idx) - assert torch.all(manager.kv_buffer[:, source_indexes] == 30 + req_idx) - assert big_pages.get_free_cache_num() == 0 - assert big_pages.alloc_one_state_cache() is None - for slot in [manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID, manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID]: - with pytest.raises(AssertionError): - big_pages.free_state_cache([slot]) From 69dcee8ca527b045f36755d05d36b1b7c474ccd7 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Wed, 9 Sep 2026 07:01:42 +0000 Subject: [PATCH 11/14] refactor: streamline hybrid sliding-window cache lifecycle --- lightllm/common/basemodel/basemodel.py | 13 +- lightllm/common/basemodel/infer_struct.py | 8 + .../sliding_window_cpu_cache_copy.py | 105 +++++++---- .../triton_kernel/sliding_window_state.py | 178 +++++++++--------- .../hybrid_sliding_mem_manager.py | 123 ++++++------ .../operator/hybrid_sliding.py | 27 +-- lightllm/common/req_manager/hybrid_att.py | 18 +- lightllm/common/req_manager/linear_att.py | 27 +-- lightllm/common/req_manager/sliding_window.py | 116 ++---------- .../sliding_window_cache_manager/config.py | 14 +- .../state_cache.py | 7 +- lightllm/models/gemma4/infer_struct.py | 38 +++- lightllm/models/gemma4/kv_layout.py | 3 +- .../layer_infer/transformer_layer_infer.py | 76 +++----- lightllm/models/gemma4/model.py | 60 +----- .../context_attention_fwd_gemma4_mm.py | 44 ++--- .../server/router/model_infer/infer_batch.py | 4 +- lightllm/utils/kv_cache_utils.py | 36 ++-- .../test_sliding_window_cpu_cache_copy.py | 52 ++++- test/kernel/test_sliding_window_prefill.py | 49 +++-- test/kernel/test_sliding_window_state.py | 57 +++--- 21 files changed, 486 insertions(+), 569 deletions(-) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e80f2b552f..a3ae29e144 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -604,6 +604,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, @@ -674,11 +675,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" @@ -847,6 +848,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, @@ -936,8 +939,6 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 infer_state1=infer_state1, ) - 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) else: model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) @@ -962,9 +963,11 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 infer_state1.init_att_state() model_output0, model_output1 = self._overlap_tpsp_token_forward(infer_state0, infer_state1=infer_state1) - 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() + 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) 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 index 0769efdcd8..cec6227d35 100644 --- a/lightllm/common/basemodel/triton_kernel/sliding_window_cpu_cache_copy.py +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_cpu_cache_copy.py @@ -2,8 +2,6 @@ import triton import triton.language as tl -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig - @triton.jit def _copy_sliding_window_cpu_cache( @@ -12,7 +10,7 @@ def _copy_sliding_window_cpu_cache( page_readies, big_page_buffer_ids, gpu_full_att_kv_state, - gpu_sliding_state, + cpu_kv_sliding_state, cpu_cache, page_num, full_stride_l, @@ -64,14 +62,14 @@ def _copy_sliding_window_cpu_cache( 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 - gpu_ptr = gpu_sliding_state + big_page * WINDOW_RANK_SIZE + offsets + 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(gpu_ptr, valid, other=0) + 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(gpu_ptr, value, valid) + tl.store(state_ptr, value, valid) def _copy_state_cache( @@ -80,12 +78,11 @@ def _copy_state_cache( page_readies, big_page_buffer_ids, gpu_full_att_kv_state, - gpu_sliding_state, + cpu_kv_sliding_state, cpu_cache_tensor, tp_rank, tp_world_size, big_page_token_num, - sliding_config, offload, grid_num, ): @@ -93,34 +90,24 @@ def _copy_state_cache( 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 - assert 0 <= tp_rank < tp_world_size - assert big_page_token_num > 0 and grid_num > 0 if page_num == 0: return - assert gpu_full_att_kv_state.shape[0] == sliding_config.full_layer_num - assert gpu_full_att_kv_state.shape[2:] == ( - 2 * sliding_config.full_head_num, - sliding_config.full_head_dim, - ) - assert gpu_sliding_state.shape[1:] == sliding_config.get_state_shape() - assert gpu_full_att_kv_state.dtype == gpu_sliding_state.dtype == sliding_config.dtype - assert gpu_full_att_kv_state.is_contiguous() and gpu_sliding_state.is_contiguous() + assert gpu_full_att_kv_state.is_contiguous() and cpu_kv_sliding_state.is_contiguous() assert cpu_cache_tensor.is_contiguous() - cpu_cache = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint8) - full_bytes = sliding_config.get_cpu_cache_full_att_bytes(big_page_token_num, tp_world_size) - window_bytes = sliding_config.get_cpu_cache_state_bytes(tp_world_size) - assert cpu_cache.shape[1] == sliding_config.get_cpu_cache_big_page_bytes(big_page_token_num, tp_world_size) # Packing preserves the original bit patterns, including BF16/FP16 NaNs. - # Gemma's K+V head rows and checkpoint tensors are all uint64-aligned. + # 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 = gpu_sliding_state.flatten(1).view(torch.uint64) - cpu_cache = cpu_cache.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] - assert full_rank_size * tp_world_size * 8 == full_bytes - assert window_rank_size * tp_world_size * 8 == window_bytes + 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, @@ -135,10 +122,10 @@ def _copy_state_cache( full_state.stride(1), cpu_cache.stride(0), tp_rank, - FULL_LAYER_NUM=sliding_config.full_layer_num, + 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_bytes // 8, + 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, @@ -152,27 +139,25 @@ def copy_kv_buffer_to_cpu_cache( page_readies: torch.Tensor, big_page_buffer_ids: torch.Tensor, gpu_full_att_kv_state: torch.Tensor, - gpu_sliding_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, - sliding_config: SlidingWindowCacheConfig, grid_num: int = 12, ): - """Pack full KV and a raw ring checkpoint into this TP rank's CPU page slices.""" + """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, - gpu_sliding_state, + cpu_kv_sliding_state, cpu_cache_tensor, tp_rank, tp_world_size, big_page_token_num, - sliding_config, offload=True, grid_num=grid_num, ) @@ -183,27 +168,69 @@ def copy_cpu_cache_to_kv_buffer( page_indexes: torch.Tensor, big_page_buffer_ids: torch.Tensor, gpu_full_att_kv_state: torch.Tensor, - gpu_sliding_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, - sliding_config: SlidingWindowCacheConfig, grid_num: int = 12, ): - """Restore full KV and the unchanged ring checkpoint from a packed CPU page.""" + """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, - gpu_sliding_state, + cpu_kv_sliding_state, cpu_cache_tensor, tp_rank, tp_world_size, big_page_token_num, - sliding_config, 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 index 2b57246682..9871bab1c5 100644 --- a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py @@ -1,132 +1,130 @@ +import math import torch import triton import triton.language as tl @triton.jit -def _move_sliding_window( - Pool, +def _build_sliding_window_page_table( + PageTable, + BKVStart, BReqIdx, BSeqLen, BReadyCacheLen, BQStartLoc, - runtime_token_start, - stride_layer, - stride_token, + MemIndexes, + table_width, WINDOW: tl.constexpr, - KV_DIM: tl.constexpr, - COMPACT: tl.constexpr, BLOCK: tl.constexpr, ): - batch_idx = tl.program_id(0) - layer_idx = tl.program_id(1).to(tl.int64) - offsets = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) - req_idx = tl.load(BReqIdx + batch_idx).to(tl.int64) - history_len = tl.load(BReadyCacheLen + batch_idx).to(tl.int64) - q_start = tl.load(BQStartLoc + batch_idx).to(tl.int64) - current_start = tl.cast(runtime_token_start, tl.int64) + q_start + (batch_idx + 1) * WINDOW - if COMPACT: - end = tl.load(BSeqLen + batch_idx).to(tl.int64) - else: - end = history_len - positions = end - WINDOW + offsets // KV_DIM - canonical = req_idx * WINDOW + positions % WINDOW - active = current_start + positions - history_len - if COMPACT: - src, dst = active, canonical - else: - src, dst = canonical, active - layer_offset = layer_idx * tl.cast(stride_layer, tl.int64) - mask = (offsets < WINDOW * KV_DIM) & (positions >= 0) - values = tl.load(Pool + layer_offset + src * stride_token + offsets % KV_DIM, mask=mask, other=0) - tl.store(Pool + layer_offset + dst * stride_token + offsets % KV_DIM, values, mask=mask) + batch = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + req_idx = tl.load(BReqIdx + batch) + history = tl.load(BReadyCacheLen + batch) + seq_len = tl.load(BSeqLen + batch) + q_start = tl.load(BQStartLoc + batch) + kv_start = tl.maximum(history - WINDOW, 0) + positions = kv_start + offsets + is_new = positions >= history + new_index = tl.load(MemIndexes + q_start + positions - history, mask=is_new & (positions < seq_len), other=0) + index = tl.where(is_new, new_index, req_idx * WINDOW + positions % WINDOW) + tl.store( + PageTable + batch * table_width + offsets, tl.where(positions < seq_len, index, -1), mask=offsets < table_width + ) + if tl.program_id(1) == 0: + tl.store(BKVStart + batch, kv_start) @torch.no_grad() -def move_sliding_window( - pool, - b_req_idx, - b_seq_len, - b_ready_cache_len, - b_q_start_loc, - window, - runtime_token_start, - compact=False, +def build_sliding_window_page_table( + b_req_idx, b_seq_len, b_ready_cache_len, b_q_start_loc, mem_indexes, window, max_q_seq_len ): - """Move every owner's window between canonical rings and the prefill region.""" - assert pool.ndim == 4 and pool.stride(-1) == 1 and pool.stride(-2) == pool.shape[-1] - if not b_req_idx.numel(): - return - kv_dim = pool.shape[2] * pool.shape[3] - grid = (b_req_idx.numel(), pool.shape[0], triton.cdiv(window * kv_dim, 1024)) - _move_sliding_window[grid]( - pool, + """Batch-local, chronological table over fixed ring history and this prefill's new KV. + + Column zero represents b_kv_start_pos, not absolute token zero. No KV is moved. + """ + page_table = torch.empty((b_req_idx.numel(), window + max_q_seq_len), dtype=torch.int32, device=b_req_idx.device) + b_kv_start_pos = torch.empty_like(b_req_idx) + _build_sliding_window_page_table[(b_req_idx.numel(), triton.cdiv(page_table.shape[1], 256))]( + page_table, + b_kv_start_pos, b_req_idx, b_seq_len, b_ready_cache_len, b_q_start_loc, - runtime_token_start, - pool.stride(0), - pool.stride(1), + mem_indexes, + page_table.shape[1], WINDOW=window, - KV_DIM=kv_dim, - COMPACT=compact, - BLOCK=1024, + BLOCK=256, num_warps=4, ) + return page_table, b_kv_start_pos @triton.jit -def _get_sliding_window_mem_indexes( - Out, +def _commit_sliding_window_kv( + Pool, + MemIndexes, BReqIdx, BSeqLen, - BQSeqLen, + BReadyCacheLen, BQStartLoc, - runtime_token_start, + stride_layer, + stride_token, WINDOW: tl.constexpr, - IS_PREFILL: tl.constexpr, + KV_DIM: tl.constexpr, BLOCK: tl.constexpr, ): - batch_idx = tl.program_id(0) - if IS_PREFILL: - q_start = tl.load(BQStartLoc + batch_idx).to(tl.int64) - q_len = tl.load(BQSeqLen + batch_idx) - offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - current_start = tl.cast(runtime_token_start, tl.int64) + q_start + (batch_idx + 1) * WINDOW - tl.store(Out + q_start + offsets, current_start + offsets, mask=offsets < q_len) - else: - req_idx = tl.load(BReqIdx + batch_idx).to(tl.int64) - seq_len = tl.load(BSeqLen + batch_idx).to(tl.int64) - tl.store(Out + batch_idx, req_idx * WINDOW + (seq_len - 1) % WINDOW) + batch, layer = tl.program_id(0), tl.program_id(1) + offsets = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + req_idx = tl.load(BReqIdx + batch).to(tl.int64) + history = tl.load(BReadyCacheLen + batch) + seq_len = tl.load(BSeqLen + batch) + q_start = tl.load(BQStartLoc + batch) + positions = seq_len - WINDOW + offsets // KV_DIM + # Retained old history is already in place. Copy only this chunk's newest W KV. + mask = (offsets < WINDOW * KV_DIM) & (positions >= history) + src_index = tl.load(MemIndexes + q_start + positions - history, mask=mask, other=0).to(tl.int64) + dst_index = req_idx * WINDOW + positions % WINDOW + layer_ptr = Pool + layer.to(tl.int64) * stride_layer + values = tl.load(layer_ptr + src_index * stride_token + offsets % KV_DIM, mask=mask, other=0) + tl.store(layer_ptr + dst_index * stride_token + offsets % KV_DIM, values, mask=mask) @torch.no_grad() -def get_sliding_window_mem_indexes( - b_req_idx, - b_seq_len, - b_q_seq_len, - b_q_start_loc, - window, - runtime_token_start, - total_token_num, - max_q_seq_len, - is_prefill, -): - indexes = torch.empty(total_token_num, dtype=torch.int32, device=b_req_idx.device) - if not b_req_idx.numel(): - return indexes - grid = (b_req_idx.numel(), triton.cdiv(max_q_seq_len, 256) if is_prefill else 1) - _get_sliding_window_mem_indexes[grid]( - indexes, +def commit_sliding_window_kv(pool, mem_indexes, b_req_idx, b_seq_len, b_ready_cache_len, b_q_start_loc, window): + """Commit [layer, slot, ...payload] tails after every prefill reader has finished. + + The contiguous per-token payload can be KV heads, an MLA vector or packed bytes. + """ + kv_dim = math.prod(pool.shape[2:]) + _commit_sliding_window_kv[(b_req_idx.numel(), pool.shape[0], triton.cdiv(window * kv_dim, 1024))]( + pool, + mem_indexes, b_req_idx, b_seq_len, - b_q_seq_len, + b_ready_cache_len, b_q_start_loc, - runtime_token_start, + pool.stride(0), + pool.stride(1), WINDOW=window, - IS_PREFILL=is_prefill, - BLOCK=256, - num_warps=4 if is_prefill else 1, + KV_DIM=kv_dim, + BLOCK=1024, + num_warps=4, ) + + +@triton.jit +def _get_sliding_window_decode_indexes(Out, BReqIdx, BSeqLen, WINDOW: tl.constexpr): + batch = tl.program_id(0) + req_idx = tl.load(BReqIdx + batch) + seq_len = tl.load(BSeqLen + batch) + tl.store(Out + batch, req_idx * WINDOW + (seq_len - 1) % WINDOW) + + +@torch.no_grad() +def get_sliding_window_decode_indexes(b_req_idx, b_seq_len, window): + """Single-token decode writes directly into each request's fixed ring.""" + indexes = torch.empty_like(b_req_idx) + _get_sliding_window_decode_indexes[(b_req_idx.numel(),)](indexes, b_req_idx, b_seq_len, WINDOW=window, num_warps=1) return indexes 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 index 0e15b5d329..df906dbae0 100644 --- a/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -1,17 +1,13 @@ import torch -import torch.distributed as dist import triton from lightllm.common.sliding_window_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 lightllm.utils.log_utils import init_logger -from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory from .mem_manager import MemoryManager from .operator.hybrid_sliding import HybridSlidingMemOperator -logger = init_logger(__name__) - class HybridSlidingMemoryManager(MemoryManager): """Token-granular full KV plus request-granular sliding-window KV.""" @@ -19,11 +15,25 @@ class HybridSlidingMemoryManager(MemoryManager): operator_class = HybridSlidingMemOperator def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): - self.sliding_config = sliding_config args = get_env_start_args() - self.enable_prompt_cache = not args.disable_dynamic_prompt_cache - self.small_page_num = args.linear_att_cache_size if self.enable_prompt_cache else 0 - self.cpu_cache_temp_page_num = 2 if args.enable_cpu_cache else 0 + self.sliding_config = sliding_config + self.sliding_prefill_start = (args.running_max_req_size + 1) * sliding_config.sliding_window + # Both microbatches share batch_max_tokens; allow TP/dummy padding for each. + self.max_sliding_prefill_tokens = args.batch_max_tokens + 2 * get_dp_world_size() + self._sliding_prefill_used = 0 + self._sliding_prefill_batches = 0 + # One layer-first pool: fixed request rings followed by this batch's new KV. + # Reserve it before profiling how much memory can be given to full attention. + self.sliding_kv_buffer = torch.zeros( + ( + sliding_config.sliding_layer_num, + self.sliding_prefill_start + self.max_sliding_prefill_tokens, + 2 * sliding_config.sliding_head_num, + sliding_config.sliding_head_dim, + ), + dtype=sliding_config.dtype, + device="cuda", + ) self.big_page_token_num = args.linear_att_page_block_num * args.linear_att_hash_page_size super().__init__( size=size, @@ -35,70 +45,57 @@ def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): mem_fraction=mem_fraction, ) - def profile_size(self, mem_fraction): - torch.cuda.empty_cache() - world_size = dist.get_world_size() - available_memory = get_available_gpu_memory(world_size) - if self.size is None: - available_memory -= get_total_gpu_memory() * (1 - mem_fraction) - available_bytes = int(available_memory * 1024 ** 3) - cell_size = self.get_cell_size() - state_bytes = self.sliding_config.get_state_nbytes() - # Runtime windows already exist. Reserve the hold token, small pages - # and CPU-transfer slots before sizing full KV and big checkpoints. - fixed_bytes = cell_size + (self.small_page_num + self.cpu_cache_temp_page_num) * state_bytes - big_page_state_bytes = state_bytes if self.enable_prompt_cache else 0 - if self.size is None: - if available_bytes < fixed_bytes + cell_size + big_page_state_bytes: - raise ValueError( - "Insufficient GPU memory for sliding-window checkpoints and full KV; " - "reduce --linear_att_cache_size or --running_max_req_size." - ) - # Each complete page costs B full-KV tokens plus one checkpoint. - # A partial page also needs one checkpoint before it can hold tokens. - page_bytes = self.big_page_token_num * cell_size + big_page_state_bytes - page_num, tail_bytes = divmod(available_bytes - fixed_bytes, page_bytes) - self.size = page_num * self.big_page_token_num + max(0, (tail_bytes - big_page_state_bytes) // cell_size) - if world_size > 1: - size_tensor = torch.tensor(self.size, dtype=torch.int64, device="cuda") - dist.all_reduce(size_tensor, op=dist.ReduceOp.MIN) - self.size = size_tensor.item() + def alloc_sliding_prefill(self, token_num: int) -> torch.Tensor: + # Overlapping microbatches lease disjoint ranges of the same token budget. + assert ( + self._sliding_prefill_used + token_num <= self.max_sliding_prefill_tokens + ), "sliding prefill pool exhausted" + start = self.sliding_prefill_start + self._sliding_prefill_used + indexes = torch.arange(start, start + token_num, dtype=torch.int32, device="cuda") + self._sliding_prefill_used += token_num + self._sliding_prefill_batches += 1 + return indexes - big_page_num = triton.cdiv(self.size, self.big_page_token_num) if self.enable_prompt_cache else 0 - cache_bytes = fixed_bytes + self.size * cell_size + big_page_num * state_bytes - if cache_bytes > available_bytes: - raise ValueError( - "Requested full KV and sliding-window checkpoints exceed available GPU memory; " - "reduce --max_total_token_num, --linear_att_cache_size or --running_max_req_size." - ) - logger.info( - f"Sliding-window cache budget: {self.size} full-KV tokens, " - f"{big_page_num} big pages, {self.small_page_num} small pages, " - f"{self.cpu_cache_temp_page_num} CPU-cache staging states, " - f"{cache_bytes / 1024 ** 3:.2f} GiB (runtime windows already allocated)" - ) + def free_sliding_prefill(self): + assert self._sliding_prefill_batches > 0 + self._sliding_prefill_batches -= 1 + if self._sliding_prefill_batches == 0: + self._sliding_prefill_used = 0 + + def free_all(self): + super().free_all() + # Also discard leases when warmup/error cleanup resets all requests. + self._sliding_prefill_used = 0 + self._sliding_prefill_batches = 0 def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): super()._init_buffers(size, dtype, head_num, head_dim, layer_num) - big_page_num = triton.cdiv(size, self.big_page_token_num) if self.enable_prompt_cache else 0 - # Keep the existing radix-cache contract; no second alias is needed. + # Match linear attention: CPU checkpoints plus two reserved tail-transfer slots. self.linear_att_big_page_buffers = SlidingWindowStateCacheManager( - size=big_page_num + self.cpu_cache_temp_page_num, - sliding_config=self.sliding_config, - keep_num=self.cpu_cache_temp_page_num, - ) - if self.cpu_cache_temp_page_num: - self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.linear_att_big_page_buffers.size - 2 - self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.linear_att_big_page_buffers.size - 1 - self.sliding_small_page_buffers = SlidingWindowStateCacheManager( - size=self.small_page_num, + 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.linear_att_big_page_buffers.size - 2 + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.linear_att_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.linear_att_big_page_buffers + self.linear_att_big_page_buffers = None + try: + return super().write_to_shm(req_manager) + finally: + self.linear_att_big_page_buffers = big_page_buffers def get_att_input_params(self, layer_index: int): - return super().get_att_input_params(self.sliding_config.get_full_layer_index(layer_index)) + 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.linear_att_big_page_buffers = None - self.sliding_small_page_buffers = None diff --git a/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py index da27b7efec..e47de29c79 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid_sliding.py @@ -2,7 +2,6 @@ import triton from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size -from lightllm.utils.envs_utils import get_env_start_args from .normal import NormalMemOperator @@ -11,7 +10,7 @@ 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.get_full_layer_index(layer_index) + 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): @@ -20,16 +19,11 @@ def load_cpu_cache_to_gpu(self, mem_indexes, page_indexes, cpu_cache_client, req ) from lightllm.server.router.model_infer.infer_batch import g_infer_context - args = get_env_start_args() - page_size = args.cpu_cache_token_page_size - assert mem_indexes.is_cuda and page_indexes.is_cuda - assert page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert triton.cdiv(len(mem_indexes), page_size) == len(page_indexes) 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 = [] @@ -54,12 +48,11 @@ def load_cpu_cache_to_gpu(self, mem_indexes, page_indexes, cpu_cache_client, req page_indexes=page_indexes, big_page_buffer_ids=big_page_ids_gpu, gpu_full_att_kv_state=mem_manager.kv_buffer, - gpu_sliding_state=mem_manager.linear_att_big_page_buffers.state_cache, + cpu_kv_sliding_state=mem_manager.linear_att_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, - sliding_config=mem_manager.sliding_config, ) # Loads and this restore use the inference stream. The next load may # reuse its reserved slot only after this copy has been queued. @@ -68,19 +61,15 @@ def load_cpu_cache_to_gpu(self, mem_indexes, page_indexes, cpu_cache_client, 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 - args = get_env_start_args() - page_size = args.cpu_cache_token_page_size - assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda - assert page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert triton.cdiv(len(mem_indexes), page_size) == len(page_indexes) == len(page_readies) 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 @@ -96,10 +85,9 @@ def offload_gpu_kv_to_cpu_cache(self, mem_indexes, page_indexes, page_readies, c src_state = radix_cache.linear_att_small_page_buffers.get_state_cache( req.tail_linear_att_small_page_buffer_id ) - mem_manager.linear_att_big_page_buffers.get_state_cache(temp_id).copy_(src_state, non_blocking=True) + copy_sliding_window_state(src_state, mem_manager.linear_att_big_page_buffers.get_state_cache(temp_id)) big_page_ids.append(temp_id) - assert len(big_page_ids) == len(page_indexes) 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 @@ -110,12 +98,11 @@ def offload_gpu_kv_to_cpu_cache(self, mem_indexes, page_indexes, page_readies, c page_readies=page_readies, big_page_buffer_ids=big_page_ids_gpu, gpu_full_att_kv_state=mem_manager.kv_buffer, - gpu_sliding_state=mem_manager.linear_att_big_page_buffers.state_cache, + cpu_kv_sliding_state=mem_manager.linear_att_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, - sliding_config=mem_manager.sliding_config, ) def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): diff --git a/lightllm/common/req_manager/hybrid_att.py b/lightllm/common/req_manager/hybrid_att.py index 8f37ee07df..3fd36e5683 100644 --- a/lightllm/common/req_manager/hybrid_att.py +++ b/lightllm/common/req_manager/hybrid_att.py @@ -28,18 +28,22 @@ def create_state_cache_manager(self, size: int): def init_hybrid_attention_state(self, req: "InferReq"): """Initialize request runtime state when no prefix cache is restored.""" - @abstractmethod def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - """Restore runtime state from a big-page checkpoint.""" + self.restore_state(req, self.mem_manager.linear_att_big_page_buffers, big_page_buffer_idx) - @abstractmethod def restore_small_page_state(self, req: "InferReq", small_page_buffers): - """Restore runtime state from a small-page checkpoint.""" + self.restore_state(req, small_page_buffers, req.shared_kv_node.small_page_buffer_idx) @abstractmethod + def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): + """Restore the same request-state payload from either checkpoint pool.""" + def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - """Save selected checkpoints; CPU request IDs avoid device-to-host synchronization.""" + """Default checkpoint copies; models may override with a batched kernel.""" + for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): + if buffer_idx != -1: + self.save_state(req_idx, buffer_idx, self.mem_manager.linear_att_big_page_buffers) @abstractmethod - def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffers): - """Save a request's final small-page checkpoint in the layout's storage.""" + def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager): + """Save a request's payload into either checkpoint pool.""" diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index e6beeea461..50f09fa44e 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -54,10 +54,6 @@ def create_state_cache_manager(self, size: int): return LinearAttCacheManager(size=size, linear_config=self.linear_config) def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - assert len(b_req_idx) == len(buffer_indexes) - if not any(buffer_idx != -1 for buffer_idx in buffer_indexes): - return - from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) @@ -73,12 +69,12 @@ def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], ) return - def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffers: LinearAttCacheManager): + def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearAttCacheManager): # Preserve main's small-page copies, including the MTP conv-state crop. conv_cache_width = self.linear_config.get_conv_state_shape()[-1] gpu_conv_state = self.req_to_conv_state.buffer[:, req_idx, ..., :conv_cache_width] gpu_ssm_state = self.req_to_ssm_state.buffer[:, req_idx * (self.mtp_step + 1), ...] - dst_conv_state, dst_ssm_state = small_page_buffers.get_state_cache(buffer_idx=buffer_idx) + dst_conv_state, dst_ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) dst_conv_state.copy_(gpu_conv_state, non_blocking=True) dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) @@ -102,26 +98,11 @@ def get_mamba_cache(self, layer_idx_in_all: int): ssm_states = self.req_to_ssm_state.buffer[layer_idx_in_linear] return conv_states, ssm_states - def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - big_page_buffers: LinearAttCacheManager = self.mem_manager.linear_att_big_page_buffers - - conv_state, ssm_state = big_page_buffers.get_state_cache(buffer_idx=big_page_buffer_idx) - conv_dest = req.req_idx - ssm_dest = req.req_idx * (self.mtp_step + 1) - conv_cache_width = conv_state.shape[-1] - self.req_to_conv_state.buffer[:, conv_dest, ..., :conv_cache_width] = conv_state - self.req_to_ssm_state.buffer[:, ssm_dest, ...] = ssm_state - if self.req_to_mtp_state_index is not None: - self.req_to_mtp_state_index[req.req_idx] = 0 - return - - def restore_small_page_state(self, req: "InferReq", small_page_buffers: LinearAttCacheManager): - conv_state, ssm_state = small_page_buffers.get_state_cache(buffer_idx=req.shared_kv_node.small_page_buffer_idx) + def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheManager, buffer_idx: int): + conv_state, ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) conv_dest = req.req_idx ssm_dest = req.req_idx * (self.mtp_step + 1) conv_cache_width = conv_state.shape[-1] - # TODO 下面这个从 cpu cache 拷贝数据的 gpu的操作，是否是阻塞的操作。 - # 同时，非连续对象的拷贝，可能存在效率问题。 self.req_to_conv_state.buffer[:, conv_dest, ..., :conv_cache_width] = conv_state self.req_to_ssm_state.buffer[:, ssm_dest, ...] = ssm_state if self.req_to_mtp_state_index is not None: diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index fb633144aa..8522da63fa 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -1,11 +1,6 @@ -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, Optional -import torch - -from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( - get_sliding_window_mem_indexes, - move_sliding_window, -) +from lightllm.common.basemodel.triton_kernel.sliding_window_cpu_cache_copy import copy_sliding_window_state from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager from .hybrid_att import HybridAttentionReqManager @@ -26,112 +21,31 @@ def __init__( max_sequence_length: int, mem_manager: Optional["HybridSlidingMemoryManager"], sliding_config: "SlidingWindowCacheConfig", - max_prefill_token_num: int, - prefill_microbatch_num: int = 1, ): super().__init__(max_request_num, max_sequence_length, mem_manager) self.sliding_config = sliding_config self.sliding_window = sliding_config.sliding_window - self.max_prefill_token_num = max_prefill_token_num - self.runtime_token_start = (max_request_num + 1) * self.sliding_window - # One KV pool. Prefill expands each active window to history + chunk; - # decode uses its compact request window. Both attention paths read this pool. - self.prefill_capacity = ( - max_prefill_token_num + min(max_request_num + 1, max_prefill_token_num) * self.sliding_window - ) - self.sliding_kv_buffer = torch.zeros( - ( - sliding_config.sliding_layer_num, - self.runtime_token_start + prefill_microbatch_num * self.prefill_capacity, - 2 * sliding_config.sliding_head_num, - sliding_config.sliding_head_dim, - ), - dtype=sliding_config.dtype, - device="cuda", - ) - self.req_to_sliding_window = self.sliding_kv_buffer[:, : self.runtime_token_start].view( - sliding_config.sliding_layer_num, - max_request_num + 1, - self.sliding_window, - 2 * sliding_config.sliding_head_num, - sliding_config.sliding_head_dim, - ) + + @property + def req_to_sliding_window(self): + # A view, not a second allocation. req_idx owns the same W slots for its lifetime. + pool = self.mem_manager.sliding_kv_buffer + return pool[:, : self.mem_manager.sliding_prefill_start].unflatten(1, (-1, self.sliding_window)) def create_state_cache_manager(self, size: int): - # Allocated with full KV and big pages, within the same GPU budget. - return self.mem_manager.sliding_small_page_buffers + return SlidingWindowStateCacheManager(size=size, sliding_config=self.sliding_config) def init_hybrid_attention_state(self, req: "InferReq"): self.req_to_sliding_window[:, req.req_idx].zero_() - def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - self._restore_state(req.req_idx, self.mem_manager.linear_att_big_page_buffers, big_page_buffer_idx) - - def restore_small_page_state(self, req: "InferReq", small_page_buffers): - self._restore_state(req.req_idx, small_page_buffers, req.shared_kv_node.small_page_buffer_idx) - - def _restore_state(self, req_idx: int, state_cache_manager, buffer_idx: int): - self.req_to_sliding_window[:, req_idx].copy_( + def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): + copy_sliding_window_state( state_cache_manager.get_state_cache(buffer_idx), - non_blocking=True, + self.req_to_sliding_window[:, req.req_idx], ) - def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - assert len(req_indexes) == len(buffer_indexes) - for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): - if buffer_idx == -1: - continue - self.save_small_page_state(req_idx, buffer_idx, self.mem_manager.linear_att_big_page_buffers) - - def save_small_page_state(self, req_idx: int, buffer_idx: int, small_page_buffers: SlidingWindowStateCacheManager): - small_page_buffers.get_state_cache(buffer_idx).copy_( + def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: SlidingWindowStateCacheManager): + copy_sliding_window_state( self.req_to_sliding_window[:, req_idx], - non_blocking=True, - ) - - def prepare_sliding_window(self, infer_state): - token_num = infer_state.input_ids.shape[0] - infer_state.sliding_window_runtime_start = ( - self.runtime_token_start + infer_state.microbatch_index * self.prefill_capacity - ) - if infer_state.is_prefill: - assert token_num <= self.max_prefill_token_num - move_sliding_window( - self.sliding_kv_buffer, - infer_state.b_req_idx, - infer_state.b_seq_len, - infer_state.b_ready_cache_len, - infer_state.b_q_start_loc, - self.sliding_window, - infer_state.sliding_window_runtime_start, - ) - infer_state.sliding_window_mem_index = get_sliding_window_mem_indexes( - infer_state.b_req_idx, - infer_state.b_seq_len, - infer_state.b_q_seq_len, - infer_state.b_q_start_loc, - self.sliding_window, - infer_state.sliding_window_runtime_start, - token_num, - infer_state.max_q_seq_len, - infer_state.is_prefill, - ) - - def finish_prefill(self, infer_state): - # Compact all physical layers together, after every shared reader. - move_sliding_window( - self.sliding_kv_buffer, - infer_state.b_req_idx, - infer_state.b_seq_len, - infer_state.b_ready_cache_len, - infer_state.b_q_start_loc, - self.sliding_window, - infer_state.sliding_window_runtime_start, - compact=True, + state_cache_manager.get_state_cache(buffer_idx), ) - - def get_layer_kv(self, layer_index: int): - local_layer = self.sliding_config.get_sliding_layer_index(layer_index) - head_num = self.sliding_config.sliding_head_num - layer_buffer = self.sliding_kv_buffer[local_layer] - return layer_buffer[:, :head_num], layer_buffer[:, head_num:] diff --git a/lightllm/common/sliding_window_cache_manager/config.py b/lightllm/common/sliding_window_cache_manager/config.py index e9599efa7b..8bff8833ad 100644 --- a/lightllm/common/sliding_window_cache_manager/config.py +++ b/lightllm/common/sliding_window_cache_manager/config.py @@ -1,4 +1,5 @@ import dataclasses +import math from typing import Dict import torch @@ -26,12 +27,6 @@ def __post_init__(self): 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_sliding_layer_index(self, layer_index: int) -> int: - return self.sliding_layer_to_cache_index[layer_index] - - def get_full_layer_index(self, layer_index: int) -> int: - return self.full_layer_to_cache_index[layer_index] - def get_state_shape(self): return ( self.sliding_layer_num, @@ -41,13 +36,9 @@ def get_state_shape(self): ) def get_state_nbytes(self): - elements = 1 - for dim in self.get_state_shape(): - elements *= dim - return elements * self.dtype.itemsize + 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): - assert big_page_token_num > 0 and tp_world_size > 0 return ( big_page_token_num * self.full_layer_num @@ -59,7 +50,6 @@ def get_cpu_cache_full_att_bytes(self, big_page_token_num: int, tp_world_size: i ) def get_cpu_cache_state_bytes(self, tp_world_size: int): - assert tp_world_size > 0 return self.get_state_nbytes() * tp_world_size def get_cpu_cache_big_page_bytes(self, big_page_token_num: int, tp_world_size: int): diff --git a/lightllm/common/sliding_window_cache_manager/state_cache.py b/lightllm/common/sliding_window_cache_manager/state_cache.py index 1ad8c76cc2..385e5cce7c 100644 --- a/lightllm/common/sliding_window_cache_manager/state_cache.py +++ b/lightllm/common/sliding_window_cache_manager/state_cache.py @@ -7,14 +7,17 @@ class SlidingWindowStateCacheManager: - """GPU storage for immutable request-level sliding-window checkpoints.""" + """Pinned CPU checkpoints, size-first: [page, layer, window, 2 * heads, dim].""" def __init__(self, size: int, sliding_config: SlidingWindowCacheConfig, keep_num: int = 0): self.size = size self.keep_num = keep_num assert 0 <= keep_num <= size self.state_cache = torch.empty( - (size, *sliding_config.get_state_shape()), dtype=sliding_config.dtype, device="cuda" + (size, *sliding_config.get_state_shape()), + dtype=sliding_config.dtype, + device="cpu", + pin_memory=True, ) self.clear_to_init_state() diff --git a/lightllm/models/gemma4/infer_struct.py b/lightllm/models/gemma4/infer_struct.py index d18a470dc5..91bf0b76f5 100644 --- a/lightllm/models/gemma4/infer_struct.py +++ b/lightllm/models/gemma4/infer_struct.py @@ -1,5 +1,10 @@ import torch from lightllm.common.basemodel import InferStateInfo +from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( + build_sliding_window_page_table, + commit_sliding_window_kv, + get_sliding_window_decode_indexes, +) from lightllm.models.gemma4.triton_kernel.build_b_image_token_end import build_b_image_token_end @@ -22,6 +27,8 @@ def __init__(self): # image token 可以看到自己当前这个token以及后面的 image token。 self.b_image_token_end = None self.sliding_window_mem_index = None + self.sliding_window_page_table = None + self.b_sliding_kv_start = None def init_some_extra_state(self, model): super().init_some_extra_state(model) @@ -39,13 +46,38 @@ 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() + self.sliding_window_mem_index = self.mem_manager.alloc_sliding_prefill(self.input_ids.shape[0]) + self.sliding_window_page_table, self.b_sliding_kv_start = build_sliding_window_page_table( + self.b_req_idx, + self.b_seq_len, + self.b_ready_cache_len, + self.b_q_start_loc, + self.sliding_window_mem_index, + self.req_manager.sliding_window, + self.max_q_seq_len, + ) else: - self.b_q_start_loc = self.b1_cu_q_seq_len[:-1] - self.req_manager.prepare_sliding_window(self) + self.sliding_window_mem_index = get_sliding_window_decode_indexes( + self.b_req_idx, self.b_seq_len, self.req_manager.sliding_window + ) return + def finish_forward(self): + if not self.is_prefill: + return + # One commit after all KV-sharing readers, also outside CUDA graph's attention probes. + commit_sliding_window_kv( + self.mem_manager.sliding_kv_buffer, + self.sliding_window_mem_index, + self.b_req_idx, + self.b_seq_len, + self.b_ready_cache_len, + self.b_q_start_loc, + self.req_manager.sliding_window, + ) + self.mem_manager.free_sliding_prefill() + 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) diff --git a/lightllm/models/gemma4/kv_layout.py b/lightllm/models/gemma4/kv_layout.py index c637fea777..186686caf3 100644 --- a/lightllm/models/gemma4/kv_layout.py +++ b/lightllm/models/gemma4/kv_layout.py @@ -13,7 +13,7 @@ def get_kv_cache_layout(config): cache_map = layer_maps[layer_type] if layer_index < cutoff: last_owner[layer_type] = layer_index - cache_map[layer_index] = len(set(cache_map.values())) + cache_map[layer_index] = len(cache_map) else: cache_map[layer_index] = cache_map[last_owner[layer_type]] owner = last_owner[layer_type] @@ -25,7 +25,6 @@ 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 tp_world_size > 0 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) diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 1e482e5665..d86cb79a56 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -4,6 +4,7 @@ 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 @@ -18,7 +19,7 @@ class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): """ Gemma-4 decoder block. Full-attention KV stays token granular, while - sliding attention reads one runtime KV pool with an adjustable window. + sliding attention reads fixed request rings and new-token KV from one pool. """ def __init__(self, layer_num, network_config): @@ -50,16 +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_ - # 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). @@ -69,13 +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_owners = get_kv_cache_layout(network_config) + 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_share_target_layer_ = kv_owner if self.is_kv_shared_ else None - self.finish_sliding_prefill_ = self.is_sliding and not any( - kind == "sliding_attention" for kind in network_config["layer_types"][layer_num + 1 :] - ) + 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 @@ -145,35 +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: - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv - - layer_idx = infer_state.req_manager.sliding_config.get_sliding_layer_index(self.layer_num_) + # 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.req_manager.sliding_kv_buffer[layer_idx], + infer_state.mem_manager.sliding_kv_buffer[self.sliding_cache_index_], ) return 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_ - if self.is_sliding: - return infer_state.req_manager.get_layer_kv(layer_idx) - return infer_state.mem_manager.get_att_input_params(layer_index=layer_idx) - def _context_attention_kernel( self, q: torch.Tensor, @@ -182,13 +164,19 @@ 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: + if not self.is_kv_shared_: + # Use the callback's live indices on prefill graph replay. History stays in the ring. + destindex_copy_kv( + kv, + infer_state.sliding_window_mem_index, + infer_state.mem_manager.sliding_kv_buffer[self.sliding_cache_index_], + ) # 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, @@ -199,21 +187,15 @@ def _context_attention_kernel( infer_state.b_seq_len, infer_state.b_ready_cache_len, infer_state.max_q_seq_len, - None, + infer_state.sliding_window_page_table, infer_state.b_image_token_end, - sliding_window=sw, - runtime_token_start=infer_state.sliding_window_runtime_start, + sliding_window=(self.sliding_window_ - 1, 0), + b_kv_start_pos=infer_state.b_sliding_kv_start, ) - # The final sliding reader compacts all physical windows together. - # Graph shape probing must not mutate state; replay uses fresh metadata. - if self.finish_sliding_prefill_ and not torch.cuda.is_current_stream_capturing(): - infer_state.req_manager.finish_prefill(infer_state) 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 + o_tensor = infer_state.prefill_att_state.prefill_att( + q=_q, k=_k, v=_v, att_control=AttControl(), alloc_func=self.alloc_tensor ) return o_tensor.view(q.shape) @@ -224,7 +206,7 @@ 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_) if self.is_sliding: from lightllm.models.gemma4.triton_kernel.sliding_window_decode import sliding_window_decode_attention @@ -240,8 +222,8 @@ def _token_attention_kernel( alloc_tensor_func=self.alloc_tensor, ) else: - o_tensor = infer_state.decode_att_state1.decode_att( - q=_q, k=_k, v=_v, att_control=self._att_control(), alloc_func=self.alloc_tensor + o_tensor = infer_state.decode_att_state.decode_att( + q=_q, k=_k, v=_v, att_control=AttControl(), alloc_func=self.alloc_tensor ) return o_tensor.view(q.shape) diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index eeacae6ce1..2a5ad196bf 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -69,25 +69,12 @@ def _init_config(self): 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']}" - ) - if kv_shared: - # Shared-KV microbatch overlap needs separate lifecycle validation. + 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 shared sliding-window KV does not support microbatch overlap yet" + ), "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" @@ -98,57 +85,28 @@ def _verify_params(self): assert not args.diverse_mode, "Gemma-4 sliding-window state does not support diverse mode yet" return - def _get_sliding_cache_config(self): - if hasattr(self, "sliding_cache_config"): - return self.sliding_cache_config - self.sliding_cache_config = build_sliding_cache_config(self.config, self.tp_world_size_, self.data_type) - return self.sliding_cache_config - def _init_req_manager(self): - args = get_env_start_args() - create_max_seq_len = max(int(self.batch_max_tokens or 0), int(self.max_seq_length or 0)) - max_prefill_token_num = max( - int(self.batch_max_tokens or 0), - int(args.prefill_cudagraph_max_handle_token or 0) if args.enable_prefill_cudagraph else 0, - ) - if args.enable_tpsp_mix_mode: - max_prefill_token_num = ( - (max(1, max_prefill_token_num) + self.tp_world_size_ - 1) // self.tp_world_size_ * self.tp_world_size_ - ) self.req_manager = ReqManagerForSlidingWindow( max_request_num=self.max_req_num, - max_sequence_length=create_max_seq_len, + max_sequence_length=max(self.batch_max_tokens, self.max_seq_length), mem_manager=None, - sliding_config=self._get_sliding_cache_config(), - max_prefill_token_num=max_prefill_token_num, - prefill_microbatch_num=2 if args.enable_prefill_microbatch_overlap else 1, + sliding_config=self.sliding_cache_config, ) def _init_mem_manager(self): self.mem_manager = HybridSlidingMemoryManager( size=self.max_total_token_num, - sliding_config=self._get_sliding_cache_config(), + sliding_config=self.sliding_cache_config, mem_fraction=self.mem_fraction, ) return 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. - # - # Sliding layers read their runtime KV pool through model-local kernels. - # The framework still requires primary attention states. + # Full attention uses the standard backend. Sliding attention uses + # Gemma's local kernels for the compact page table and image mask. 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) - def _init_custom(self): self._init_to_get_rotary_gemma4() if self.config.get("enable_moe_block", False): diff --git a/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py b/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py index bf159b42ba..c6beeab700 100644 --- a/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py +++ b/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py @@ -43,6 +43,7 @@ def _fwd_kernel( Req_to_tokens, B_req_idx, B_Image_Token_End, + B_KV_Start_Pos, stride_qbs, stride_qh, stride_qd, @@ -65,7 +66,7 @@ def _fwd_kernel( BLOCK_N: tl.constexpr, USE_SLIDING_WINDOW: tl.constexpr, SLIDING_WINDOW_LEFT: tl.constexpr, - RUNTIME_TOKEN_START: tl.constexpr, + COMPACT_KV: tl.constexpr, ): start_m = tl.program_id(0) cur_bh = tl.program_id(1) @@ -78,7 +79,8 @@ def _fwd_kernel( prompt_cache_len = tl.load(b_prompt_cache_len + cur_batch) total_len = tl.load(B_Seqlen + cur_batch) cur_batch_seq_len = total_len - prompt_cache_len # new tokens this step - cur_batch_req_idx = tl.load(B_req_idx + cur_batch) + cur_batch_req_idx = cur_batch if COMPACT_KV else tl.load(B_req_idx + cur_batch) + table_start = tl.load(B_KV_Start_Pos + cur_batch) if COMPACT_KV else 0 block_start_loc = BLOCK_M * start_m if block_start_loc >= cur_batch_seq_len: @@ -127,20 +129,11 @@ def _fwd_kernel( k_pos = kv_start_index + start_n + offs_n # [N] k_valid = k_pos < block_end_loc - if RUNTIME_TOKEN_START is not None: - kv_loc = ( - RUNTIME_TOKEN_START - + cur_batch_in_all_start_index.to(tl.int64) - + (cur_batch + 1) * (SLIDING_WINDOW_LEFT + 1) - + k_pos.to(tl.int64) - - prompt_cache_len.to(tl.int64) - ) - else: - kv_loc = tl.load( - Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * k_pos, - mask=k_valid, - other=0, - ).to(tl.int64) + kv_loc = tl.load( + Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * (k_pos - table_start), + mask=k_valid, + other=0, + ).to(tl.int64) k_ptr = K + kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd k = tl.load(k_ptr, mask=k_valid[None, :], other=0.0) qk = tl.dot(q, k) @@ -197,7 +190,7 @@ def context_attention_fwd_gemma4_mm( req_to_token_indexs, b_image_token_end, sliding_window=(-1, -1), - runtime_token_start=None, + b_kv_start_pos=None, ): """Prefill attention with image bidirectional masking on sliding layers. @@ -208,9 +201,9 @@ def context_attention_fwd_gemma4_mm( position (in the flattened new-token layout), value is the image span's end index (in absolute request position) if the token is inside an image span, else 0. - runtime_token_start: Start of the single prefill KV region, where each - request's W history tokens precede its current tokens. The token - index table is unused and may be None. + b_kv_start_pos: Absolute position of column zero in a batch-local compact KV table. + None means the table is indexed by request ID and covers the full history. Query and image positions + always remain absolute; only the page-table lookup is rebased. """ BLOCK_M = 128 if not is_tesla() else 64 Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] @@ -238,8 +231,8 @@ def context_attention_fwd_gemma4_mm( assert int(sliding_window[1]) == 0, "sliding_window right must be 0" sliding_window_left = int(sliding_window[0]) - if runtime_token_start is not None: - assert use_sliding_window and sliding_window_left >= 0, "runtime KV requires a finite sliding window" + if b_kv_start_pos is not None: + assert use_sliding_window and sliding_window_left >= 0, "compact KV requires a finite sliding window" _fwd_kernel[grid]( q, @@ -252,6 +245,7 @@ def context_attention_fwd_gemma4_mm( req_to_token_indexs, b_req_idx, b_image_token_end, + b_kv_start_pos, q.stride(0), q.stride(1), q.stride(2), @@ -264,8 +258,8 @@ def context_attention_fwd_gemma4_mm( o.stride(0), o.stride(1), o.stride(2), - 0 if runtime_token_start is not None else req_to_token_indexs.stride(0), - 0 if runtime_token_start is not None else req_to_token_indexs.stride(1), + req_to_token_indexs.stride(0), + req_to_token_indexs.stride(1), kv_group_num=kv_group_num, b_prompt_cache_len=b_prompt_cache_len, H=head, @@ -274,7 +268,7 @@ def context_attention_fwd_gemma4_mm( BLOCK_N=BLOCK_N, USE_SLIDING_WINDOW=use_sliding_window, SLIDING_WINDOW_LEFT=sliding_window_left, - RUNTIME_TOKEN_START=runtime_token_start, + COMPACT_KV=b_kv_start_pos is not None, num_warps=num_warps, num_stages=num_stages, ) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index c7f535c2e9..ff0703facf 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -439,10 +439,10 @@ def copy_linear_att_state_to_cache_buffer(self, b_req_idx: torch.Tensor, reqs: L ) if req.tail_linear_att_small_page_buffer_id is not None: dst_buffer_idx = req.tail_linear_att_small_page_buffer_id - self.req_manager.save_small_page_state( + self.req_manager.save_state( req_idx=req.req_idx, buffer_idx=dst_buffer_idx, - small_page_buffers=self.radix_cache.linear_att_small_page_buffers, + state_cache_manager=self.radix_cache.linear_att_small_page_buffers, ) return diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 8c5e887b9d..fdd7fab72a 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -81,33 +81,25 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": else: mem_manager_class = select_mem_manager_class() - if mem_manager_class is HybridSlidingMemoryManager: - from lightllm.models.gemma4.kv_layout import build_sliding_cache_config - - model_config = get_config_json(args.model_dir) - text_config = model_config.get("text_config", model_config) - tp_world_size = args.tp // args.dp - sliding_config = build_sliding_cache_config(text_config, tp_world_size, get_llm_data_type()) - 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 - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=1, - layer_num=1, - num_heads=1, - head_dim=sliding_config.get_cpu_cache_big_page_bytes(big_page_token_num, tp_world_size), - data_type=torch.uint8, - scale_head_dim=0, - scale_data_type=get_llm_data_type(), - ) - elif mem_manager_class is Qwen3NextMemManager: - linear_config = LinearAttCacheConfig.load_from_args() + if mem_manager_class in (Qwen3NextMemManager, HybridSlidingMemoryManager): + if mem_manager_class is Qwen3NextMemManager: + page_bytes = LinearAttCacheConfig.load_from_args().get_cpu_cache_big_page_bytes() + else: + from lightllm.models.gemma4.kv_layout import build_sliding_cache_config + + model_config = get_config_json(args.model_dir) + text_config = model_config.get("text_config", model_config) + tp_world_size = args.tp // args.dp + sliding_config = build_sliding_cache_config(text_config, tp_world_size, get_llm_data_type()) + 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 + page_bytes = sliding_config.get_cpu_cache_big_page_bytes(big_page_token_num, tp_world_size) cpu_cache_meta = CpuKVCacheMeta( page_num=0, token_page_size=1, layer_num=1, num_heads=1, - head_dim=linear_config.get_cpu_cache_big_page_bytes(), + head_dim=page_bytes, data_type=torch.uint8, scale_head_dim=0, scale_data_type=get_llm_data_type(), diff --git a/test/kernel/test_sliding_window_cpu_cache_copy.py b/test/kernel/test_sliding_window_cpu_cache_copy.py index 7f16fc64ca..ec25399b03 100644 --- a/test/kernel/test_sliding_window_cpu_cache_copy.py +++ b/test/kernel/test_sliding_window_cpu_cache_copy.py @@ -6,6 +6,7 @@ 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.sliding_window_cache_manager import SlidingWindowCacheConfig @@ -84,7 +85,7 @@ def test_multi_page_round_trip_preserves_tp_slices_tail_and_ring(tp_world_size, 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) + 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() @@ -104,12 +105,11 @@ def test_multi_page_round_trip_preserves_tp_slices_tail_and_ring(tp_world_size, page_readies=page_readies.cuda(), big_page_buffer_ids=big_page_ids.cuda(), gpu_full_att_kv_state=full_cpu.cuda(), - gpu_sliding_state=window_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, - sliding_config=config, grid_num=3, ) torch.cuda.synchronize() @@ -124,7 +124,7 @@ def test_multi_page_round_trip_preserves_tp_slices_tail_and_ring(tp_world_size, 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_gpu = expected_full.cuda(), expected_window.cuda() + 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: @@ -137,34 +137,66 @@ def test_multi_page_round_trip_preserves_tp_slices_tail_and_ring(tp_world_size, page_indexes=load_pages, big_page_buffer_ids=load_slots.cuda(), gpu_full_att_kv_state=full_gpu, - gpu_sliding_state=window_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, - sliding_config=config, grid_num=3, ) torch.cuda.synchronize() _assert_same_bits(full_gpu, expected_full) - _assert_same_bits(window_gpu, expected_window) + _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(): - config = SlidingWindowCacheConfig({0: 0}, {1: 0}, 8, 1, 64, 1, 64, torch.bfloat16) 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, - gpu_sliding_state=None, + cpu_kv_sliding_state=None, cpu_cache_tensor=None, tp_rank=0, tp_world_size=1, big_page_token_num=16, - sliding_config=config, ) 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_prefill.py b/test/kernel/test_sliding_window_prefill.py index 3d72864502..d33853826b 100644 --- a/test/kernel/test_sliding_window_prefill.py +++ b/test/kernel/test_sliding_window_prefill.py @@ -2,6 +2,7 @@ import torch from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import context_attention_fwd_gemma4_mm +from lightllm.common.basemodel.triton_kernel.sliding_window_state import build_sliding_window_page_table pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -26,13 +27,13 @@ def _compare_runtime_and_paged(window, q_len, dtype, head_dim=64, image_span=Non offset += length runtime_start = req_slots * window - runtime = torch.full( - (runtime_start + query_num + len(req_ids) * window, 2 * kv_heads, head_dim), -3, device="cuda", dtype=dtype - ) - for batch, (req_id, history, length, start) in enumerate(zip(req_ids, histories, lengths, starts)): - positions = torch.arange(max(0, history - window), length, device="cuda") - current_start = runtime_start + start + (batch + 1) * window - runtime[current_start + positions - history] = reference[mapping[req_id, positions].long()] + runtime = torch.full((runtime_start + query_num, 2 * kv_heads, head_dim), -3, device="cuda", dtype=dtype) + for req_id, history, length, start in zip(req_ids, histories, lengths, starts): + positions = torch.arange(max(0, history - window), history, device="cuda") + runtime[req_id * window + positions % window] = reference[mapping[req_id, positions].long()] + runtime[runtime_start + start : runtime_start + start + length - history] = reference[ + mapping[req_id, history:length].long() + ] image_ends = torch.zeros(query_num, device="cuda", dtype=torch.int32) if image_span is not None: @@ -58,12 +59,22 @@ def _compare_runtime_and_paged(window, q_len, dtype, head_dim=64, image_span=Non req_to_token_indexs=mapping, **kwargs, ) + indexes = torch.arange(runtime_start, runtime_start + query_num, device="cuda", dtype=torch.int32) + page_table, kv_start = build_sliding_window_page_table( + kwargs["b_req_idx"], + kwargs["b_seq_len"], + kwargs["b_prompt_cache_len"], + kwargs["b_start_loc"], + indexes, + window, + max(q_lens), + ) context_attention_fwd_gemma4_mm( k=runtime[:, :kv_heads], v=runtime[:, kv_heads:], o=actual, - req_to_token_indexs=None, - runtime_token_start=runtime_start, + req_to_token_indexs=page_table, + b_kv_start_pos=kv_start, **kwargs, ) torch.testing.assert_close(actual, expected, atol=0, rtol=0) @@ -105,7 +116,8 @@ def test_runtime_prefill_cuda_graph_replay_reads_updated_request_metadata(window req_slots, head_dim, max_q_len = 6, 64, window + 33 query_num = 2 * (window + 64) + 97 runtime_start = req_slots * window - runtime = torch.randn((runtime_start + query_num + 2 * window, 4, head_dim), device="cuda", dtype=torch.bfloat16) + runtime = torch.randn((runtime_start + query_num, 4, head_dim), device="cuda", dtype=torch.bfloat16) + indexes = torch.arange(runtime_start, runtime_start + query_num, device="cuda", dtype=torch.int32) q = torch.randn((query_num, 4, head_dim), device="cuda", dtype=torch.bfloat16) out = torch.full_like(q, -11) int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) @@ -128,7 +140,15 @@ def test_runtime_prefill_cuda_graph_replay_reads_updated_request_metadata(window ) def forward(): - context_attention_fwd_gemma4_mm(o=out, req_to_token_indexs=None, runtime_token_start=runtime_start, **kwargs) + page_table, kv_start = build_sliding_window_page_table( + b_req, b_seq, b_history, b_start, indexes, window, max_q_len + ) + context_attention_fwd_gemma4_mm( + o=out, + req_to_token_indexs=page_table, + b_kv_start_pos=kv_start, + **kwargs, + ) forward() torch.cuda.synchronize() @@ -150,9 +170,10 @@ def forward(): # Materialize a table only for the independent old-path reference, after # changing every piece of GPU metadata used by the captured runtime kernel. mapping = torch.full((req_slots, max(lengths)), -1, device="cuda", dtype=torch.int32) - for batch, (req_id, history, length, start) in enumerate(zip(req_ids, histories, lengths, starts)): - positions = torch.arange(max(0, history - window), length, device="cuda", dtype=torch.int32) - mapping[req_id, positions.long()] = runtime_start + start + (batch + 1) * window + positions - history + for req_id, history, length, start in zip(req_ids, histories, lengths, starts): + positions = torch.arange(max(0, history - window), history, device="cuda", dtype=torch.int32) + mapping[req_id, positions.long()] = req_id * window + positions % window + mapping[req_id, history:length] = indexes[start : start + length - history] expected = torch.full_like(q, -11) context_attention_fwd_gemma4_mm(o=expected, req_to_token_indexs=mapping, **kwargs) # Include gaps to verify the captured grid respects the new query lengths. diff --git a/test/kernel/test_sliding_window_state.py b/test/kernel/test_sliding_window_state.py index a99fb4567b..81bc47d7a3 100644 --- a/test/kernel/test_sliding_window_state.py +++ b/test/kernel/test_sliding_window_state.py @@ -2,56 +2,51 @@ import torch from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( - get_sliding_window_mem_indexes, - move_sliding_window, + get_sliding_window_decode_indexes, + commit_sliding_window_kv, ) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("window,q_lengths", [(32, [1, 31, 65]), (512, [4096, 1, 513]), (1024, [8192, 7, 1023])]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -def test_prefill_moves_all_layers_and_preserves_canonical_snapshot(window, q_lengths, dtype): +@pytest.mark.parametrize( + "window,q_lengths", [(32, [1, 31, 65]), (128, [1, 127, 513]), (512, [4096, 1, 513]), (1024, [8192, 7, 1023])] +) +@pytest.mark.parametrize( + "dtype,payload_shape", + [(torch.bfloat16, (4, 32)), (torch.float16, (4, 32)), (torch.bfloat16, (512,)), (torch.uint8, (584,))], +) +def test_prefill_commits_all_layers_to_fixed_rings(window, q_lengths, dtype, payload_shape): torch.manual_seed(42) req_ids, histories = [3, 0, 5], [0, window - 1, 2 * window + 3] starts = [11, 18 + q_lengths[0], 31 + q_lengths[0] + q_lengths[1]] lengths = [history + q_len for history, q_len in zip(histories, q_lengths)] total_tokens = starts[-1] + q_lengths[-1] + 7 - runtime_start = 6 * window + 17 - pool = torch.full((3, runtime_start + total_tokens + 3 * window, 4, 32), -11, device="cuda", dtype=dtype) - references = [torch.randn((3, length, 4, 32), device="cuda", dtype=dtype) for length in lengths] + prefill_start = 6 * window + pool = torch.full((3, prefill_start + total_tokens, *payload_shape), 11, device="cuda", dtype=dtype) + state = pool[:, :prefill_start].unflatten(1, (6, window)) + references = [ + torch.randint(0, 256, (3, length, *payload_shape), device="cuda", dtype=dtype) + if dtype == torch.uint8 + else torch.randn((3, length, *payload_shape), device="cuda", dtype=dtype) + for length in lengths + ] int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - b_req, b_seq, b_history, b_q, b_start = map(int_tensor, [req_ids, lengths, histories, q_lengths, starts]) + b_req, b_seq, b_history, b_start = map(int_tensor, [req_ids, lengths, histories, starts]) for req, history, reference in zip(req_ids, histories, references): positions = torch.arange(max(0, history - window), history, device="cuda") - pool[:, req * window + positions % window] = reference[:, positions] + state[:, req, positions % window] = reference[:, positions] - expected = pool.clone() - for batch, (history, start, reference) in enumerate(zip(histories, starts, references)): - positions = torch.arange(max(0, history - window), history, device="cuda") - current_start = runtime_start + start + (batch + 1) * window - expected[:, current_start + positions - history] = reference[:, positions] - move_sliding_window(pool, b_req, b_seq, b_history, b_start, window, runtime_start) - torch.testing.assert_close(pool, expected, atol=0, rtol=0) - - indexes = get_sliding_window_mem_indexes( - b_req, b_seq, b_q, b_start, window, runtime_start, total_tokens, max(q_lengths), is_prefill=True - ) - for batch, (start, q_len, history, reference) in enumerate(zip(starts, q_lengths, histories, references)): - expected_indexes = runtime_start + start + (batch + 1) * window + torch.arange(q_len, device="cuda") - torch.testing.assert_close(indexes[start : start + q_len].long(), expected_indexes, atol=0, rtol=0) + indexes = torch.arange(prefill_start, prefill_start + total_tokens, device="cuda", dtype=torch.int32) + for start, q_len, history, reference in zip(starts, q_lengths, histories, references): pool[:, indexes[start : start + q_len].long()] = reference[:, history:] - expected = pool.clone() for req, length, reference in zip(req_ids, lengths, references): positions = torch.arange(max(0, length - window), length, device="cuda") expected[:, req * window + positions % window] = reference[:, positions] - move_sliding_window(pool, b_req, b_seq, b_history, b_start, window, runtime_start, compact=True) + commit_sliding_window_kv(pool, indexes, b_req, b_seq, b_history, b_start, window) + # Check every layer, untouched requests, query gaps and the temporary region. torch.testing.assert_close(pool, expected, atol=0, rtol=0) - # The page format remains a raw W-slot ring, independent of the active area. - snapshot = pool[:, 5 * window : 6 * window].clone() - pool[:, runtime_start:].zero_() - torch.testing.assert_close(snapshot, expected[:, 5 * window : 6 * window], atol=0, rtol=0) @pytest.mark.parametrize("window", [32, 512, 1024]) @@ -60,7 +55,7 @@ def test_decode_indexes_cuda_graph_replay(window): b_req, b_seq = int_tensor([2, 0, 5]), int_tensor([1, window, 2 * window + 1]) def forward(): - return get_sliding_window_mem_indexes(b_req, b_seq, None, None, window, 6 * window, 3, 1, False) + return get_sliding_window_decode_indexes(b_req, b_seq, window) torch.testing.assert_close(forward(), int_tensor([2 * window, window - 1, 5 * window]), atol=0, rtol=0) torch.cuda.synchronize() From 8300ce62e9a4cd01fd0eb9d1ac3290f361de73a1 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:31:49 +0000 Subject: [PATCH 12/14] refactor: simplify sliding-window cache with a private KV pool --- lightllm/common/basemodel/basemodel.py | 6 +- .../triton_kernel/sliding_window_state.py | 165 +++++------------ .../common/kv_cache_mem_manager/allocator.py | 28 +-- .../hybrid_sliding_mem_manager.py | 45 +---- .../kv_cache_mem_manager/mem_manager.py | 6 +- lightllm/common/req_manager/__init__.py | 2 +- lightllm/common/req_manager/base.py | 5 +- lightllm/common/req_manager/hybrid_att.py | 49 ----- lightllm/common/req_manager/hybrid_base.py | 66 +++++++ lightllm/common/req_manager/linear_att.py | 33 ++-- lightllm/common/req_manager/sliding_window.py | 100 +++++++++-- .../state_cache.py | 6 +- lightllm/models/gemma4/infer_struct.py | 78 +++++--- .../layer_infer/transformer_layer_infer.py | 27 +-- lightllm/models/gemma4/model.py | 35 +++- .../context_attention_fwd_gemma4_mm.py | 25 +-- .../triton_kernel/sliding_window_decode.py | 165 ----------------- .../server/router/model_infer/infer_batch.py | 10 +- .../model_infer/mode_backend/base_backend.py | 2 +- test/kernel/test_sliding_window_decode.py | 168 ++++-------------- test/kernel/test_sliding_window_prefill.py | 105 ++--------- test/kernel/test_sliding_window_state.py | 85 +++------ 22 files changed, 421 insertions(+), 790 deletions(-) delete mode 100644 lightllm/common/req_manager/hybrid_att.py create mode 100644 lightllm/common/req_manager/hybrid_base.py delete mode 100644 lightllm/models/gemma4/triton_kernel/sliding_window_decode.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index a3ae29e144..f9c674cf5c 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -939,6 +939,8 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 infer_state1=infer_state1, ) + 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) else: model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) @@ -963,11 +965,11 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 infer_state1.init_att_state() model_output0, model_output1 = self._overlap_tpsp_token_forward(infer_state0, infer_state1=infer_state1) + 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() - 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) return model_output0, model_output1 @final diff --git a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py index 9871bab1c5..3b8ac3c0e1 100644 --- a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py @@ -1,130 +1,57 @@ -import math import torch import triton import triton.language as tl @triton.jit -def _build_sliding_window_page_table( - PageTable, - BKVStart, - BReqIdx, - BSeqLen, - BReadyCacheLen, - BQStartLoc, - MemIndexes, - table_width, - WINDOW: tl.constexpr, - BLOCK: tl.constexpr, -): - batch = tl.program_id(0) - offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - req_idx = tl.load(BReqIdx + batch) - history = tl.load(BReadyCacheLen + batch) - seq_len = tl.load(BSeqLen + batch) - q_start = tl.load(BQStartLoc + batch) - kv_start = tl.maximum(history - WINDOW, 0) - positions = kv_start + offsets - is_new = positions >= history - new_index = tl.load(MemIndexes + q_start + positions - history, mask=is_new & (positions < seq_len), other=0) - index = tl.where(is_new, new_index, req_idx * WINDOW + positions % WINDOW) - tl.store( - PageTable + batch * table_width + offsets, tl.where(positions < seq_len, index, -1), mask=offsets < table_width - ) - if tl.program_id(1) == 0: - tl.store(BKVStart + batch, kv_start) - - -@torch.no_grad() -def build_sliding_window_page_table( - b_req_idx, b_seq_len, b_ready_cache_len, b_q_start_loc, mem_indexes, window, max_q_seq_len -): - """Batch-local, chronological table over fixed ring history and this prefill's new KV. - - Column zero represents b_kv_start_pos, not absolute token zero. No KV is moved. - """ - page_table = torch.empty((b_req_idx.numel(), window + max_q_seq_len), dtype=torch.int32, device=b_req_idx.device) - b_kv_start_pos = torch.empty_like(b_req_idx) - _build_sliding_window_page_table[(b_req_idx.numel(), triton.cdiv(page_table.shape[1], 256))]( - page_table, - b_kv_start_pos, - b_req_idx, - b_seq_len, - b_ready_cache_len, - b_q_start_loc, - mem_indexes, - page_table.shape[1], - WINDOW=window, - BLOCK=256, - num_warps=4, - ) - return page_table, b_kv_start_pos - - -@triton.jit -def _commit_sliding_window_kv( +def _copy_sliding_window_checkpoint( Pool, - MemIndexes, - BReqIdx, - BSeqLen, - BReadyCacheLen, - BQStartLoc, - stride_layer, - stride_token, + ReqToTokens, + seq_len, + Checkpoint, + req_idx, + table_stride, + pool_stride_l, + pool_stride_t, WINDOW: tl.constexpr, - KV_DIM: tl.constexpr, + TOKEN_BYTES: tl.constexpr, + TOTAL_BYTES: tl.constexpr, + RESTORE: tl.constexpr, BLOCK: tl.constexpr, ): - batch, layer = tl.program_id(0), tl.program_id(1) - offsets = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) - req_idx = tl.load(BReqIdx + batch).to(tl.int64) - history = tl.load(BReadyCacheLen + batch) - seq_len = tl.load(BSeqLen + batch) - q_start = tl.load(BQStartLoc + batch) - positions = seq_len - WINDOW + offsets // KV_DIM - # Retained old history is already in place. Copy only this chunk's newest W KV. - mask = (offsets < WINDOW * KV_DIM) & (positions >= history) - src_index = tl.load(MemIndexes + q_start + positions - history, mask=mask, other=0).to(tl.int64) - dst_index = req_idx * WINDOW + positions % WINDOW - layer_ptr = Pool + layer.to(tl.int64) * stride_layer - values = tl.load(layer_ptr + src_index * stride_token + offsets % KV_DIM, mask=mask, other=0) - tl.store(layer_ptr + dst_index * stride_token + offsets % KV_DIM, values, mask=mask) - - -@torch.no_grad() -def commit_sliding_window_kv(pool, mem_indexes, b_req_idx, b_seq_len, b_ready_cache_len, b_q_start_loc, window): - """Commit [layer, slot, ...payload] tails after every prefill reader has finished. - - The contiguous per-token payload can be KV heads, an MLA vector or packed bytes. - """ - kv_dim = math.prod(pool.shape[2:]) - _commit_sliding_window_kv[(b_req_idx.numel(), pool.shape[0], triton.cdiv(window * kv_dim, 1024))]( - pool, - mem_indexes, - b_req_idx, - b_seq_len, - b_ready_cache_len, - b_q_start_loc, - pool.stride(0), - pool.stride(1), - WINDOW=window, - KV_DIM=kv_dim, - BLOCK=1024, - num_warps=4, + for block in range(tl.program_id(0), tl.cdiv(TOTAL_BYTES, BLOCK), tl.num_programs(0)): + offsets = tl.cast(block, tl.int64) * BLOCK + tl.arange(0, BLOCK) + layer = offsets // (WINDOW * TOKEN_BYTES) + ring_pos = offsets // TOKEN_BYTES % WINDOW + # Keep the existing CPU checkpoint order: absolute token position % W. + position = seq_len - 1 - (seq_len - 1 - ring_pos + WINDOW) % WINDOW + valid = (offsets < TOTAL_BYTES) & (position >= 0) + slot = tl.load(ReqToTokens + tl.cast(req_idx, tl.int64) * table_stride + position, valid, other=0).to(tl.int64) + pool_ptr = Pool + layer * pool_stride_l + slot * pool_stride_t + offsets % TOKEN_BYTES + if RESTORE: + value = tl.load(Checkpoint + offsets, valid, other=0) + tl.store(pool_ptr, value, valid) + else: + value = tl.load(pool_ptr, valid, other=0) + tl.store(Checkpoint + offsets, value, offsets < TOTAL_BYTES) + + +def copy_sliding_window_checkpoint(pool, req_to_tokens, seq_len: int, req_idx: int, checkpoint, restore=False): + """Gather/scatter GPU slots to a pinned CPU [layer, W, ...] checkpoint, byte-exact.""" + pool_bytes = pool.view(torch.uint8) + checkpoint_bytes = checkpoint.view(torch.uint8) + _copy_sliding_window_checkpoint[(16,)]( + pool_bytes, + req_to_tokens, + seq_len, + checkpoint_bytes, + req_idx, + req_to_tokens.stride(0), + pool_bytes.stride(0), + pool_bytes.stride(1), + WINDOW=checkpoint.shape[1], + TOKEN_BYTES=pool_bytes.stride(1), + TOTAL_BYTES=checkpoint_bytes.numel(), + RESTORE=restore, + BLOCK=4096, ) - - -@triton.jit -def _get_sliding_window_decode_indexes(Out, BReqIdx, BSeqLen, WINDOW: tl.constexpr): - batch = tl.program_id(0) - req_idx = tl.load(BReqIdx + batch) - seq_len = tl.load(BSeqLen + batch) - tl.store(Out + batch, req_idx * WINDOW + (seq_len - 1) % WINDOW) - - -@torch.no_grad() -def get_sliding_window_decode_indexes(b_req_idx, b_seq_len, window): - """Single-token decode writes directly into each request's fixed ring.""" - indexes = torch.empty_like(b_req_idx) - _get_sliding_window_decode_indexes[(b_req_idx.numel(),)](indexes, b_req_idx, b_seq_len, WINDOW=window, num_warps=1) - return indexes 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 index df906dbae0..9067a93dac 100644 --- a/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -2,7 +2,6 @@ import triton from lightllm.common.sliding_window_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 .mem_manager import MemoryManager @@ -10,30 +9,15 @@ class HybridSlidingMemoryManager(MemoryManager): - """Token-granular full KV plus request-granular sliding-window KV.""" + """管理 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.sliding_prefill_start = (args.running_max_req_size + 1) * sliding_config.sliding_window - # Both microbatches share batch_max_tokens; allow TP/dummy padding for each. - self.max_sliding_prefill_tokens = args.batch_max_tokens + 2 * get_dp_world_size() - self._sliding_prefill_used = 0 - self._sliding_prefill_batches = 0 - # One layer-first pool: fixed request rings followed by this batch's new KV. - # Reserve it before profiling how much memory can be given to full attention. - self.sliding_kv_buffer = torch.zeros( - ( - sliding_config.sliding_layer_num, - self.sliding_prefill_start + self.max_sliding_prefill_tokens, - 2 * sliding_config.sliding_head_num, - sliding_config.sliding_head_dim, - ), - dtype=sliding_config.dtype, - device="cuda", - ) self.big_page_token_num = args.linear_att_page_block_num * args.linear_att_hash_page_size super().__init__( size=size, @@ -45,29 +29,6 @@ def __init__(self, size, sliding_config, always_copy=False, mem_fraction=0.9): mem_fraction=mem_fraction, ) - def alloc_sliding_prefill(self, token_num: int) -> torch.Tensor: - # Overlapping microbatches lease disjoint ranges of the same token budget. - assert ( - self._sliding_prefill_used + token_num <= self.max_sliding_prefill_tokens - ), "sliding prefill pool exhausted" - start = self.sliding_prefill_start + self._sliding_prefill_used - indexes = torch.arange(start, start + token_num, dtype=torch.int32, device="cuda") - self._sliding_prefill_used += token_num - self._sliding_prefill_batches += 1 - return indexes - - def free_sliding_prefill(self): - assert self._sliding_prefill_batches > 0 - self._sliding_prefill_batches -= 1 - if self._sliding_prefill_batches == 0: - self._sliding_prefill_used = 0 - - def free_all(self): - super().free_all() - # Also discard leases when warmup/error cleanup resets all requests. - self._sliding_prefill_used = 0 - self._sliding_prefill_batches = 0 - 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. diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 658d3e899c..59c5a5d131 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -26,7 +26,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 @@ -36,7 +38,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/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index 2d8f8113b4..75d7bccb8e 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -1,5 +1,5 @@ from .base import ReqManager -from .hybrid_att import HybridAttentionReqManager +from .hybrid_base import HybridAttentionReqManager from .linear_att import ReqManagerForMamba from .req_sampling_params import ReqSamplingParamsManager from .sliding_window import ReqManagerForSlidingWindow 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/hybrid_att.py b/lightllm/common/req_manager/hybrid_att.py deleted file mode 100644 index 3fd36e5683..0000000000 --- a/lightllm/common/req_manager/hybrid_att.py +++ /dev/null @@ -1,49 +0,0 @@ -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List - -import torch - -from .base import ReqManager - - -if TYPE_CHECKING: - from lightllm.server.router.model_infer.infer_batch import InferReq - - -class HybridAttentionReqManager(ReqManager, ABC): - """混合 attention 的请求运行态与大小页 checkpoint 管理接口。 - - 大小页沿同一虚拟 token 索引空间匹配前缀，full attention KV 保持 token 粒度存储。 - linear/sliding-window 状态在大页边界及请求可缓存尾部的小页边界保存 checkpoint， - 缓存命中后，再将相应 checkpoint 恢复到请求运行态。 - - 公共缓存流程负责大小页分配、边界、匹配与淘汰；各实现负责状态存储和保存/恢复。 - """ - - @abstractmethod - def create_state_cache_manager(self, size: int): - """Return checkpoint storage used by request-state page boundaries.""" - - @abstractmethod - def init_hybrid_attention_state(self, req: "InferReq"): - """Initialize request runtime state when no prefix cache is restored.""" - - def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - self.restore_state(req, self.mem_manager.linear_att_big_page_buffers, big_page_buffer_idx) - - def restore_small_page_state(self, req: "InferReq", small_page_buffers): - self.restore_state(req, small_page_buffers, req.shared_kv_node.small_page_buffer_idx) - - @abstractmethod - def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): - """Restore the same request-state payload from either checkpoint pool.""" - - def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - """Default checkpoint copies; models may override with a batched kernel.""" - for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): - if buffer_idx != -1: - self.save_state(req_idx, buffer_idx, self.mem_manager.linear_att_big_page_buffers) - - @abstractmethod - def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager): - """Save a request's payload into either checkpoint pool.""" diff --git a/lightllm/common/req_manager/hybrid_base.py b/lightllm/common/req_manager/hybrid_base.py new file mode 100644 index 0000000000..512a40ead2 --- /dev/null +++ b/lightllm/common/req_manager/hybrid_base.py @@ -0,0 +1,66 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, List + +import torch + +from .base import ReqManager + + +if TYPE_CHECKING: + from lightllm.server.router.model_infer.infer_batch import InferReq + + +class HybridAttentionReqManager(ReqManager, ABC): + """混合 attention 的请求运行态与大小页 checkpoint 管理接口。 + + 大小页沿同一虚拟 token 索引空间匹配前缀，full attention KV 保持 token 粒度存储。 + linear/sliding-window 状态在大页边界及请求可缓存尾部的小页边界保存 checkpoint， + 缓存命中后，再将相应 checkpoint 恢复到请求运行态。 + + 请求的 GPU 计算状态由具体实现管理；small_page_buffers 持有 CPU 小页快照； + big_page_buffers 引用 mem_manager 的 CPU 大页快照，与 full KV 容量一起创建、调整。 + 大小页不包含额外的 GPU 运行态，命中后恢复到请求状态，不重新分配 checkpoint 池。 + 公共缓存流程负责 checkpoint 槽位分配、边界、匹配与淘汰；本接口负责运行态与 checkpoint 保存/恢复。 + """ + + def __init__(self, max_request_num, max_sequence_length, mem_manager): + super().__init__(max_request_num, max_sequence_length, mem_manager) + self.small_page_buffers = None + + @property + def big_page_buffers(self): + return self.mem_manager.linear_att_big_page_buffers + + @abstractmethod + def create_small_page_cache_manager(self, size: int): + """创建并持有 CPU 小页池，返回同一池供 radix 使用；size 是槽位数，不是 token 数。""" + + @abstractmethod + def init_hybrid_attention_state(self, req: "InferReq"): + """无前缀缓存命中时，初始化已分配请求槽位的 GPU 运行态。""" + + def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): + """将指定大页槽位的 CPU checkpoint 恢复到请求 GPU 运行态。""" + self.restore_state(req, self.big_page_buffers, big_page_buffer_idx) + + def restore_small_page_state(self, req: "InferReq"): + """将 req.shared_kv_node 对应的小页 checkpoint 恢复到请求 GPU 运行态。""" + self.restore_state(req, self.small_page_buffers, req.shared_kv_node.small_page_buffer_idx) + + @abstractmethod + def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): + """CPU checkpoint → 请求 GPU 运行态；大小页共用，不负责前缀匹配或 full KV 索引恢复。""" + + def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): + """批量保存请求 GPU 运行态到已分配的大页槽位，buffer_indexes 中的 -1 表示跳过。 + + b_req_idx 与 req_indexes 分别为同一批请求的 GPU 索引张量和 CPU 索引列表。 + 默认逐请求保存，模型可覆盖为批量拷贝算子。 + """ + for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): + if buffer_idx != -1: + self.save_state(req_idx, buffer_idx, self.big_page_buffers) + + @abstractmethod + def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager): + """请求 GPU 运行态 → 指定 CPU checkpoint 槽位；大小页共用，调用方负责分配槽位。""" diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index 50f09fa44e..e5d4e85eeb 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -7,7 +7,7 @@ from lightllm.common.linear_att_cache_manager.linear_att_buffer_manager import LinearAttCacheManager from lightllm.utils.envs_utils import get_env_start_args -from .hybrid_att import HybridAttentionReqManager +from .hybrid_base import HybridAttentionReqManager if TYPE_CHECKING: @@ -50,14 +50,26 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_con ) return - def create_state_cache_manager(self, size: int): - return LinearAttCacheManager(size=size, linear_config=self.linear_config) + def init_hybrid_attention_state(self, req: "InferReq"): + conv_index = req.req_idx + ssm_start = req.req_idx * (self.mtp_step + 1) + self.req_to_conv_state.buffer[:, conv_index, ...].fill_(0) + # #17: zero the FULL (mtp_step + 1)-row SSM block, not just canonical row +0, so a future + # first-step verify reading offset>0 after fresh init never hits a never-written row (NaN). + self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + (self.mtp_step + 1), ...].fill_(0) + if self.req_to_mtp_state_index is not None: + self.req_to_mtp_state_index[req.req_idx] = 0 + return + + def create_small_page_cache_manager(self, size: int): + self.small_page_buffers = LinearAttCacheManager(size=size, linear_config=self.linear_config) + return self.small_page_buffers def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) - state_cache_manager = self.mem_manager.linear_att_big_page_buffers + state_cache_manager = self.big_page_buffers copy_linear_att_state_to_kv_buffer( b_req_idx=b_req_idx, big_page_buffer_ids=buffer_indexes, @@ -70,7 +82,7 @@ def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], return def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearAttCacheManager): - # Preserve main's small-page copies, including the MTP conv-state crop. + # checkpoint 只保存标准 conv 窗口和请求的基准 SSM 状态，不包含 MTP 扩展运行态。 conv_cache_width = self.linear_config.get_conv_state_shape()[-1] gpu_conv_state = self.req_to_conv_state.buffer[:, req_idx, ..., :conv_cache_width] gpu_ssm_state = self.req_to_ssm_state.buffer[:, req_idx * (self.mtp_step + 1), ...] @@ -78,17 +90,6 @@ def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearA dst_conv_state.copy_(gpu_conv_state, non_blocking=True) dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) - def init_hybrid_attention_state(self, req: "InferReq"): - conv_index = req.req_idx - ssm_start = req.req_idx * (self.mtp_step + 1) - self.req_to_conv_state.buffer[:, conv_index, ...].fill_(0) - # #17: zero the FULL (mtp_step + 1)-row SSM block, not just canonical row +0, so a future - # first-step verify reading offset>0 after fresh init never hits a never-written row (NaN). - self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + (self.mtp_step + 1), ...].fill_(0) - if self.req_to_mtp_state_index is not None: - self.req_to_mtp_state_index[req.req_idx] = 0 - return - def get_mamba_cache(self, layer_idx_in_all: int): assert ( 0 <= layer_idx_in_all < self.linear_config.all_layer_num diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index 8522da63fa..1ac36e6efb 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -1,9 +1,14 @@ from typing import TYPE_CHECKING, Optional -from lightllm.common.basemodel.triton_kernel.sliding_window_cpu_cache_copy import copy_sliding_window_state +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.sliding_window_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_att import HybridAttentionReqManager +from .hybrid_base import HybridAttentionReqManager if TYPE_CHECKING: @@ -13,7 +18,7 @@ class ReqManagerForSlidingWindow(HybridAttentionReqManager): - """按请求保存窗口运行态，并在大小页边界保存和恢复 checkpoint。""" + """管理请求的窗口索引与 checkpoint；私有 MemoryManager 负责 GPU KV 和物理槽位。""" def __init__( self, @@ -25,27 +30,90 @@ def __init__( 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() - @property - def req_to_sliding_window(self): - # A view, not a second allocation. req_idx owns the same W slots for its lifetime. - pool = self.mem_manager.sliding_kv_buffer - return pool[:, : self.mem_manager.sliding_prefill_start].unflatten(1, (-1, self.sliding_window)) - - def create_state_cache_manager(self, size: int): - return SlidingWindowStateCacheManager(size=size, sliding_config=self.sliding_config) + 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) + self._sliding_req_indexes = [torch.empty(0, dtype=torch.int32) for _ in range(self.max_request_num)] + self._sliding_seq_lens = [0] * self.max_request_num def init_hybrid_attention_state(self, req: "InferReq"): - self.req_to_sliding_window[:, req.req_idx].zero_() + # A cache miss owns no history slots; the forward allocates only its new tokens. + self._release_sliding_window(req.req_idx) + + def update_sliding_window(self, req_idx: int, seq_len: int, new_indexes: torch.Tensor): + """所有层读取后只保留最后 W 个槽位，无需移动 KV 或清空过期映射。""" + indexes = torch.cat((self._sliding_req_indexes[req_idx], new_indexes)) + expired = max(0, indexes.numel() - self.sliding_window) + if expired: + self.sliding_mem_manager.free(indexes[:expired]) + # Retain only the suffix, not the whole forward's pinned index buffer. + self._sliding_req_indexes[req_idx] = indexes[expired:].clone() + self._sliding_seq_lens[req_idx] = seq_len + + def _release_sliding_window(self, req_idx: int): + self.sliding_mem_manager.free(self._sliding_req_indexes[req_idx]) + self._sliding_req_indexes[req_idx] = torch.empty(0, dtype=torch.int32) + self._sliding_seq_lens[req_idx] = 0 + + def free_req(self, free_req_index: int): + self._release_sliding_window(free_req_index) + super().free_req(free_req_index) + + def free_all(self): + self.sliding_mem_manager.free_all() + self._sliding_req_indexes = [torch.empty(0, dtype=torch.int32) for _ in range(self.max_request_num)] + 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): - copy_sliding_window_state( + 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) + self._release_sliding_window(req.req_idx) + window_len = min(cache_len, self.sliding_window) + # Own the indices: MemoryManager.alloc() returns a reusable staging-buffer view. + indexes = torch.empty(window_len, dtype=torch.int32, device="cpu", pin_memory=True) + indexes.copy_(self.sliding_mem_manager.alloc(window_len)) + self._sliding_req_indexes[req.req_idx] = indexes + 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( + self.sliding_mem_manager.kv_buffer, + self.req_to_sliding_window, + cache_len, + req.req_idx, state_cache_manager.get_state_cache(buffer_idx), - self.req_to_sliding_window[:, req.req_idx], + restore=True, ) def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: SlidingWindowStateCacheManager): - copy_sliding_window_state( - self.req_to_sliding_window[:, req_idx], + copy_sliding_window_checkpoint( + self.sliding_mem_manager.kv_buffer, + self.req_to_sliding_window, + self._sliding_seq_lens[req_idx], + req_idx, state_cache_manager.get_state_cache(buffer_idx), ) diff --git a/lightllm/common/sliding_window_cache_manager/state_cache.py b/lightllm/common/sliding_window_cache_manager/state_cache.py index 385e5cce7c..2bdeafdb61 100644 --- a/lightllm/common/sliding_window_cache_manager/state_cache.py +++ b/lightllm/common/sliding_window_cache_manager/state_cache.py @@ -7,7 +7,11 @@ class SlidingWindowStateCacheManager: - """Pinned CPU checkpoints, size-first: [page, layer, window, 2 * heads, dim].""" + """大小页共用的 CPU pinned checkpoint 存储，两个池独立分配。 + + 布局为 size-first: [slot, layer, window, 2 * heads, dim]。 + 本类只管理状态存储与空闲槽位，不判断页面大小或缓存边界，也不持有 GPU 运行态。 + """ def __init__(self, size: int, sliding_config: SlidingWindowCacheConfig, keep_num: int = 0): self.size = size diff --git a/lightllm/models/gemma4/infer_struct.py b/lightllm/models/gemma4/infer_struct.py index 91bf0b76f5..8208fb368f 100644 --- a/lightllm/models/gemma4/infer_struct.py +++ b/lightllm/models/gemma4/infer_struct.py @@ -1,10 +1,10 @@ +import copy +from types import SimpleNamespace + import torch from lightllm.common.basemodel import InferStateInfo -from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( - build_sliding_window_page_table, - commit_sliding_window_kv, - get_sliding_window_decode_indexes, -) +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 @@ -27,8 +27,8 @@ def __init__(self): # image token 可以看到自己当前这个token以及后面的 image token。 self.b_image_token_end = None self.sliding_window_mem_index = None - self.sliding_window_page_table = None - self.b_sliding_kv_start = 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) @@ -47,36 +47,64 @@ def init_some_extra_state(self, model): ) if self.is_prefill: self._build_b_image_token_end() - self.sliding_window_mem_index = self.mem_manager.alloc_sliding_prefill(self.input_ids.shape[0]) - self.sliding_window_page_table, self.b_sliding_kv_start = build_sliding_window_page_table( + 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.append(sliding_mem_manager.alloc(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 allocator views once 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.req_manager.sliding_window, self.max_q_seq_len, ) else: - self.sliding_window_mem_index = get_sliding_window_decode_indexes( - self.b_req_idx, self.b_seq_len, self.req_manager.sliding_window + 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, ) + # A metadata view gives the unchanged attention backend its sliding table. + # Tensor metadata is shared with this state, including CUDA graph updates. + sliding_state = copy.copy(self) + sliding_state.req_manager = SimpleNamespace(req_to_token_indexs=self.req_manager.req_to_sliding_window) + self.decode_att_state1 = model.decode_att_backend.create_att_decode_state(infer_state=sliding_state) return def finish_forward(self): - if not self.is_prefill: - return - # One commit after all KV-sharing readers, also outside CUDA graph's attention probes. - commit_sliding_window_kv( - self.mem_manager.sliding_kv_buffer, - self.sliding_window_mem_index, - self.b_req_idx, - self.b_seq_len, - self.b_ready_cache_len, - self.b_q_start_loc, - self.req_manager.sliding_window, - ) - self.mem_manager.free_sliding_prefill() + # 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 diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index d86cb79a56..0280dafd80 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -19,7 +19,7 @@ class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): """ Gemma-4 decoder block. Full-attention KV stays token granular, while - sliding attention reads fixed request rings and new-token KV from one pool. + sliding attention indexes a bounded KV pool through its own token table. """ def __init__(self, layer_num, network_config): @@ -168,7 +168,7 @@ def _context_attention_kernel( _k, _v = infer_state.mem_manager.get_att_input_params(self.kv_cache_layer_index_) if self.is_sliding: if not self.is_kv_shared_: - # Use the callback's live indices on prefill graph replay. History stays in the ring. + # Use the callback's live indices on prefill graph replay. destindex_copy_kv( kv, infer_state.sliding_window_mem_index, @@ -187,10 +187,9 @@ def _context_attention_kernel( infer_state.b_seq_len, infer_state.b_ready_cache_len, infer_state.max_q_seq_len, - infer_state.sliding_window_page_table, + infer_state.req_manager.req_to_sliding_window, infer_state.b_image_token_end, sliding_window=(self.sliding_window_ - 1, 0), - b_kv_start_pos=infer_state.b_sliding_kv_start, ) return o_tensor.view(q.shape) @@ -209,22 +208,12 @@ def _token_attention_kernel( _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_) if self.is_sliding: - from lightllm.models.gemma4.triton_kernel.sliding_window_decode import sliding_window_decode_attention - - o_tensor = sliding_window_decode_attention( - q=_q, - k=_k, - v=_v, - b_req_idx=infer_state.b_req_idx, - b_seq_len=infer_state.b_seq_len, - sliding_window=self.sliding_window_, - out=out, - alloc_tensor_func=self.alloc_tensor, - ) + att_state = infer_state.decode_att_state1 + att_control = AttControl(use_sliding_window=True, sliding_window=(self.sliding_window_ - 1, 0)) else: - o_tensor = infer_state.decode_att_state.decode_att( - q=_q, k=_k, v=_v, att_control=AttControl(), alloc_func=self.alloc_tensor - ) + 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 2a5ad196bf..b6bf936266 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -99,11 +99,42 @@ def _init_mem_manager(self): 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): - # Full attention uses the standard backend. Sliding attention uses - # Gemma's local kernels for the compact page table and image mask. + # Both pools use main's token-indexed attention. Gemma's sliding prefill + # retains its image mask; full layers' head_dim=512 still requires Triton. self.prefill_att_backend = TritonAttBackend(model=self) self.decode_att_backend = TritonAttBackend(model=self) diff --git a/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py b/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py index c6beeab700..dee10e96d3 100644 --- a/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py +++ b/lightllm/models/gemma4/triton_kernel/context_attention_fwd_gemma4_mm.py @@ -43,7 +43,6 @@ def _fwd_kernel( Req_to_tokens, B_req_idx, B_Image_Token_End, - B_KV_Start_Pos, stride_qbs, stride_qh, stride_qd, @@ -66,7 +65,6 @@ def _fwd_kernel( BLOCK_N: tl.constexpr, USE_SLIDING_WINDOW: tl.constexpr, SLIDING_WINDOW_LEFT: tl.constexpr, - COMPACT_KV: tl.constexpr, ): start_m = tl.program_id(0) cur_bh = tl.program_id(1) @@ -79,8 +77,7 @@ def _fwd_kernel( prompt_cache_len = tl.load(b_prompt_cache_len + cur_batch) total_len = tl.load(B_Seqlen + cur_batch) cur_batch_seq_len = total_len - prompt_cache_len # new tokens this step - cur_batch_req_idx = cur_batch if COMPACT_KV else tl.load(B_req_idx + cur_batch) - table_start = tl.load(B_KV_Start_Pos + cur_batch) if COMPACT_KV else 0 + cur_batch_req_idx = tl.load(B_req_idx + cur_batch) block_start_loc = BLOCK_M * start_m if block_start_loc >= cur_batch_seq_len: @@ -130,12 +127,13 @@ def _fwd_kernel( k_valid = k_pos < block_end_loc kv_loc = tl.load( - Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * (k_pos - table_start), + Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * k_pos, mask=k_valid, other=0, ).to(tl.int64) - k_ptr = K + kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd - k = tl.load(k_ptr, mask=k_valid[None, :], other=0.0) + + off_k = kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd + k = tl.load(K + off_k, mask=k_valid[None, :], other=0.0) qk = tl.dot(q, k) if USE_SLIDING_WINDOW: @@ -160,8 +158,8 @@ def _fwd_kernel( l_i = l_i * alpha + l_ij acc = acc * alpha[:, None] - v_ptr = V + kv_loc[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd - v = tl.load(v_ptr, mask=k_valid[:, None], other=0.0) + off_v = kv_loc[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :] * stride_vd + v = tl.load(V + off_v, mask=k_valid[:, None], other=0.0) p = p.to(v.dtype) acc = tl.dot(p, v, acc) @@ -190,7 +188,6 @@ def context_attention_fwd_gemma4_mm( req_to_token_indexs, b_image_token_end, sliding_window=(-1, -1), - b_kv_start_pos=None, ): """Prefill attention with image bidirectional masking on sliding layers. @@ -201,9 +198,6 @@ def context_attention_fwd_gemma4_mm( position (in the flattened new-token layout), value is the image span's end index (in absolute request position) if the token is inside an image span, else 0. - b_kv_start_pos: Absolute position of column zero in a batch-local compact KV table. - None means the table is indexed by request ID and covers the full history. Query and image positions - always remain absolute; only the page-table lookup is rebased. """ BLOCK_M = 128 if not is_tesla() else 64 Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] @@ -231,9 +225,6 @@ def context_attention_fwd_gemma4_mm( assert int(sliding_window[1]) == 0, "sliding_window right must be 0" sliding_window_left = int(sliding_window[0]) - if b_kv_start_pos is not None: - assert use_sliding_window and sliding_window_left >= 0, "compact KV requires a finite sliding window" - _fwd_kernel[grid]( q, k, @@ -245,7 +236,6 @@ def context_attention_fwd_gemma4_mm( req_to_token_indexs, b_req_idx, b_image_token_end, - b_kv_start_pos, q.stride(0), q.stride(1), q.stride(2), @@ -268,7 +258,6 @@ def context_attention_fwd_gemma4_mm( BLOCK_N=BLOCK_N, USE_SLIDING_WINDOW=use_sliding_window, SLIDING_WINDOW_LEFT=sliding_window_left, - COMPACT_KV=b_kv_start_pos is not None, num_warps=num_warps, num_stages=num_stages, ) diff --git a/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py b/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py deleted file mode 100644 index f96365b580..0000000000 --- a/lightllm/models/gemma4/triton_kernel/sliding_window_decode.py +++ /dev/null @@ -1,165 +0,0 @@ -"""Gemma sliding decode over each request's canonical KV ring.""" - -import torch -import triton -import triton.language as tl - -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage2 import ( - flash_decode_stage2, -) - - -@triton.jit -def _sliding_window_decode_stage1( - Q, - K, - V, - BReqIdx, - BSeqLen, - MidO, - MidLogSumExp, - sm_scale, - stride_qb, - stride_qh, - stride_qd, - stride_kt, - stride_kh, - stride_kd, - stride_vt, - stride_vh, - stride_vd, - stride_ob, - stride_oh, - stride_os, - stride_od, - stride_lb, - stride_lh, - stride_ls, - gqa_group_size, - WINDOW: tl.constexpr, - Q_HEAD_NUM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, -): - batch_idx = tl.program_id(0) - kv_head = tl.program_id(1) - block_idx = tl.program_id(2) - grid_block_num = tl.num_programs(2) - - seq_len = tl.load(BSeqLen + batch_idx).to(tl.int64) - kv_start = tl.maximum(seq_len - WINDOW, 0) - window_len = seq_len - kv_start - total_blocks = tl.cdiv(window_len, BLOCK_SEQ) - if block_idx >= total_blocks: - return - - req_idx = tl.load(BReqIdx + batch_idx).to(tl.int64) - head_offsets = tl.arange(0, Q_HEAD_NUM) - q_heads = kv_head * gqa_group_size + head_offsets - q_heads = tl.where(head_offsets < gqa_group_size, q_heads, kv_head * gqa_group_size) - offs_d = tl.arange(0, BLOCK_DMODEL) - q = tl.load(Q + batch_idx * stride_qb + q_heads[:, None] * stride_qh + offs_d[None, :] * stride_qd) - - # Match the common GQA stage1's tiling and online-softmax arithmetic. - sum_exp = tl.zeros([Q_HEAD_NUM], dtype=tl.float32) - max_logic = tl.zeros([Q_HEAD_NUM], dtype=tl.float32) - float("inf") - acc = tl.zeros([Q_HEAD_NUM, BLOCK_DMODEL], dtype=tl.float32) - for block in range(block_idx, total_blocks, grid_block_num): - block_start = block * BLOCK_SEQ - block_end = tl.minimum(window_len, block_start + BLOCK_SEQ) - offs_n = block_start + tl.arange(0, BLOCK_N) - for tile in range(0, tl.cdiv(block_end - block_start, BLOCK_N)): - positions = tile * BLOCK_N + offs_n - mask = positions < block_end - token_pos = kv_start + positions - k_loc = req_idx * WINDOW + token_pos % WINDOW - k = tl.load( - K + k_loc[None, :] * stride_kt + kv_head * stride_kh + offs_d[:, None] * stride_kd, - mask=mask[None, :], - other=0.0, - ) - att_value = tl.dot(q, k.to(q.dtype)) - att_value *= sm_scale - att_value = tl.where(mask[None, :], att_value, float("-inf")) - v = tl.load( - V + k_loc[:, None] * stride_vt + kv_head * stride_vh + offs_d[None, :] * stride_vd, - mask=mask[:, None], - other=0.0, - ) - cur_max_logic = tl.max(att_value, axis=1) - new_max_logic = tl.maximum(cur_max_logic, max_logic) - exp_logic = tl.exp(att_value - new_max_logic[:, None]) - logic_scale = tl.exp(max_logic - new_max_logic) - acc *= logic_scale[:, None] - acc += tl.dot(exp_logic.to(v.dtype), v) - sum_exp = sum_exp * logic_scale + tl.sum(exp_logic, axis=1) - max_logic = new_max_logic - - out_offsets = ( - batch_idx * stride_ob + q_heads[:, None] * stride_oh + block_idx * stride_os + offs_d[None, :] * stride_od - ) - log_offsets = batch_idx * stride_lb + q_heads * stride_lh + block_idx * stride_ls - tl.store(MidO + out_offsets, acc / sum_exp[:, None], mask=(head_offsets < gqa_group_size)[:, None]) - tl.store(MidLogSumExp + log_offsets, max_logic + tl.log(sum_exp), mask=head_offsets < gqa_group_size) - - -@torch.no_grad() -def sliding_window_decode_attention( - q, - k, - v, - b_req_idx, - b_seq_len, - sliding_window: int, - out=None, - alloc_tensor_func=torch.empty, -): - """Decode one token per request after its current KV has been written to the ring.""" - batch_size, q_head_num, head_dim = q.shape - assert k.shape == v.shape and k.shape[-1] == head_dim - assert head_dim in {16, 32, 64, 128, 256, 512} - assert q_head_num % k.shape[1] == 0 - assert b_req_idx.shape == b_seq_len.shape == (batch_size,) - assert sliding_window > 0 - assert q.dtype == k.dtype == v.dtype - - # Keep the common GQA wrapper's launch and reduction schedule unchanged. - block_seq = 256 - block_num = 128 if batch_size <= 16 else (64 if batch_size <= 64 else 32) - mid_o = alloc_tensor_func([batch_size, q_head_num, block_num, head_dim], dtype=q.dtype, device=q.device) - mid_logsumexp = alloc_tensor_func([batch_size, q_head_num, block_num], dtype=torch.float32, device=q.device) - out = alloc_tensor_func(q.shape, dtype=q.dtype, device=q.device) if out is None else out - group_size = q_head_num // k.shape[1] - _sliding_window_decode_stage1[(batch_size, k.shape[1], block_num)]( - q, - k, - v, - b_req_idx, - b_seq_len, - mid_o, - mid_logsumexp, - 1.0 / (head_dim ** 0.5), - *q.stride(), - *k.stride(), - *v.stride(), - *mid_o.stride(), - *mid_logsumexp.stride(), - group_size, - WINDOW=sliding_window, - Q_HEAD_NUM=max(16, triton.next_power_of_2(group_size)), - BLOCK_SEQ=block_seq, - BLOCK_DMODEL=head_dim, - BLOCK_N=16, - num_warps=4, - num_stages=2, - ) - flash_decode_stage2( - mid_out=mid_o, - mid_out_logexpsum=mid_logsumexp, - B_Seqlen=b_seq_len, - out=out, - block_seq=block_seq, - sliding_window=(sliding_window - 1, 0), - ) - return out diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index ff0703facf..e68d2811b7 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -710,10 +710,7 @@ def _linear_match_radix_cache(self): self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 assert self.tail_linear_att_small_page_buffer_id is None # 恢复linear att 状态 - g_infer_context.req_manager.restore_small_page_state( - req=self, - small_page_buffers=g_infer_context.radix_cache.linear_att_small_page_buffers, - ) + g_infer_context.req_manager.restore_small_page_state(req=self) else: # 如果 大页本质是被启用的，则需要使用小页的匹配结果, 将小页的kv 复制到的新申请的kv位置，同时释放 # 对应的小页对应的节点，递归找到对应最近的大叶节点进行返回,然后赋值到req.shared_node 对象上 @@ -743,10 +740,7 @@ def _linear_match_radix_cache(self): ) self.shared_kv_node = share_node # 只是为了保证 restore_small_page_state 正确调用 - g_infer_context.req_manager.restore_small_page_state( - req=self, - small_page_buffers=g_infer_context.radix_cache.linear_att_small_page_buffers, - ) + g_infer_context.req_manager.restore_small_page_state(req=self) self.shared_kv_node = None big_page_shared_node = radix_cache.deref_to_first_big_page_node(node=share_node) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index d1f97c7b22..3aa966fa9c 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -154,7 +154,7 @@ def init_model(self, kvargs): self.is_hybrid_att_mixed_model = isinstance(self.model.req_manager, HybridAttentionReqManager) if self.is_hybrid_att_mixed_model: - self.linear_att_cache_manager = self.model.req_manager.create_state_cache_manager( + self.linear_att_cache_manager = self.model.req_manager.create_small_page_cache_manager( size=self.args.linear_att_cache_size ) else: diff --git a/test/kernel/test_sliding_window_decode.py b/test/kernel/test_sliding_window_decode.py index 774b6494f8..630c381cb8 100644 --- a/test/kernel/test_sliding_window_decode.py +++ b/test/kernel/test_sliding_window_decode.py @@ -6,155 +6,47 @@ from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding import ( gqa_token_decode_attention_flash_decoding, ) -from lightllm.models.gemma4.triton_kernel.sliding_window_decode import sliding_window_decode_attention pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def _int_tensor(values, dtype=torch.int32): - return torch.tensor(values, device="cuda", dtype=dtype) - - -def _table_reference(q, k, v, req_ids, seq_lengths, window): - # Only the visible suffix matters. Rebase very long sequences to avoid - # allocating a token table proportional to their virtual token positions. - indexes = torch.zeros((max(req_ids) + 1, window), device="cuda", dtype=torch.int64) - visible_lengths = [min(length, window) for length in seq_lengths] - for req_idx, length, visible_len in zip(req_ids, seq_lengths, visible_lengths): - positions = torch.arange(length - visible_len, length, device="cuda", dtype=torch.int64) - indexes[req_idx, :visible_len] = req_idx * window + positions % window - state = SimpleNamespace( - batch_size=len(req_ids), - b_req_idx=_int_tensor(req_ids), - b_seq_len=_int_tensor(visible_lengths), - max_kv_seq_len=window, - req_manager=SimpleNamespace(req_to_token_indexs=indexes), - ) - return gqa_token_decode_attention_flash_decoding( - q=q, - infer_state=state, - # The common kernel requires equal K/V strides and contiguous head_dim. - cache_k=k.contiguous(), - cache_v=v.contiguous(), - out=torch.empty_like(q), - sliding_window=(window - 1, 0), - ) - - -@pytest.fixture(autouse=True) -def _use_default_gqa_schedule(monkeypatch): - # Compare identical math and tiling rather than a machine-specific tune. - from lightllm.common.triton_utils import autotuner - - monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.CLOSE_AUTOTUNE) - - @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_formula_decode_matches_table_gqa_exactly(dtype, window, q_heads, kv_heads, head_dim): +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] - seq_lengths = [1, 2, window - 1, window, window + 1, 3 * window + 7] - runtime = torch.randn((9 * window, 2 * kv_heads, head_dim), device="cuda", dtype=dtype) + 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) - output = torch.empty_like(q) - actual = sliding_window_decode_attention( - q, - runtime[:, :kv_heads], - runtime[:, kv_heads:], - _int_tensor(req_ids), - _int_tensor(seq_lengths), - window, - out=output, - ) - expected = _table_reference(q, runtime[:, :kv_heads], runtime[:, kv_heads:], req_ids, seq_lengths, window) - assert actual is output - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - - -@pytest.mark.parametrize("batch_size", [1, 17, 65]) -def test_formula_decode_preserves_gqa_batch_schedule(batch_size): - window, head_dim = 512, 64 - req_ids = list(reversed(range(batch_size))) - seq_lengths = [2 * window + i + 1 for i in range(batch_size)] - runtime = torch.randn(((batch_size + 1) * window, 2, head_dim), device="cuda", dtype=torch.bfloat16) - q = torch.randn((batch_size, 4, head_dim), device="cuda", dtype=torch.bfloat16) - actual = sliding_window_decode_attention( - q, - runtime[:, :1], - runtime[:, 1:], - _int_tensor(req_ids), - _int_tensor(seq_lengths), - window, + 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 = _table_reference(q, runtime[:, :1], runtime[:, 1:], req_ids, seq_lengths, window) - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - - -def test_formula_decode_supports_int64_virtual_token_positions(): - window = 512 - req_ids, seq_lengths = [2, 0], [2 ** 31 + 17, 2 ** 32 + 31] - runtime = torch.randn((4 * window, 2, 64), device="cuda", dtype=torch.bfloat16) - q = torch.randn((2, 4, 64), device="cuda", dtype=torch.bfloat16) - actual = sliding_window_decode_attention( - q, - runtime[:, :1], - runtime[:, 1:], - _int_tensor(req_ids), - _int_tensor(seq_lengths, dtype=torch.int64), - window, + expected = gqa_token_decode_attention_flash_decoding( + q, state, full[:, :kv_heads], full[:, kv_heads:], out=torch.empty_like(q), sliding_window=(window - 1, 0) ) - expected = _table_reference(q, runtime[:, :1], runtime[:, 1:], req_ids, seq_lengths, window) - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -@pytest.mark.parametrize("window", [512, 1024]) -def test_formula_decode_cuda_graph_replay_with_changed_requests_and_padding(dtype, window): - runtime = torch.randn((8 * window, 4, 256), device="cuda", dtype=dtype) - q = torch.randn((4, 8, 256), device="cuda", dtype=dtype) - b_req = _int_tensor([4, 1, 7, 7]) - b_seq = _int_tensor([window + 7, 3, 2, 2]) - out = torch.empty_like(q) - - def forward(): - sliding_window_decode_attention(q, runtime[:, :2], runtime[:, 2:], b_req, b_seq, window, out=out) - - forward() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - forward() - req_ids, seq_lengths = [2, 5, 7, 7], [2 * window + 3, 1, 2, 2] - b_req.copy_(_int_tensor(req_ids)) - b_seq.copy_(_int_tensor(seq_lengths)) - q.mul_(0.5) - runtime.mul_(0.75) - graph.replay() - expected = _table_reference(q, runtime[:, :2], runtime[:, 2:], req_ids, seq_lengths, window) - # Padding may share the hold request ID; its outputs are intentionally discarded. - torch.testing.assert_close(out[:2], expected[:2], atol=0, rtol=0) - assert torch.isfinite(out).all() - - -def test_formula_decode_reads_independently_strided_kv_without_updating_runtime(): - window, kv_heads, head_dim = 512, 2, 64 - req_ids, seq_lengths = [2, 0], [window + 17, 1] - k = torch.randn((4 * window, kv_heads, head_dim * 2), device="cuda", dtype=torch.bfloat16)[..., ::2] - v = torch.randn((kv_heads, 4 * window, head_dim), device="cuda", dtype=torch.bfloat16).transpose(0, 1) - original_k, original_v = k.clone(), v.clone() - assert k.stride() != v.stride() - q = torch.randn((2, 8, head_dim), device="cuda", dtype=torch.bfloat16) - actual = sliding_window_decode_attention( - q, - k, - v, - _int_tensor(req_ids), - _int_tensor(seq_lengths), - window, + state.req_manager.req_to_token_indexs = window_table + actual = gqa_token_decode_attention_flash_decoding( + q, state, pool[:, :kv_heads], pool[:, kv_heads:], out=torch.empty_like(q), sliding_window=(window - 1, 0) ) - expected = _table_reference(q, k, v, req_ids, seq_lengths, window) torch.testing.assert_close(actual, expected, atol=0, rtol=0) - torch.testing.assert_close(k, original_k, atol=0, rtol=0) - torch.testing.assert_close(v, original_v, atol=0, rtol=0) diff --git a/test/kernel/test_sliding_window_prefill.py b/test/kernel/test_sliding_window_prefill.py index d33853826b..0389d3abf1 100644 --- a/test/kernel/test_sliding_window_prefill.py +++ b/test/kernel/test_sliding_window_prefill.py @@ -2,7 +2,6 @@ import torch from lightllm.models.gemma4.triton_kernel.context_attention_fwd_gemma4_mm import context_attention_fwd_gemma4_mm -from lightllm.common.basemodel.triton_kernel.sliding_window_state import build_sliding_window_page_table pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -26,14 +25,19 @@ def _compare_runtime_and_paged(window, q_len, dtype, head_dim=64, image_span=Non mapping[req_id, :length] = shuffled_indexes[offset : offset + length].to(torch.int32) offset += length - runtime_start = req_slots * window - runtime = torch.full((runtime_start + query_num, 2 * kv_heads, head_dim), -3, device="cuda", dtype=dtype) - for req_id, history, length, start in zip(req_ids, histories, lengths, starts): - positions = torch.arange(max(0, history - window), history, device="cuda") - runtime[req_id * window + positions % window] = reference[mapping[req_id, positions].long()] - runtime[runtime_start + start : runtime_start + start + length - history] = reference[ - mapping[req_id, history:length].long() - ] + 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: @@ -59,22 +63,11 @@ def _compare_runtime_and_paged(window, q_len, dtype, head_dim=64, image_span=Non req_to_token_indexs=mapping, **kwargs, ) - indexes = torch.arange(runtime_start, runtime_start + query_num, device="cuda", dtype=torch.int32) - page_table, kv_start = build_sliding_window_page_table( - kwargs["b_req_idx"], - kwargs["b_seq_len"], - kwargs["b_prompt_cache_len"], - kwargs["b_start_loc"], - indexes, - window, - max(q_lens), - ) context_attention_fwd_gemma4_mm( k=runtime[:, :kv_heads], v=runtime[:, kv_heads:], o=actual, - req_to_token_indexs=page_table, - b_kv_start_pos=kv_start, + req_to_token_indexs=runtime_mapping, **kwargs, ) torch.testing.assert_close(actual, expected, atol=0, rtol=0) @@ -108,73 +101,3 @@ def test_runtime_prefill_preserves_image_bidirectional_mask(window, dtype, image # 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) - - -@pytest.mark.parametrize("window", [512, 1024]) -def test_runtime_prefill_cuda_graph_replay_reads_updated_request_metadata(window): - torch.manual_seed(43) - req_slots, head_dim, max_q_len = 6, 64, window + 33 - query_num = 2 * (window + 64) + 97 - runtime_start = req_slots * window - runtime = torch.randn((runtime_start + query_num, 4, head_dim), device="cuda", dtype=torch.bfloat16) - indexes = torch.arange(runtime_start, runtime_start + query_num, device="cuda", dtype=torch.int32) - q = torch.randn((query_num, 4, head_dim), device="cuda", dtype=torch.bfloat16) - out = torch.full_like(q, -11) - int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - b_req = int_tensor([1, 4]) - b_history = int_tensor([window + 3, 0]) - b_seq = int_tensor([2 * window + 20, 31]) - b_start = int_tensor([13, max_q_len + 47]) - image_ends = torch.zeros(query_num, device="cuda", dtype=torch.int32) - kwargs = dict( - q=q, - k=runtime[:, :2], - v=runtime[:, 2:], - b_req_idx=b_req, - b_start_loc=b_start, - b_seq_len=b_seq, - b_prompt_cache_len=b_history, - max_input_len=max_q_len, - b_image_token_end=image_ends, - sliding_window=(window - 1, 0), - ) - - def forward(): - page_table, kv_start = build_sliding_window_page_table( - b_req, b_seq, b_history, b_start, indexes, window, max_q_len - ) - context_attention_fwd_gemma4_mm( - o=out, - req_to_token_indexs=page_table, - b_kv_start_pos=kv_start, - **kwargs, - ) - - forward() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - forward() - - req_ids, histories, q_lens, starts = [4, 2], [2 * window + 57, window - 5], [max_q_len, 17], [5, max_q_len + 59] - lengths = [history + count for history, count in zip(histories, q_lens)] - b_req.copy_(int_tensor(req_ids)) - b_history.copy_(int_tensor(histories)) - b_seq.copy_(int_tensor(lengths)) - b_start.copy_(int_tensor(starts)) - q.mul_(0.5) - runtime.mul_(0.75) - out.fill_(-11) - graph.replay() - - # Materialize a table only for the independent old-path reference, after - # changing every piece of GPU metadata used by the captured runtime kernel. - mapping = torch.full((req_slots, max(lengths)), -1, device="cuda", dtype=torch.int32) - for req_id, history, length, start in zip(req_ids, histories, lengths, starts): - positions = torch.arange(max(0, history - window), history, device="cuda", dtype=torch.int32) - mapping[req_id, positions.long()] = req_id * window + positions % window - mapping[req_id, history:length] = indexes[start : start + length - history] - expected = torch.full_like(q, -11) - context_attention_fwd_gemma4_mm(o=expected, req_to_token_indexs=mapping, **kwargs) - # Include gaps to verify the captured grid respects the new query lengths. - torch.testing.assert_close(out, expected, atol=0, rtol=0) diff --git a/test/kernel/test_sliding_window_state.py b/test/kernel/test_sliding_window_state.py index 81bc47d7a3..ff793750bf 100644 --- a/test/kernel/test_sliding_window_state.py +++ b/test/kernel/test_sliding_window_state.py @@ -1,68 +1,33 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.sliding_window_state import ( - get_sliding_window_decode_indexes, - commit_sliding_window_kv, -) +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( - "window,q_lengths", [(32, [1, 31, 65]), (128, [1, 127, 513]), (512, [4096, 1, 513]), (1024, [8192, 7, 1023])] -) -@pytest.mark.parametrize( - "dtype,payload_shape", - [(torch.bfloat16, (4, 32)), (torch.float16, (4, 32)), (torch.bfloat16, (512,)), (torch.uint8, (584,))], -) -def test_prefill_commits_all_layers_to_fixed_rings(window, q_lengths, dtype, payload_shape): - torch.manual_seed(42) - req_ids, histories = [3, 0, 5], [0, window - 1, 2 * window + 3] - starts = [11, 18 + q_lengths[0], 31 + q_lengths[0] + q_lengths[1]] - lengths = [history + q_len for history, q_len in zip(histories, q_lengths)] - total_tokens = starts[-1] + q_lengths[-1] + 7 - prefill_start = 6 * window - pool = torch.full((3, prefill_start + total_tokens, *payload_shape), 11, device="cuda", dtype=dtype) - state = pool[:, :prefill_start].unflatten(1, (6, window)) - references = [ - torch.randint(0, 256, (3, length, *payload_shape), device="cuda", dtype=dtype) - if dtype == torch.uint8 - else torch.randn((3, length, *payload_shape), device="cuda", dtype=dtype) - for length in lengths - ] - int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - b_req, b_seq, b_history, b_start = map(int_tensor, [req_ids, lengths, histories, starts]) - for req, history, reference in zip(req_ids, histories, references): - positions = torch.arange(max(0, history - window), history, device="cuda") - state[:, req, positions % window] = reference[:, positions] - - indexes = torch.arange(prefill_start, prefill_start + total_tokens, device="cuda", dtype=torch.int32) - for start, q_len, history, reference in zip(starts, q_lengths, histories, references): - pool[:, indexes[start : start + q_len].long()] = reference[:, history:] - expected = pool.clone() - for req, length, reference in zip(req_ids, lengths, references): - positions = torch.arange(max(0, length - window), length, device="cuda") - expected[:, req * window + positions % window] = reference[:, positions] - commit_sliding_window_kv(pool, indexes, b_req, b_seq, b_history, b_start, window) - # Check every layer, untouched requests, query gaps and the temporary region. - torch.testing.assert_close(pool, expected, atol=0, rtol=0) - - -@pytest.mark.parametrize("window", [32, 512, 1024]) -def test_decode_indexes_cuda_graph_replay(window): - int_tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - b_req, b_seq = int_tensor([2, 0, 5]), int_tensor([1, window, 2 * window + 1]) - - def forward(): - return get_sliding_window_decode_indexes(b_req, b_seq, window) - - torch.testing.assert_close(forward(), int_tensor([2 * window, window - 1, 5 * window]), atol=0, rtol=0) +@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() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - indexes = forward() - b_req.copy_(int_tensor([5, 1, 3])) - b_seq.copy_(int_tensor([window + 3, 7, 4 * window])) - graph.replay() - torch.testing.assert_close(indexes, int_tensor([5 * window + 2, window + 6, 4 * window - 1]), atol=0, rtol=0) + 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) From 5c6edc6119031c274e045a6006a9ec51da008f04 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:27:03 +0000 Subject: [PATCH 13/14] refactor: unify hybrid checkpoint cache managers Share checkpoint slot management across linear and sliding attention, and clarify sliding checkpoint copy names. --- .../linear_att_cpu_cache_copy.py | 2 +- .../triton_kernel/sliding_window_state.py | 98 +++++++++++-------- .../hybrid_sliding_mem_manager.py | 2 +- .../operator/linear_att.py | 2 +- .../qwen3next_mem_manager.py | 2 +- .../linear_att_cache_manager/__init__.py | 3 - .../linear_att_buffer_manager.py | 83 ---------------- lightllm/common/req_manager/linear_att.py | 4 +- lightllm/common/req_manager/sliding_window.py | 25 ++--- .../sliding_window_cache_manager/__init__.py | 5 - .../state_cache.py | 53 ---------- .../common/state_cache_manager/__init__.py | 16 +++ lightllm/common/state_cache_manager/base.py | 57 +++++++++++ .../layer_cache.py | 0 .../common/state_cache_manager/linear_att.py | 43 ++++++++ .../linear_att_config.py} | 0 .../state_cache_manager/sliding_window.py | 24 +++++ .../sliding_window_config.py} | 0 lightllm/models/gemma4/kv_layout.py | 2 +- lightllm/models/qwen3next/model.py | 2 +- .../dynamic_prompt/linear_att_radix_cache.py | 6 +- lightllm/utils/backend_validator.py | 2 +- lightllm/utils/kv_cache_utils.py | 2 +- test/cpu_cache_kernel/test_speed.py | 2 +- .../test_sliding_window_cpu_cache_copy.py | 2 +- .../basemodel/attention/linear/test_gdn.py | 2 +- 26 files changed, 226 insertions(+), 213 deletions(-) delete mode 100644 lightllm/common/linear_att_cache_manager/__init__.py delete mode 100644 lightllm/common/linear_att_cache_manager/linear_att_buffer_manager.py delete mode 100644 lightllm/common/sliding_window_cache_manager/__init__.py delete mode 100644 lightllm/common/sliding_window_cache_manager/state_cache.py create mode 100644 lightllm/common/state_cache_manager/__init__.py create mode 100644 lightllm/common/state_cache_manager/base.py rename lightllm/common/{linear_att_cache_manager => state_cache_manager}/layer_cache.py (100%) create mode 100644 lightllm/common/state_cache_manager/linear_att.py rename lightllm/common/{linear_att_cache_manager/config_objs.py => state_cache_manager/linear_att_config.py} (100%) create mode 100644 lightllm/common/state_cache_manager/sliding_window.py rename lightllm/common/{sliding_window_cache_manager/config.py => state_cache_manager/sliding_window_config.py} (100%) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py b/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py index ed0e742d73..b0c9a60649 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py @@ -1,7 +1,7 @@ import torch import triton import triton.language as tl -from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.state_cache_manager import LinearAttCacheConfig @triton.jit diff --git a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py index 3b8ac3c0e1..4305f96802 100644 --- a/lightllm/common/basemodel/triton_kernel/sliding_window_state.py +++ b/lightllm/common/basemodel/triton_kernel/sliding_window_state.py @@ -5,53 +5,71 @@ @triton.jit def _copy_sliding_window_checkpoint( - Pool, - ReqToTokens, - seq_len, - 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, - table_stride, - pool_stride_l, - pool_stride_t, - WINDOW: tl.constexpr, - TOKEN_BYTES: tl.constexpr, - TOTAL_BYTES: tl.constexpr, + 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: tl.constexpr, + BLOCK_BYTES: tl.constexpr, ): - for block in range(tl.program_id(0), tl.cdiv(TOTAL_BYTES, BLOCK), tl.num_programs(0)): - offsets = tl.cast(block, tl.int64) * BLOCK + tl.arange(0, BLOCK) - layer = offsets // (WINDOW * TOKEN_BYTES) - ring_pos = offsets // TOKEN_BYTES % WINDOW - # Keep the existing CPU checkpoint order: absolute token position % W. - position = seq_len - 1 - (seq_len - 1 - ring_pos + WINDOW) % WINDOW - valid = (offsets < TOTAL_BYTES) & (position >= 0) - slot = tl.load(ReqToTokens + tl.cast(req_idx, tl.int64) * table_stride + position, valid, other=0).to(tl.int64) - pool_ptr = Pool + layer * pool_stride_l + slot * pool_stride_t + offsets % TOKEN_BYTES + 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: - value = tl.load(Checkpoint + offsets, valid, other=0) - tl.store(pool_ptr, value, valid) + 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: - value = tl.load(pool_ptr, valid, other=0) - tl.store(Checkpoint + offsets, value, offsets < TOTAL_BYTES) + 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(pool, req_to_tokens, seq_len: int, req_idx: int, checkpoint, restore=False): - """Gather/scatter GPU slots to a pinned CPU [layer, W, ...] checkpoint, byte-exact.""" - pool_bytes = pool.view(torch.uint8) - checkpoint_bytes = checkpoint.view(torch.uint8) +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,)]( - pool_bytes, - req_to_tokens, - seq_len, - checkpoint_bytes, - req_idx, - req_to_tokens.stride(0), - pool_bytes.stride(0), - pool_bytes.stride(1), - WINDOW=checkpoint.shape[1], - TOKEN_BYTES=pool_bytes.stride(1), - TOTAL_BYTES=checkpoint_bytes.numel(), + 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=4096, + BLOCK_BYTES=4096, ) 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 index 9067a93dac..19d6fa3f8a 100644 --- a/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/hybrid_sliding_mem_manager.py @@ -1,7 +1,7 @@ import torch import triton -from lightllm.common.sliding_window_cache_manager import SlidingWindowStateCacheManager +from lightllm.common.state_cache_manager import SlidingWindowStateCacheManager from lightllm.utils.envs_utils import get_env_start_args from .mem_manager import MemoryManager diff --git a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py index 497c042314..64faa41ba0 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py +++ b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py @@ -6,7 +6,7 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size from lightllm.utils.log_utils import init_logger -from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.state_cache_manager import LinearAttCacheConfig if TYPE_CHECKING: from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py index 907cc494a6..67c16a694a 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -3,7 +3,7 @@ from lightllm.utils.log_utils import init_logger from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.linear_att_cache_manager import LinearAttCacheConfig, LinearAttCacheManager +from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager from .operator import LinearAttMemOperator from typing import Tuple, Any, List diff --git a/lightllm/common/linear_att_cache_manager/__init__.py b/lightllm/common/linear_att_cache_manager/__init__.py deleted file mode 100644 index ab3c8e2cd9..0000000000 --- a/lightllm/common/linear_att_cache_manager/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .linear_att_buffer_manager import LinearAttCacheManager -from .config_objs import LinearAttCacheConfig -from .layer_cache import LayerCache diff --git a/lightllm/common/linear_att_cache_manager/linear_att_buffer_manager.py b/lightllm/common/linear_att_cache_manager/linear_att_buffer_manager.py deleted file mode 100644 index 30dc4d937c..0000000000 --- a/lightllm/common/linear_att_cache_manager/linear_att_buffer_manager.py +++ /dev/null @@ -1,83 +0,0 @@ -import torch -import collections -from lightllm.utils.log_utils import init_logger -from .layer_cache import LayerCache -from typing import List, Optional, Tuple, Union -from .config_objs import LinearAttCacheConfig - -logger = init_logger(__name__) - - -class LinearAttCacheManager: - def __init__( - self, - size: int, - linear_config: LinearAttCacheConfig, - keep_num: int = 0, # 用于记录需要保留的缓存数量，用于支持含有 linear_att 的如qwen3.5 模型的cpu cache的碎页处理。 - ): - # init the mem state - self.size = size - self.linear_config = linear_config - self.keep_num = keep_num - assert 0 <= self.keep_num <= self.size, f"invalid keep_num {self.keep_num} for size {self.size}" - # init the layer cache - self.conv_state_cache = LayerCache( - size=self.size, - dtype=self.linear_config.conv_state_dtype, - shape=self.linear_config.get_conv_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cpu", - size_first=True, - ) - self.ssm_state_cache = LayerCache( - size=self.size, - dtype=self.linear_config.ssm_state_dtype, - shape=self.linear_config.get_ssm_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cpu", - size_first=True, - ) - self.clear_to_init_state() - return - - def get_state_cache(self, buffer_idx: int): - return self.conv_state_cache.buffer[buffer_idx, ...], self.ssm_state_cache.buffer[buffer_idx, ...] - - def alloc_one_state_cache(self) -> Optional[int]: - if len(self.free_list) == 0: - return None - - alloc_index = self.free_list.popleft() - return alloc_index - - def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: - if need_size > len(self.free_list): - logger.error(f"warn no enough cache need_size {need_size} free_size {len(self.free_list)}") - return None - - alloc_indexes = [self.free_list.popleft() for _ in range(need_size)] - return alloc_indexes - - def free_state_cache(self, free_indexes: List[int]): - alloc_upper_bound = self.size - self.keep_num - for idx in free_indexes: - assert 0 <= idx < alloc_upper_bound, ( - f"free index {idx} out of alloc range [0, {alloc_upper_bound}), " f"reserved tail num {self.keep_num}" - ) - self.free_list.extend(free_indexes) - assert ( - len(self.free_list) <= alloc_upper_bound - ), f"free cache num {len(self.free_list)} should not be larger than alloc size {alloc_upper_bound}" - return - - def get_free_cache_num(self): - return len(self.free_list) - - def get_used_cache_num(self): - return self.size - len(self.free_list) - - def clear_to_init_state(self): - self.conv_state_cache.buffer.zero_() - self.ssm_state_cache.buffer.zero_() - self.free_list = collections.deque(range(self.size - self.keep_num)) - return diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index e5d4e85eeb..a17b60ae48 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -2,9 +2,7 @@ import torch -from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig -from lightllm.common.linear_att_cache_manager.layer_cache import LayerCache -from lightllm.common.linear_att_cache_manager.linear_att_buffer_manager import LinearAttCacheManager +from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig, LinearAttCacheManager from lightllm.utils.envs_utils import get_env_start_args from .hybrid_base import HybridAttentionReqManager diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index 1ac36e6efb..609ab5f22b 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -4,7 +4,7 @@ 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.sliding_window_cache_manager import SlidingWindowStateCacheManager +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 @@ -13,7 +13,7 @@ if TYPE_CHECKING: from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager - from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig + from lightllm.common.state_cache_manager import SlidingWindowCacheConfig from lightllm.server.router.model_infer.infer_batch import InferReq @@ -101,19 +101,20 @@ def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): 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( - self.sliding_mem_manager.kv_buffer, - self.req_to_sliding_window, - cache_len, - req.req_idx, - state_cache_manager.get_state_cache(buffer_idx), + 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( - self.sliding_mem_manager.kv_buffer, - self.req_to_sliding_window, - self._sliding_seq_lens[req_idx], - req_idx, - state_cache_manager.get_state_cache(buffer_idx), + 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/sliding_window_cache_manager/__init__.py b/lightllm/common/sliding_window_cache_manager/__init__.py deleted file mode 100644 index b05865d3bc..0000000000 --- a/lightllm/common/sliding_window_cache_manager/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .config import SlidingWindowCacheConfig -from .state_cache import SlidingWindowStateCacheManager - - -__all__ = ["SlidingWindowCacheConfig", "SlidingWindowStateCacheManager"] diff --git a/lightllm/common/sliding_window_cache_manager/state_cache.py b/lightllm/common/sliding_window_cache_manager/state_cache.py deleted file mode 100644 index 2bdeafdb61..0000000000 --- a/lightllm/common/sliding_window_cache_manager/state_cache.py +++ /dev/null @@ -1,53 +0,0 @@ -import collections -from typing import List, Optional - -import torch - -from .config import SlidingWindowCacheConfig - - -class SlidingWindowStateCacheManager: - """大小页共用的 CPU pinned checkpoint 存储，两个池独立分配。 - - 布局为 size-first: [slot, layer, window, 2 * heads, dim]。 - 本类只管理状态存储与空闲槽位，不判断页面大小或缓存边界，也不持有 GPU 运行态。 - """ - - def __init__(self, size: int, sliding_config: SlidingWindowCacheConfig, keep_num: int = 0): - self.size = size - self.keep_num = keep_num - assert 0 <= keep_num <= size - self.state_cache = torch.empty( - (size, *sliding_config.get_state_shape()), - dtype=sliding_config.dtype, - device="cpu", - pin_memory=True, - ) - self.clear_to_init_state() - - def get_state_cache(self, buffer_idx: int): - return self.state_cache[buffer_idx] - - def alloc_one_state_cache(self) -> Optional[int]: - return None if not self.free_list else self.free_list.popleft() - - def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: - if need_size > len(self.free_list): - return None - return [self.free_list.popleft() for _ in range(need_size)] - - def free_state_cache(self, free_indexes: List[int]): - alloc_size = self.size - self.keep_num - assert all(0 <= idx < alloc_size for idx in free_indexes) - self.free_list.extend(free_indexes) - assert len(self.free_list) <= alloc_size - - def get_free_cache_num(self): - return len(self.free_list) - - def get_used_cache_num(self): - return self.size - len(self.free_list) - - def clear_to_init_state(self): - self.state_cache.zero_() - self.free_list = collections.deque(range(self.size - self.keep_num)) diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py new file mode 100644 index 0000000000..ae65b9a334 --- /dev/null +++ b/lightllm/common/state_cache_manager/__init__.py @@ -0,0 +1,16 @@ +from .base import StateCacheManager +from .layer_cache import LayerCache +from .linear_att import LinearAttCacheManager +from .linear_att_config import LinearAttCacheConfig +from .sliding_window import SlidingWindowStateCacheManager +from .sliding_window_config import SlidingWindowCacheConfig + + +__all__ = [ + "StateCacheManager", + "LayerCache", + "LinearAttCacheManager", + "LinearAttCacheConfig", + "SlidingWindowStateCacheManager", + "SlidingWindowCacheConfig", +] diff --git a/lightllm/common/state_cache_manager/base.py b/lightllm/common/state_cache_manager/base.py new file mode 100644 index 0000000000..aa131fe784 --- /dev/null +++ b/lightllm/common/state_cache_manager/base.py @@ -0,0 +1,57 @@ +import collections +from abc import ABC, abstractmethod +from typing import List, Optional + +from lightllm.utils.log_utils import init_logger + + +logger = init_logger(__name__) + + +class StateCacheManager(ABC): + """CPU checkpoint 槽位池；大小页分别实例化，具体状态布局由子类负责。 + + 尾部 keep_num 个槽位保留给 CPU cache 碎页传输，不参与普通分配。 + 本类不管理 GPU 运行态，也不判断页面边界、前缀匹配或淘汰策略。 + """ + + def __init__(self, size: int, keep_num: int = 0): + self.size = size + self.keep_num = keep_num + assert 0 <= keep_num <= size, f"invalid keep_num {keep_num} for size {size}" + self.free_list = collections.deque(range(size - keep_num)) + + @abstractmethod + def get_state_cache(self, buffer_idx: int): + """返回指定槽位的状态视图；可以是单个 Tensor 或多个 Tensor。""" + + def alloc_one_state_cache(self) -> Optional[int]: + return None if not self.free_list else self.free_list.popleft() + + def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: + if need_size > len(self.free_list): + logger.error(f"warn no enough cache need_size {need_size} free_size {len(self.free_list)}") + return None + return [self.free_list.popleft() for _ in range(need_size)] + + def free_state_cache(self, free_indexes: List[int]): + alloc_upper_bound = self.size - self.keep_num + for idx in free_indexes: + assert 0 <= idx < alloc_upper_bound, ( + f"free index {idx} out of alloc range [0, {alloc_upper_bound}), " f"reserved tail num {self.keep_num}" + ) + self.free_list.extend(free_indexes) + assert ( + len(self.free_list) <= alloc_upper_bound + ), f"free cache num {len(self.free_list)} should not be larger than alloc size {alloc_upper_bound}" + + def get_free_cache_num(self): + return len(self.free_list) + + def get_used_cache_num(self): + # Preserve the existing accounting: reserved slots count as used. + return self.size - len(self.free_list) + + def clear_to_init_state(self): + """重置空闲槽位；子类同时清零自身的 checkpoint buffer。""" + self.free_list = collections.deque(range(self.size - self.keep_num)) diff --git a/lightllm/common/linear_att_cache_manager/layer_cache.py b/lightllm/common/state_cache_manager/layer_cache.py similarity index 100% rename from lightllm/common/linear_att_cache_manager/layer_cache.py rename to lightllm/common/state_cache_manager/layer_cache.py diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py new file mode 100644 index 0000000000..a36f970cc3 --- /dev/null +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -0,0 +1,43 @@ +from .base import StateCacheManager +from .layer_cache import LayerCache +from .linear_att_config import LinearAttCacheConfig + + +class LinearAttCacheManager(StateCacheManager): + """CPU pinned conv/SSM checkpoint，两个 buffer 均保持 size-first 布局。""" + + def __init__( + self, + size: int, + linear_config: LinearAttCacheConfig, + keep_num: int = 0, # 用于记录需要保留的缓存数量，用于支持含有 linear_att 的如qwen3.5 模型的cpu cache的碎页处理。 + ): + super().__init__(size, keep_num) + self.linear_config = linear_config + # init the layer cache + self.conv_state_cache = LayerCache( + size=self.size, + dtype=self.linear_config.conv_state_dtype, + shape=self.linear_config.get_conv_state_shape(), + layer_num=self.linear_config.linear_layer_num, + device="cpu", + size_first=True, + ) + self.ssm_state_cache = LayerCache( + size=self.size, + dtype=self.linear_config.ssm_state_dtype, + shape=self.linear_config.get_ssm_state_shape(), + layer_num=self.linear_config.linear_layer_num, + device="cpu", + size_first=True, + ) + return + + def get_state_cache(self, buffer_idx: int): + return self.conv_state_cache.buffer[buffer_idx, ...], self.ssm_state_cache.buffer[buffer_idx, ...] + + def clear_to_init_state(self): + self.conv_state_cache.buffer.zero_() + self.ssm_state_cache.buffer.zero_() + super().clear_to_init_state() + return diff --git a/lightllm/common/linear_att_cache_manager/config_objs.py b/lightllm/common/state_cache_manager/linear_att_config.py similarity index 100% rename from lightllm/common/linear_att_cache_manager/config_objs.py rename to lightllm/common/state_cache_manager/linear_att_config.py 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..35b8d6ee3b --- /dev/null +++ b/lightllm/common/state_cache_manager/sliding_window.py @@ -0,0 +1,24 @@ +import torch + +from .base import StateCacheManager +from .sliding_window_config import SlidingWindowCacheConfig + + +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/common/sliding_window_cache_manager/config.py b/lightllm/common/state_cache_manager/sliding_window_config.py similarity index 100% rename from lightllm/common/sliding_window_cache_manager/config.py rename to lightllm/common/state_cache_manager/sliding_window_config.py diff --git a/lightllm/models/gemma4/kv_layout.py b/lightllm/models/gemma4/kv_layout.py index 186686caf3..87dd516553 100644 --- a/lightllm/models/gemma4/kv_layout.py +++ b/lightllm/models/gemma4/kv_layout.py @@ -1,4 +1,4 @@ -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig +from lightllm.common.state_cache_manager import SlidingWindowCacheConfig def get_kv_cache_layout(config): diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 5e443abc95..dd1a9c883a 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -16,7 +16,7 @@ from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.common.req_manager import ReqManagerForMamba -from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.state_cache_manager import LinearAttCacheConfig logger = init_logger(__name__) diff --git a/lightllm/server/router/dynamic_prompt/linear_att_radix_cache.py b/lightllm/server/router/dynamic_prompt/linear_att_radix_cache.py index 6a8e0a3917..1bb22b1f17 100644 --- a/lightllm/server/router/dynamic_prompt/linear_att_radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/linear_att_radix_cache.py @@ -2,7 +2,7 @@ import numpy as np from typing import Tuple, Dict, Set, List, Optional from sortedcontainers import SortedSet, SortedDict -from lightllm.common.linear_att_cache_manager import LinearAttCacheManager +from lightllm.common.state_cache_manager import StateCacheManager from .shared_arr import SharedArray from .radix_cache import time_gen @@ -132,7 +132,7 @@ def __init__( self.mem_manager: MemoryManager = kv_cache_mem_manager - self.linear_att_big_page_buffers: LinearAttCacheManager = self.mem_manager.linear_att_big_page_buffers + self.linear_att_big_page_buffers: StateCacheManager = self.mem_manager.linear_att_big_page_buffers self._key_dtype = torch.int64 self._value_dtype = torch.int64 @@ -153,7 +153,7 @@ def __init__( f"{unique_name}_tree_total_tokens_num_{rank_in_node}", (1,), dtype=np.int64 ) self.tree_total_tokens_num.arr[0] = 0 - self.linear_att_small_page_buffers: LinearAttCacheManager = linear_att_small_page_buffers + self.linear_att_small_page_buffers: StateCacheManager = linear_att_small_page_buffers def _discard_node(self, node: LinearAttPagedTreeNode): if node.is_leaf(): diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index cd535871f0..0903b9b33c 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -99,7 +99,7 @@ def _validate_flashqla(): from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( chunk_gated_delta_rule as fla_chunk_gated_delta_rule, ) - from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig + from lightllm.common.state_cache_manager import LinearAttCacheConfig linear_config = LinearAttCacheConfig.load_from_args() num_k_heads = linear_config.num_linear_k_heads diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index fdd7fab72a..f5d8a9fcaa 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -38,7 +38,7 @@ from tqdm import tqdm from lightllm.utils.auto_shm_cleanup import register_sysv_shm_for_cleanup from lightllm.utils.dist_utils import get_current_device_id -from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.state_cache_manager import LinearAttCacheConfig from lightllm.common.kv_cache_mem_manager.hybrid_sliding_mem_manager import HybridSlidingMemoryManager logger = init_logger(__name__) diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index 254142050c..f187088883 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -39,7 +39,7 @@ # --------------------------------------------------------------------------- # Step 1 – build LinearAttCacheConfig directly (avoids needing a real model dir) # --------------------------------------------------------------------------- -from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.state_cache_manager import LinearAttCacheConfig linear_config = LinearAttCacheConfig( tp_world_size=8, diff --git a/test/kernel/test_sliding_window_cpu_cache_copy.py b/test/kernel/test_sliding_window_cpu_cache_copy.py index ec25399b03..81720c66ef 100644 --- a/test/kernel/test_sliding_window_cpu_cache_copy.py +++ b/test/kernel/test_sliding_window_cpu_cache_copy.py @@ -8,7 +8,7 @@ copy_kv_buffer_to_cpu_cache, copy_sliding_window_state, ) -from lightllm.common.sliding_window_cache_manager import SlidingWindowCacheConfig +from lightllm.common.state_cache_manager import SlidingWindowCacheConfig pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 12b6996a83..d6574de151 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -10,7 +10,7 @@ import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.common.linear_att_cache_manager.config_objs import LinearAttCacheConfig +from lightllm.common.state_cache_manager import LinearAttCacheConfig from lightllm.server.api_cli import make_argument_parser import lightllm.utils.backend_validator as backend_validator From 47c67bd3a437b7bdb8572d20233e414fcc0c26e2 Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:57:08 +0000 Subject: [PATCH 14/14] Optimize Gemma4 sliding-window slot reuse and attention state --- lightllm/common/req_manager/sliding_window.py | 63 ++++--- .../common/state_cache_manager/__init__.py | 3 +- .../state_cache_manager/sliding_window.py | 79 ++++++++- .../sliding_window_config.py | 79 --------- lightllm/models/gemma4/infer_struct.py | 19 +- .../layer_infer/transformer_layer_infer.py | 46 ++--- lightllm/models/gemma4/model.py | 7 +- .../models/gemma4/test_sliding_runtime.py | 167 ++++++++++++++++++ 8 files changed, 329 insertions(+), 134 deletions(-) delete mode 100644 lightllm/common/state_cache_manager/sliding_window_config.py create mode 100644 unit_tests/models/gemma4/test_sliding_runtime.py diff --git a/lightllm/common/req_manager/sliding_window.py b/lightllm/common/req_manager/sliding_window.py index 609ab5f22b..48782e6f0c 100644 --- a/lightllm/common/req_manager/sliding_window.py +++ b/lightllm/common/req_manager/sliding_window.py @@ -48,35 +48,57 @@ def _init_runtime_buffer(self): # 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) - self._sliding_req_indexes = [torch.empty(0, dtype=torch.int32) for _ in range(self.max_request_num)] + # 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"): - # A cache miss owns no history slots; the forward allocates only its new tokens. - self._release_sliding_window(req.req_idx) + # 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): - """所有层读取后只保留最后 W 个槽位，无需移动 KV 或清空过期映射。""" - indexes = torch.cat((self._sliding_req_indexes[req_idx], new_indexes)) - expired = max(0, indexes.numel() - self.sliding_window) - if expired: - self.sliding_mem_manager.free(indexes[:expired]) - # Retain only the suffix, not the whole forward's pinned index buffer. - self._sliding_req_indexes[req_idx] = indexes[expired:].clone() + """所有层读取后将借用的尾部槽纳入请求窗口，归还被替换和过期的槽；不移动 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 _release_sliding_window(self, req_idx: int): - self.sliding_mem_manager.free(self._sliding_req_indexes[req_idx]) - self._sliding_req_indexes[req_idx] = torch.empty(0, dtype=torch.int32) - self._sliding_seq_lens[req_idx] = 0 - def free_req(self, free_req_index: int): - self._release_sliding_window(free_req_index) + 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_req_indexes = [torch.empty(0, dtype=torch.int32) for _ in range(self.max_request_num)] 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) @@ -92,12 +114,11 @@ def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): # 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) - self._release_sliding_window(req.req_idx) window_len = min(cache_len, self.sliding_window) - # Own the indices: MemoryManager.alloc() returns a reusable staging-buffer view. + # 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_mem_manager.alloc(window_len)) - self._sliding_req_indexes[req.req_idx] = indexes + 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( diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py index ef18ad2280..cbaf8fbe3b 100644 --- a/lightllm/common/state_cache_manager/__init__.py +++ b/lightllm/common/state_cache_manager/__init__.py @@ -1,8 +1,7 @@ from .base import StateCacheManager from .layer_cache import LayerCache from .linear_att import LinearAttCacheConfig, LinearAttCacheManager -from .sliding_window import SlidingWindowStateCacheManager -from .sliding_window_config import SlidingWindowCacheConfig +from .sliding_window import SlidingWindowCacheConfig, SlidingWindowStateCacheManager def get_hybrid_cache_config(): diff --git a/lightllm/common/state_cache_manager/sliding_window.py b/lightllm/common/state_cache_manager/sliding_window.py index 35b8d6ee3b..e4566ac71a 100644 --- a/lightllm/common/state_cache_manager/sliding_window.py +++ b/lightllm/common/state_cache_manager/sliding_window.py @@ -1,7 +1,84 @@ +import dataclasses +import math +from typing import Dict + import torch from .base import StateCacheManager -from .sliding_window_config import SlidingWindowCacheConfig + + +@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): diff --git a/lightllm/common/state_cache_manager/sliding_window_config.py b/lightllm/common/state_cache_manager/sliding_window_config.py deleted file mode 100644 index 04bf294486..0000000000 --- a/lightllm/common/state_cache_manager/sliding_window_config.py +++ /dev/null @@ -1,79 +0,0 @@ -import dataclasses -import math -from typing import Dict - -import torch - - -@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()) diff --git a/lightllm/models/gemma4/infer_struct.py b/lightllm/models/gemma4/infer_struct.py index 8208fb368f..3cc21916f7 100644 --- a/lightllm/models/gemma4/infer_struct.py +++ b/lightllm/models/gemma4/infer_struct.py @@ -26,6 +26,7 @@ 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 @@ -52,7 +53,7 @@ def init_some_extra_state(self, model): token_num = 0 for req_idx, _, q_len in self.sliding_requests: if req_idx != self.req_manager.HOLD_REQUEST_ID: - index_chunks.append(sliding_mem_manager.alloc(q_len)) + 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") @@ -65,7 +66,7 @@ def init_some_extra_state(self, model): (padding_token_num,), sliding_mem_manager.HOLD_TOKEN_MEMINDEX, dtype=torch.int32, device="cpu" ) ) - # Combine allocator views once into owned pinned storage for asynchronous H2D. + # 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 ) @@ -89,13 +90,16 @@ def init_some_extra_state(self, model): self.b_seq_len, self.sliding_window_mem_index, ) - # A metadata view gives the unchanged attention backend its sliding table. - # Tensor metadata is shared with this state, including CUDA graph updates. - sliding_state = copy.copy(self) - sliding_state.req_manager = SimpleNamespace(req_to_token_indexs=self.req_manager.req_to_sliding_window) - self.decode_att_state1 = model.decode_att_backend.create_att_decode_state(infer_state=sliding_state) 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 @@ -130,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/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 0280dafd80..3a681a6a88 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -174,28 +174,30 @@ def _context_attention_kernel( infer_state.sliding_window_mem_index, infer_state.mem_manager.sliding_kv_buffer[self.sliding_cache_index_], ) - # 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) - 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) - - o_tensor = infer_state.prefill_att_state.prefill_att( - q=_q, k=_k, v=_v, att_control=AttControl(), alloc_func=self.alloc_tensor - ) + 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( diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index b6bf936266..447cf7d035 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -133,11 +133,14 @@ def _create_inferstate(self, model_input, microbatch_index=0): return infer_state def _init_att_backend(self): - # Both pools use main's token-indexed attention. Gemma's sliding prefill - # retains its image mask; full layers' head_dim=512 still requires Triton. + # 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): + self.prefill_att_backend1 = TritonAttBackend(model=self) + self.decode_att_backend1 = TritonAttBackend(model=self) + def _init_custom(self): self._init_to_get_rotary_gemma4() if self.config.get("enable_moe_block", False): 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)