diff --git a/lightllm/common/basemodel/attention/fa3/windowed_mtp.py b/lightllm/common/basemodel/attention/fa3/windowed_mtp.py new file mode 100644 index 0000000000..77d65cde10 --- /dev/null +++ b/lightllm/common/basemodel/attention/fa3/windowed_mtp.py @@ -0,0 +1,49 @@ +import weakref + +import torch + +from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState +from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + +class WindowedMTPAttBackend(BaseAttBackend): + """Non-causal block attention over retained draft KV and temporary noise KV.""" + + def create_att_decode_state(self, infer_state): + return WindowedMTPDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +class WindowedMTPDecodeAttState(BaseDecodeAttState): + def init_state(self): + model = self.backend.model + self.block_size = model.block_size + self.b_req_idx = self.infer_state.b_req_idx[:: self.block_size].contiguous() + # The base attention-state copy refreshes these tensors on graph replay. + store = self.infer_state.mem_manager.windowed_draft_kv + self.cache_seqlens = ( + store.req_to_seq_len.index_select(0, self.b_req_idx).clamp(max=store.capacity).to(torch.int32) + + self.block_size + ) + + def prepare_kv(self, layer_index, noise_kv, alloc_func=torch.empty): + """Pack retained history and the current block in the backend's KV layout.""" + scratch = self.infer_state.mem_manager.windowed_draft_kv.pack( + layer_index - self.backend.model.draft_layer_start, + self.b_req_idx, + noise_kv, + self.block_size, + alloc_func=alloc_func, + ) + head_num = noise_kv.shape[1] // 2 + return scratch[:, :, :head_num], scratch[:, :, head_num:] + + def decode_att(self, q, k, v, att_control: AttControl = AttControl(), alloc_func=torch.empty): + return flash_attn_with_kvcache( + q.view(-1, self.block_size, q.shape[-2], q.shape[-1]), + k, + v, + cache_seqlens=self.cache_seqlens, + causal=False, + window_size=att_control.sliding_window if att_control.use_sliding_window else (-1, -1), + softmax_scale=q.shape[-1] ** -0.5, + ).view(q.shape) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 34a1c9c71c..8914010358 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -4,7 +4,6 @@ import gc import copy import math -import json import torch import torch.nn.functional as F import triton @@ -15,8 +14,7 @@ from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.common.req_manager import ReqManager -from lightllm.common.build_utils import repair_config +from lightllm.common.req_manager import ReqManager, ReqManagerForWindowedMTP from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( select_kv_index_from_req, select_kv_index_from_req_prefill, @@ -27,6 +25,7 @@ from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed from lightllm.utils.log_utils import init_logger +from lightllm.utils.config_utils import load_config from lightllm.utils.dist_utils import get_dp_world_size from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory from lightllm.utils.envs_utils import ( @@ -146,12 +145,7 @@ def __init__(self, kvargs): return def _init_config(self): - with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: - self.config = json.load(json_file) - # rename keys - repair_config(self.config, same_names=["num_attention_heads", "n_head"]) - repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) + self.config = load_config(self.weight_dir_) if self.finetune_config: self.config["vocab_size"] = self.finetune_config.vocab_size return @@ -237,7 +231,8 @@ def _init_req_manager(self): if self.max_seq_length is not None: create_max_seq_len = max(create_max_seq_len, self.max_seq_length) - self.req_manager = ReqManager(self.max_req_num, create_max_seq_len, None) + req_manager_class = ReqManagerForWindowedMTP if self.args.mtp_windowed else ReqManager + self.req_manager = req_manager_class(self.max_req_num, create_max_seq_len, None) return def _init_infer_layer(self, start_layer_index=0): diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py index 060a92bf7c..98ebe3ab2c 100644 --- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py +++ b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py @@ -29,8 +29,9 @@ def _fwd_kernel_destindex_copy_kv( k_ptrs = K + cur_index * stride_k_bs + stride_k_h * offs_h[:, None] + stride_k_d * offs_d[None, :] o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] - k = tl.load(k_ptrs, mask=offs_h[:, None] < head_num, other=0.0) - tl.store(o_ptrs, k, mask=offs_h[:, None] < head_num) + mask = (dest_index >= 0) & (offs_h[:, None] < head_num) + k = tl.load(k_ptrs, mask=mask, other=0.0) + tl.store(o_ptrs, k, mask=mask) return diff --git a/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py b/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py new file mode 100644 index 0000000000..3ba6c30b4a --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py @@ -0,0 +1,54 @@ +"""Stream-ordered copies for pinned checkpoint pages.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _copy_checkpoint_pages( + src_state, + dst_state, + src_mem_index, + dst_mem_index, + page_readies, + src_stride_s, + dst_stride_s, + state_dim, + page_num, + SKIP_READY: tl.constexpr, + BLOCK: tl.constexpr, +): + for page in range(page_num): + src_id = tl.load(src_mem_index + page).to(tl.int64) + dst_id = tl.load(dst_mem_index + page).to(tl.int64) + valid = (src_id >= 0) & (dst_id >= 0) + if SKIP_READY: + valid = valid & (tl.load(page_readies + page) == 0) + if valid: + for block in range(tl.program_id(0), tl.cdiv(state_dim, BLOCK), tl.num_programs(0)): + offsets = block * BLOCK + tl.arange(0, BLOCK) + data = tl.load(src_state + src_id * src_stride_s + offsets, offsets < state_dim, other=0) + tl.store(dst_state + dst_id * dst_stride_s + offsets, data, offsets < state_dim, cache_modifier=".wt") + + +def copy_checkpoint_pages(sources, destinations, source_ids, destination_ids, page_readies=None): + """Copy slot-indexed, contiguous rows; preserve each tensor's byte representation.""" + assert len(source_ids) == len(destination_ids) + for source, destination in zip(sources, destinations): + source = source.reshape(source.shape[0], -1).view(torch.uint8) + destination = destination.reshape(destination.shape[0], -1).view(torch.uint8) + assert source.shape[1] == destination.shape[1] + _copy_checkpoint_pages[(12,)]( + source, + destination, + source_ids, + destination_ids, + page_readies, + source.stride(0), + destination.stride(0), + source.shape[1], + len(source_ids), + SKIP_READY=page_readies is not None, + BLOCK=4096, + ) diff --git a/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py b/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py index 0fdc43ab9f..6ec91e008d 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py +++ b/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py @@ -61,6 +61,7 @@ def _offload_gpu_kv_to_cpu( token_range = block_index * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK) token_indexes = tl.load(token_indexes_ptr + token_range).to(tl.int64) + token_mask = token_indexes >= 0 head_dim_range = tl.arange(0, BLOCK_HEAD_DIM) head_dim_mask = head_dim_range < head_dim scale_head_dim_mask = head_dim_range < scale_head_dim @@ -77,7 +78,7 @@ def _offload_gpu_kv_to_cpu( + gpu_k_head_index.to(tl.int64) * gpu_stride2 + head_dim_range[None, :] ) - gpu_data = tl.load(gpu_ptr, mask=head_dim_mask[None, :], other=0.0) + gpu_data = tl.load(gpu_ptr, mask=token_mask[:, None] & head_dim_mask[None, :], other=0.0) cpu_ptr = ( cpu_kv_cache_ptr + cpu_page_index * cpu_stride0 @@ -89,7 +90,7 @@ def _offload_gpu_kv_to_cpu( tl.store( cpu_ptr, gpu_data, - mask=head_dim_mask[None, :], + mask=token_mask[:, None] & head_dim_mask[None, :], cache_modifier=".wt", ) if HAS_SCALE: @@ -127,7 +128,7 @@ def _offload_gpu_kv_to_cpu( + gpu_v_head_index.to(tl.int64) * gpu_stride2 + head_dim_range[None, :] ) - gpu_data = tl.load(gpu_ptr, mask=head_dim_mask[None, :], other=0.0) + gpu_data = tl.load(gpu_ptr, mask=token_mask[:, None] & head_dim_mask[None, :], other=0.0) cpu_ptr = ( cpu_kv_cache_ptr + cpu_page_index * cpu_stride0 @@ -139,7 +140,7 @@ def _offload_gpu_kv_to_cpu( tl.store( cpu_ptr, gpu_data, - mask=head_dim_mask[None, :], + mask=token_mask[:, None] & head_dim_mask[None, :], cache_modifier=".wt", ) if HAS_SCALE: @@ -407,6 +408,7 @@ def _load_cpu_cache_to_gpu( gpu_mem_indexes = tl.load(gpu_mem_indexes_ptr + token_range, mask=token_mask).to(tl.int64) cpu_mem_indexes = tl.load(cpu_mem_indexes_ptr + token_range, mask=token_mask).to(tl.int64) cpu_page_indexes = tl.load(cpu_page_indexes_ptr + token_range, mask=token_mask).to(tl.int64) + token_mask = token_mask & (gpu_mem_indexes >= 0) & (cpu_page_indexes >= 0) head_dim_range = tl.arange(0, BLOCK_HEAD_DIM) head_dim_mask = head_dim_range < head_dim diff --git a/lightllm/common/basemodel/triton_kernel/windowed_mtp.py b/lightllm/common/basemodel/triton_kernel/windowed_mtp.py new file mode 100644 index 0000000000..5030ad7ec5 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/windowed_mtp.py @@ -0,0 +1,151 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _prepare_window_write_indices( + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_commit_len, + out_mem_index, + req_to_seq_len, + window_size: tl.constexpr, + BLOCK: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_req_idx = tl.load(b_req_idx + cur_batch).to(tl.int64) + cur_ready_cache_len = tl.load(b_ready_cache_len + cur_batch) + cur_commit_len = tl.load(b_commit_len + cur_batch) + cur_seq_len = cur_ready_cache_len + cur_commit_len + cur_start_loc = tl.load(b_start_loc + cur_batch) + + # Only the retained tail may write: older rows would collide in the ring. + tail_start = tl.maximum(0, cur_commit_len - window_size) + for start in range(tail_start, cur_commit_len, BLOCK): + row = start + tl.arange(0, BLOCK) + position = cur_ready_cache_len + row + mem_index = cur_req_idx * window_size + position % window_size + tl.store(out_mem_index + cur_start_loc + row, mem_index, mask=row < cur_commit_len) + # Publish the new length now; the caller must order all layer KV writes before any state reads. + tl.store(req_to_seq_len + cur_req_idx, cur_seq_len) + + +def prepare_window_write_indices( + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_commit_len, + req_to_seq_len, + window_size, + token_num, + alloc_func=torch.empty, +): + """Return one destination per source token; rejected and expired rows stay -1.""" + mem_index = alloc_func((token_num,), dtype=torch.int32, device=b_req_idx.device) + mem_index.fill_(-1) + if b_req_idx.numel(): + _prepare_window_write_indices[(b_req_idx.numel(),)]( + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_commit_len, + mem_index, + req_to_seq_len, + window_size=window_size, + BLOCK=256, + ) + return mem_index + + +@triton.jit +def _pack_window_kv( + kv_buffer, + noise_kv, + b_req_idx, + req_to_seq_len, + out_kv, + stride_kv_bs, + stride_kv_h, + stride_kv_d, + stride_noise_bs, + stride_noise_h, + stride_noise_d, + stride_out_b, + stride_out_s, + stride_out_h, + stride_out_d, + window_size: tl.constexpr, + block_size: tl.constexpr, + head_num: tl.constexpr, + head_dim: tl.constexpr, + BLOCK_HEAD: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + cur_token = tl.program_id(0) + cur_batch = tl.program_id(1) + cur_req_idx = tl.load(b_req_idx + cur_batch).to(tl.int64) + cur_seq_len = tl.load(req_to_seq_len + cur_req_idx) + cur_cache_len = tl.minimum(cur_seq_len, window_size) + cur_position = tl.maximum(0, cur_seq_len - window_size) + cur_token + cur_mem_index = cur_req_idx * window_size + cur_position % window_size + cur_noise_index = cur_batch * block_size + cur_token - cur_cache_len + offs_h = tl.arange(0, BLOCK_HEAD) + offs_d = tl.arange(0, BLOCK_DMODEL) + head_mask = (offs_h[:, None] < head_num) & (offs_d[None, :] < head_dim) + + kv = tl.load( + kv_buffer + cur_mem_index * stride_kv_bs + offs_h[:, None] * stride_kv_h + offs_d[None, :] * stride_kv_d, + mask=(cur_token < cur_cache_len) & head_mask, + other=0, + ) + noise = tl.load( + noise_kv + + cur_noise_index.to(tl.int64) * stride_noise_bs + + offs_h[:, None] * stride_noise_h + + offs_d[None, :] * stride_noise_d, + mask=(cur_token >= cur_cache_len) & (cur_token < cur_cache_len + block_size) & head_mask, + other=0, + ) + out_ptr = ( + out_kv + + cur_batch * stride_out_b + + cur_token * stride_out_s + + offs_h[:, None] * stride_out_h + + offs_d[None, :] * stride_out_d + ) + tl.store(out_ptr, tl.where(cur_token < cur_cache_len, kv, noise), mask=head_mask) + + +def pack_window_kv(kv_buffer, noise_kv, b_req_idx, req_to_seq_len, window_size, block_size, alloc_func=torch.empty): + head_num, head_dim = kv_buffer.shape[-2:] + out_kv = alloc_func( + (b_req_idx.numel(), window_size + block_size, head_num, head_dim), dtype=noise_kv.dtype, device=noise_kv.device + ) + if b_req_idx.numel(): + _pack_window_kv[(window_size + block_size, b_req_idx.numel())]( + kv_buffer, + noise_kv, + b_req_idx, + req_to_seq_len, + out_kv, + stride_kv_bs=kv_buffer.stride(0), + stride_kv_h=kv_buffer.stride(1), + stride_kv_d=kv_buffer.stride(2), + stride_noise_bs=noise_kv.stride(0), + stride_noise_h=noise_kv.stride(1), + stride_noise_d=noise_kv.stride(2), + stride_out_b=out_kv.stride(0), + stride_out_s=out_kv.stride(1), + stride_out_h=out_kv.stride(2), + stride_out_d=out_kv.stride(3), + window_size=window_size, + block_size=block_size, + head_num=head_num, + head_dim=head_dim, + BLOCK_HEAD=triton.next_power_of_2(head_num), + BLOCK_DMODEL=triton.next_power_of_2(head_dim), + num_warps=4, + ) + return out_kv diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 7108f39d79..911d4e0296 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -8,6 +8,8 @@ from lightllm.utils.log_utils import init_logger from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt from .allocator import KvCacheAllocator +from .windowed_mtp import WindowKVStore, WindowKVPageHelper +from lightllm.common.state_cache_manager.windowed_mtp import load_window_state_config from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory from lightllm.utils.dist_utils import ( get_current_device_id, @@ -38,6 +40,10 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False self.layer_num = layer_num self.always_copy = always_copy self.dtype = dtype + self.windowed_draft_kv = None + self._windowed_draft_max_req_num = None + args = get_env_start_args() + self.window_state_config = load_window_state_config(args) if args.mtp_windowed else None # profile the max total token num if the size is None self.profile_size(mem_fraction) @@ -59,7 +65,12 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False self.HOLD_TOKEN_MEMINDEXES = tuple(range(self.size, self.size + self.page_size)) # 构建对外的操作类接口 - self.operator: BaseMemManagerOperator = self.operator_class(self) + operator_class = self.operator_class + if self.window_state_config is not None and operator_class is NormalMemOperator: + from .operator.windowed_mtp import WindowedMTPMemOperator + + operator_class = WindowedMTPMemOperator + self.operator: BaseMemManagerOperator = operator_class(self) def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: k = self.kv_buffer[layer_index][:, : self.head_num, :] @@ -120,6 +131,39 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): self.kv_buffer = torch.empty( (layer_num, size + self.page_size, 2 * head_num, head_dim), dtype=dtype, device="cuda" ) + if self._windowed_draft_max_req_num is not None: + self.init_windowed_draft_kv(self._windowed_draft_max_req_num) + self._init_state_cache_buffers() + + def _init_state_cache_buffers(self): + self.big_page_buffers = None + if self.window_state_config is not None: + from lightllm.common.state_cache_manager import WindowStateCacheManager + + args = get_env_start_args() + page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num + self.big_page_buffers = WindowStateCacheManager( + size=(self.size + page_tokens - 1) // page_tokens + 2, + config=self.window_state_config, + dtype=self.dtype, + keep_num=2, + ) + self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 + + def init_windowed_draft_kv(self, max_request_num): + """Allocate request windows from the shared draft-state configuration.""" + config = self.window_state_config + self._windowed_draft_max_req_num = max_request_num + self.windowed_draft_kv = WindowKVStore( + requests=max_request_num + 1, + layers=config.layers, + kv_heads=config.kv_heads, + head_dim=config.head_dim, + dtype=self.dtype, + device=self.kv_buffer.device, + window=config.window, + ) def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: self.kv_move_buffer = torch.empty( @@ -128,8 +172,13 @@ def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: self._buffer_mem_indexes_tensors = [ torch.empty((page_size,), dtype=torch.int64, device="cpu", pin_memory=True) for _ in range(page_num) ] + self._assert_att_state_page_size() return self.kv_move_buffer + def _assert_att_state_page_size(self): + if self.windowed_draft_kv is not None: + WindowKVPageHelper(self).assert_page_size() + def write_mem_to_page_kv_move_buffer( self, mem_indexes: List[int], @@ -140,6 +189,11 @@ def write_mem_to_page_kv_move_buffer( page_kind: str = "kv", req_idx: int = None, ): + if page_kind == "att_state" and self.windowed_draft_kv is not None: + dp_mems = mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] + assert len(dp_mems) == dp_world_size + WindowKVPageHelper(self).copy_req_page(page_index, req_idx, dp_mems, mode="write") + return assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" cur_page = self.kv_move_buffer[page_index] pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] @@ -172,6 +226,11 @@ def read_page_kv_move_buffer_to_mem( page_kind: str = "kv", req_idx: int = None, ): + if page_kind == "att_state" and self.windowed_draft_kv is not None: + dp_mems = mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] + assert len(dp_mems) == dp_world_size + WindowKVPageHelper(self).copy_req_page(page_index, req_idx, dp_mems, mode="read") + return assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" cur_page = self.kv_move_buffer[page_index] pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] @@ -196,6 +255,8 @@ def read_page_kv_move_buffer_to_mem( def _free_buffers(self): self.kv_buffer = None + self.windowed_draft_kv = None + self.big_page_buffers = None def alloc(self, need_size) -> torch.Tensor: assert ( @@ -208,6 +269,8 @@ def free(self, free_index: Union[torch.Tensor, List[int]]) -> None: def free_all(self): self.allocator.free_all() + if self.windowed_draft_kv is not None: + self.windowed_draft_kv.reset() def resize_mem(self, new_size): """ @@ -253,7 +316,14 @@ def write_to_shm(self, req_manager): with lock: node_world_size = get_node_world_size() shm_name = f"mem_manager_{get_current_rank_in_node()}" - obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] + # Only GPU running state is consumed by PD. Pickling CPU checkpoints + # would move their storage into ordinary shared memory and lose pinning. + big_page_buffers = getattr(self, "big_page_buffers", None) + self.big_page_buffers = None + try: + obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] + finally: + self.big_page_buffers = big_page_buffers # Each consumer needs its own IPC ref-counter; its pickle size can vary. obj_size = max(len(obj_bytes) for obj_bytes in obj_bytes_array) shm = create_or_link_shm( diff --git a/lightllm/common/kv_cache_mem_manager/operator/hybrid.py b/lightllm/common/kv_cache_mem_manager/operator/hybrid.py new file mode 100644 index 0000000000..9b80012bb8 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid.py @@ -0,0 +1,185 @@ +import torch +import triton +from .base import BaseMemManagerOperator +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.dist_utils import get_current_rank_in_dp +from lightllm.common.state_cache_manager.windowed_mtp import WindowedMTPCacheConfig +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient + from lightllm.server.router.model_infer.infer_batch import InferReq + + +class HybridAttMemOperator(BaseMemManagerOperator): + """Shared CPU transfer lifecycle for token KV with checkpoint state.""" + + def __init__(self, mem_manager): + super().__init__(mem_manager) + if mem_manager.window_state_config is not None: + self.cpu_cache_config = WindowedMTPCacheConfig.load_from_args( + linear_config=getattr(mem_manager, "linear_config", None) + ) + + 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 + ) + return + + def load_cpu_cache_to_gpu( + self, + mem_indexes: torch.Tensor, + page_indexes: torch.Tensor, + cpu_cache_client: "CpuKvCacheClient", + req: "InferReq", + ): + assert mem_indexes.is_cuda and page_indexes.is_cuda + args = get_env_start_args() + assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) + assert len(mem_indexes) % args.linear_att_hash_page_size == 0 + assert args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num + mem_manager = self.mem_manager + + big_page_num = len(mem_indexes) // args.cpu_cache_token_page_size + max_kv_len = (req.cur_kv_len // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size + assert max_kv_len % args.cpu_cache_token_page_size == 0 + + big_page_buffer_ids_cpu = [] + for i in range(big_page_num): + page_id = mem_manager.big_page_buffers.alloc_one_state_cache() + assert page_id is not None + req.hybrid_len_to_big_page_id[max_kv_len] = page_id + big_page_buffer_ids_cpu.append(page_id) + max_kv_len -= args.cpu_cache_token_page_size + assert max_kv_len % args.cpu_cache_token_page_size == 0 + + big_page_buffer_ids_cpu.reverse() + + # 碎页情况的处理 + has_tail_page = len(mem_indexes) % args.cpu_cache_token_page_size != 0 + if has_tail_page: + padded_token_num = triton.cdiv( + len(mem_indexes), args.cpu_cache_token_page_size + ) * args.cpu_cache_token_page_size - len(mem_indexes) + mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), mode="constant", value=-1) + + # 将对应的小叶数据拷贝到临时的大页上，再从大页上拷贝到对应的运行态页面上 + big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) + + big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( + non_blocking=True + ) + + assert len(big_page_buffer_ids_gpu) == len(page_indexes) + + self._load_target_pages(mem_indexes, big_page_buffer_ids_gpu, page_indexes, cpu_cache_client) + self._copy_window_pages(big_page_buffer_ids_gpu, page_indexes, cpu_cache_client) + + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + g_infer_context.req_manager.restore_big_page_state( + big_page_buffer_idx=big_page_buffer_ids_cpu[-1], + req=req, + ) + + return + + def offload_gpu_kv_to_cpu_cache( + self, + mem_indexes: torch.Tensor, + page_indexes: torch.Tensor, + page_readies: torch.Tensor, + cpu_cache_client: "CpuKvCacheClient", + req: "InferReq", + ): + args = get_env_start_args() + if not hasattr(self, "big_page_ids_buffer_store"): + self.big_page_ids_buffer_store = torch.empty((1024 * 1024 * 4,), dtype=torch.int64, device="cuda") + # 多申请3个cpu cache token page size，用于处理碎页情况，碎页情况需要将对应的大页数据拷贝到临时的大页上， + # 再从大页上拷贝到对应的运行态页面上 + self.mem_indexes_buffer = torch.empty( + (args.max_req_total_len + 3 * args.cpu_cache_token_page_size,), dtype=torch.int32, device="cuda" + ) + + assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda + + assert len(mem_indexes) % args.linear_att_hash_page_size == 0 + assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) + + mem_manager = self.mem_manager + + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + big_page_buffer_ids_cpu = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) + max_kv_len = (len(mem_indexes) // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size + start_kv_len = (len(big_page_buffer_ids_cpu) + 1) * args.cpu_cache_token_page_size + for seq_len in range(start_kv_len, max_kv_len + 1, args.cpu_cache_token_page_size): + page_id = req.hybrid_len_to_big_page_id[seq_len] + big_page_buffer_ids_cpu.append(page_id) + + if len(mem_indexes) % args.cpu_cache_token_page_size != 0: + # 存在不满大页的碎页的页面存在需要复制的情况 + dst_len = triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size + assert dst_len <= self.mem_indexes_buffer.shape[0] + dst_mem_indexes = self.mem_indexes_buffer[0:dst_len].fill_(-1) + dst_mem_indexes[0 : len(mem_indexes)].copy_(mem_indexes, non_blocking=True) + mem_indexes = dst_mem_indexes + assert req.tail_small_page_buffer_id is not None + self._copy_tail_state( + g_infer_context.radix_cache.small_page_buffers, + req.tail_small_page_buffer_id, + mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID, + ) + big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID) + + assert len(big_page_buffer_ids_cpu) == len(page_indexes) == len(page_readies) + + big_page_buffer_ids_cpu = torch.tensor( + big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu", pin_memory=True + ) + assert len(big_page_buffer_ids_cpu) <= self.big_page_ids_buffer_store.shape[0] + big_page_buffer_ids_gpu = self.big_page_ids_buffer_store[0 : len(big_page_buffer_ids_cpu)] + big_page_buffer_ids_gpu.copy_(big_page_buffer_ids_cpu, non_blocking=True) + + self._offload_target_pages(mem_indexes, big_page_buffer_ids_gpu, page_indexes, page_readies, cpu_cache_client) + self._copy_window_pages(big_page_buffer_ids_gpu, page_indexes, cpu_cache_client, page_readies) + return + + def _load_target_pages(self, mem_indexes, buffer_ids, page_indexes, cpu_cache_client): + raise NotImplementedError() + + def _offload_target_pages(self, mem_indexes, buffer_ids, page_indexes, page_readies, cpu_cache_client): + raise NotImplementedError() + + def _copy_target_state(self, source, source_id, destination_id): + # A full-attention target has no additional checkpoint payload. + return + + def _copy_tail_state(self, source, source_id, destination_id): + self._copy_target_state(source, source_id, destination_id) + window = self.mem_manager.big_page_buffers.draft_window + if window is not None: + from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( + copy_linear_att_state_to_linear_att_state, + ) + + # The existing copier handles both contiguous checkpoint tensors as bytes. + src_kv, src_seq_lens = source.draft_window.get_state_cache(source_id) + dst_kv, dst_seq_lens = window.get_state_cache(destination_id) + copy_linear_att_state_to_linear_att_state(src_kv, src_seq_lens, dst_kv, dst_seq_lens) + + def _copy_window_pages(self, buffer_ids, page_indexes, cpu_cache_client, page_readies=None): + window = self.mem_manager.big_page_buffers.draft_window + if window is None: + return + from lightllm.common.basemodel.triton_kernel.hybrid_cpu_cache_copy import copy_checkpoint_pages + + buffers = (window.kv_buffer, window.seq_lens) + pages = self.cpu_cache_config.get_window_views(cpu_cache_client.cpu_kv_cache_tensor, get_current_rank_in_dp()) + if page_readies is None: + copy_checkpoint_pages(pages, buffers, page_indexes, buffer_ids) + else: + copy_checkpoint_pages(buffers, pages, buffer_ids, page_indexes, page_readies) 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 e935539e37..ebcf5e37a0 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py +++ b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py @@ -1,20 +1,13 @@ import torch -import triton -from typing import List -from typing import TYPE_CHECKING -from .base import BaseMemManagerOperator +from .hybrid import HybridAttMemOperator 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 -if TYPE_CHECKING: - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient - from lightllm.server.router.model_infer.infer_batch import InferReq - logger = init_logger(__name__) -class LinearAttMemOperator(BaseMemManagerOperator): +class LinearAttMemOperator(HybridAttMemOperator): """ 只用于非量化的linear att 混合 full att的模型，列入 qwen3.5 """ @@ -23,54 +16,9 @@ def __init__(self, mem_manager): super().__init__(mem_manager) self.linear_config = mem_manager.linear_config - def load_cpu_cache_to_gpu( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda + def _load_target_pages(self, mem_indexes, big_page_buffer_ids_gpu, page_indexes, cpu_cache_client): args = get_env_start_args() - assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager - - mem_manager: Qwen3NextMemManager = self.mem_manager - - big_page_num = len(mem_indexes) // args.cpu_cache_token_page_size - max_kv_len = (req.cur_kv_len // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - assert max_kv_len % args.cpu_cache_token_page_size == 0 - - big_page_buffer_ids_cpu = [] - for i in range(big_page_num): - page_id = mem_manager.big_page_buffers.alloc_one_state_cache() - assert page_id is not None - req.hybrid_len_to_big_page_id[max_kv_len] = page_id - big_page_buffer_ids_cpu.append(page_id) - max_kv_len -= args.cpu_cache_token_page_size - assert max_kv_len % args.cpu_cache_token_page_size == 0 - - big_page_buffer_ids_cpu.reverse() - - # 碎页情况的处理 - has_tail_page = len(mem_indexes) % args.cpu_cache_token_page_size != 0 - if has_tail_page: - padded_token_num = triton.cdiv( - len(mem_indexes), args.cpu_cache_token_page_size - ) * args.cpu_cache_token_page_size - len(mem_indexes) - mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), mode="constant", value=-1) - - # 将对应的小叶数据拷贝到临时的大页上，再从大页上拷贝到对应的运行态页面上 - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) - - big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( - non_blocking=True - ) - - assert len(big_page_buffer_ids_gpu) == len(page_indexes) - + mem_manager = self.mem_manager from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( copy_cpu_cache_to_kv_buffer, ) @@ -89,85 +37,9 @@ def load_cpu_cache_to_gpu( linear_config=self.linear_config, ) - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=big_page_buffer_ids_cpu[-1], - req=req, - ) - - return - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): + def _offload_target_pages(self, mem_indexes, big_page_buffer_ids_gpu, page_indexes, page_readies, cpu_cache_client): args = get_env_start_args() - if not hasattr(self, "big_page_ids_buffer_store"): - self.big_page_ids_buffer_store = torch.empty((1024 * 1024 * 4,), dtype=torch.int64, device="cuda") - # 多申请3个cpu cache token page size，用于处理碎页情况，碎页情况需要将对应的大页数据拷贝到临时的大页上， - # 再从大页上拷贝到对应的运行态页面上 - self.mem_indexes_buffer = torch.empty( - (args.max_req_total_len + 3 * args.cpu_cache_token_page_size,), dtype=torch.int32, device="cuda" - ) - - assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda - - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) - - from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager - - mem_manager: Qwen3NextMemManager = self.mem_manager - - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - big_page_buffer_ids_cpu = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) - max_kv_len = (len(mem_indexes) // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - start_kv_len = (len(big_page_buffer_ids_cpu) + 1) * args.cpu_cache_token_page_size - for seq_len in range(start_kv_len, max_kv_len + 1, args.cpu_cache_token_page_size): - page_id = req.hybrid_len_to_big_page_id[seq_len] - big_page_buffer_ids_cpu.append(page_id) - - if len(mem_indexes) % args.cpu_cache_token_page_size != 0: - # 存在不满大页的碎页的页面存在需要复制的情况 - dst_len = triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - assert dst_len <= self.mem_indexes_buffer.shape[0] - dst_mem_indexes = self.mem_indexes_buffer[0:dst_len].fill_(-1) - dst_mem_indexes[0 : len(mem_indexes)].copy_(mem_indexes, non_blocking=True) - mem_indexes = dst_mem_indexes - assert req.tail_small_page_buffer_id is not None - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_linear_att_state_to_linear_att_state, - ) - - src_conv_state, src_ssm_state = g_infer_context.radix_cache.small_page_buffers.get_state_cache( - buffer_idx=req.tail_small_page_buffer_id - ) - dst_conv_state, dst_ssm_state = mem_manager.big_page_buffers.get_state_cache( - buffer_idx=mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID, - ) - copy_linear_att_state_to_linear_att_state( - src_conv_state=src_conv_state, - src_ssm_state=src_ssm_state, - dst_conv_state=dst_conv_state, - dst_ssm_state=dst_ssm_state, - ) - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID) - - assert len(big_page_buffer_ids_cpu) == len(page_indexes) == len(page_readies) - - big_page_buffer_ids_cpu = torch.tensor( - big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu", pin_memory=True - ) - assert len(big_page_buffer_ids_cpu) <= self.big_page_ids_buffer_store.shape[0] - big_page_buffer_ids_gpu = self.big_page_ids_buffer_store[0 : len(big_page_buffer_ids_cpu)] - big_page_buffer_ids_gpu.copy_(big_page_buffer_ids_cpu, non_blocking=True) - + mem_manager = self.mem_manager from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( copy_kv_buffer_to_cpu_cache, ) @@ -186,7 +58,20 @@ def offload_gpu_kv_to_cpu_cache( big_page_token_num=args.cpu_cache_token_page_size, linear_config=self.linear_config, ) - return + + def _copy_target_state(self, source, source_id, destination_id): + from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( + copy_linear_att_state_to_linear_att_state, + ) + + src_conv_state, src_ssm_state = source.get_state_cache(source_id) + dst_conv_state, dst_ssm_state = self.mem_manager.big_page_buffers.get_state_cache(destination_id) + copy_linear_att_state_to_linear_att_state( + src_conv_state=src_conv_state, + src_ssm_state=src_ssm_state, + dst_conv_state=dst_conv_state, + dst_ssm_state=dst_ssm_state, + ) def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): # Qwen3Next 需要调整 layer_index @@ -200,11 +85,3 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: destindex_copy_kv(kv, mem_index, mem_manager.kv_buffer[layer_index]) return - - 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 - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py new file mode 100644 index 0000000000..7f71145b20 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py @@ -0,0 +1,35 @@ +from .hybrid import HybridAttMemOperator +from .normal import NormalMemOperator +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size +from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu, offload_gpu_kv_to_cpu + + +class WindowedMTPMemOperator(HybridAttMemOperator, NormalMemOperator): + """Ordinary target KV plus request-local draft windows in hybrid CPU pages.""" + + def _load_target_pages(self, mem_indexes, buffer_ids, page_indexes, cpu_cache_client): + load_cpu_kv_to_gpu( + gpu_mem_indexes=mem_indexes, + gpu_kv_cache=self.mem_manager.kv_buffer, + gpu_kv_cache_scale=None, + cpu_kv_cache=self.cpu_cache_config.get_full_att_kv_view(cpu_cache_client.cpu_kv_cache_tensor), + cpu_kv_cache_scale=None, + page_indexes=page_indexes, + tp_index=get_current_rank_in_dp(), + tp_world_size=get_dp_world_size(), + grid_num=16, + ) + + def _offload_target_pages(self, mem_indexes, buffer_ids, page_indexes, page_readies, cpu_cache_client): + offload_gpu_kv_to_cpu( + token_indexes=mem_indexes, + gpu_kv_cache=self.mem_manager.kv_buffer, + gpu_kv_cache_scale=None, + cpu_kv_cache=self.cpu_cache_config.get_full_att_kv_view(cpu_cache_client.cpu_kv_cache_tensor), + cpu_kv_cache_scale=None, + page_indexes=page_indexes, + page_readies=page_readies, + tp_index=get_current_rank_in_dp(), + tp_world_size=get_dp_world_size(), + grid_num=16, + ) 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 1962fa71b0..8cda62e43a 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -11,6 +11,7 @@ FP8StaticPerTensorQuantMemOperator, LinearAttMemOperator, ) +from .windowed_mtp import WindowKVPageHelper from typing import Tuple, Any, List logger = init_logger(__name__) @@ -38,11 +39,8 @@ def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) return super().get_att_input_params(layer_index) - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - super()._init_buffers(size, dtype, head_num, head_dim, layer_num) - # TODO 初始化线性 att 对应的部分 buffer. + def _init_state_cache_buffers(self): self._init_linear_att_buffers() - return def _init_linear_att_buffers(self): big_page_token_num = ( @@ -55,21 +53,13 @@ def _init_linear_att_buffers(self): size=triton.cdiv(self.size, big_page_token_num) + 2, linear_config=self.linear_config, keep_num=2, + window_config=self.window_state_config, ) self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 return - def _free_buffers(self): - super()._free_buffers() - self._free_linear_att_buffers() - return - - def _free_linear_att_buffers(self): - self.big_page_buffers = None - return - def write_to_shm(self, req_manager): self.req_to_conv_state = req_manager.req_to_conv_state self.req_to_ssm_state = req_manager.req_to_ssm_state @@ -85,11 +75,11 @@ def write_to_shm(self, req_manager): finally: self.big_page_buffers = big_page_buffers - def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: - kv_move_buffer = super().alloc_paged_kv_move_buffer(page_num, page_size) + def _assert_att_state_page_size(self): self.att_state_page_helper = self._create_att_state_page_helper() self.att_state_page_helper.assert_page_size() - return kv_move_buffer + if self.windowed_draft_kv is not None: + WindowKVPageHelper(self, offset=self.att_state_page_helper.state_nbytes).assert_page_size() def _create_att_state_page_helper(self): return Qwen3NextLinearAttPageHelper(self) @@ -119,6 +109,10 @@ def write_mem_to_page_kv_move_buffer( helper = self.att_state_page_helper dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) helper.write_req_to_page(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) + if self.windowed_draft_kv is not None: + WindowKVPageHelper(self, offset=helper.state_nbytes).copy_req_page( + page_index, req_idx, dp_mems, mode="write" + ) return def read_page_kv_move_buffer_to_mem( @@ -146,6 +140,10 @@ def read_page_kv_move_buffer_to_mem( helper = self.att_state_page_helper dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) helper.read_page_to_req(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) + if self.windowed_draft_kv is not None: + WindowKVPageHelper(self, offset=helper.state_nbytes).copy_req_page( + page_index, req_idx, dp_mems, mode="read" + ) return diff --git a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py new file mode 100644 index 0000000000..e867ce77c8 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py @@ -0,0 +1,142 @@ +"""Request-local draft windows and their existing checkpoint/PD interfaces.""" + +import torch + + +class WindowKVStore: + """Use the normal [layer, token, head, dim] KV layout for each request's ring. + + Request r owns token slots [r * window, (r + 1) * window). Its committed + sequence length determines the visible positions; noise KV stays temporary. + """ + + def __init__(self, requests, layers, kv_heads, head_dim, dtype, device, window): + self.capacity = window + self.kv_buffer = torch.empty((layers, requests * window, 2 * kv_heads, head_dim), dtype=dtype, device=device) + self.req_to_seq_len = torch.zeros(requests, dtype=torch.int64, device=device) + + def reset(self): + self.req_to_seq_len.zero_() + + def reset_req(self, req_idx): + self.req_to_seq_len[req_idx] = 0 + + def save_checkpoint(self, req_idx, buffers, slot): + start = req_idx * self.capacity + buffers.kv_buffer[slot].copy_(self.kv_buffer[:, start : start + self.capacity], non_blocking=True) + buffers.seq_lens[slot : slot + 1].copy_(self.req_to_seq_len[req_idx : req_idx + 1], non_blocking=True) + + def restore_checkpoint(self, req_idx, buffers, slot): + start = req_idx * self.capacity + self.kv_buffer[:, start : start + self.capacity].copy_(buffers.kv_buffer[slot], non_blocking=True) + self.req_to_seq_len[req_idx : req_idx + 1].copy_(buffers.seq_lens[slot : slot + 1], non_blocking=True) + # The CPU slot can be recycled as soon as prefix matching returns. + if self.kv_buffer.is_cuda: + torch.cuda.current_stream(self.kv_buffer.device).synchronize() + + def prepare_write_indices( + self, + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_commit_len, + token_num, + alloc_func=torch.empty, + ): + """Map committed window rows to cache slots without changing the source layout. + + token_num is the source hidden row count; b_start_loc indexes that same layout. + b_commit_len counts each request's new prompt rows or accepted verify rows. + + This also updates req_to_seq_len before any layer's KV is written. The caller + must enqueue all layer writes on the same CUDA stream before reading or + checkpointing the window state; this is not a side-effect-free index query. + """ + from lightllm.common.basemodel.triton_kernel.windowed_mtp import prepare_window_write_indices + + return prepare_window_write_indices( + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_commit_len, + self.req_to_seq_len, + self.capacity, + token_num=token_num, + alloc_func=alloc_func, + ) + + def write(self, layer_index, mem_index, kv): + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + destindex_copy_kv(kv, mem_index, self.kv_buffer[layer_index]) + + def pack(self, layer_index, b_req_idx, noise_kv, block_size, alloc_func=torch.empty): + from lightllm.common.basemodel.triton_kernel.windowed_mtp import pack_window_kv + + return pack_window_kv( + self.kv_buffer[layer_index], + noise_kv, + b_req_idx, + self.req_to_seq_len, + self.capacity, + block_size, + alloc_func=alloc_func, + ) + + +class WindowKVPageHelper: + """Append a request's draft window to the existing PD att_state page. + + The page uses global KV heads, just like ordinary PD KV pages, so P and D + may use different TP sizes. The ring's physical slot order is preserved. + """ + + def __init__(self, mem_manager, offset=0): + from lightllm.utils.envs_utils import get_env_start_args + + self.mem_manager = mem_manager + store = mem_manager.windowed_draft_kv + args = get_env_start_args() + tp_world_size = args.tp // args.dp + layers, _, heads, dim = store.kv_buffer.shape + window = store.capacity + self.shape = (window, layers, heads * tp_world_size, dim) + self.kv_offset = (offset + 7) // 8 * 8 + self.kv_nbytes = window * layers * heads * tp_world_size * dim * store.kv_buffer.element_size() + self.seq_len_offset = (self.kv_offset + self.kv_nbytes + 7) // 8 * 8 + self.state_nbytes = self.seq_len_offset + 8 + + def assert_page_size(self): + page = self.mem_manager.kv_move_buffer[0] + page_nbytes = page.numel() * page.element_size() + assert page_nbytes >= self.state_nbytes, ( + f"PD page bytes {page_nbytes} is smaller than attention state bytes {self.state_nbytes}; " + "increase --pd_kv_page_size on both P and D nodes" + ) + + def copy_req_page(self, page_index, req_idx, dp_mems, mode): + from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io + + page = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) + kv_page = ( + page[self.kv_offset : self.kv_offset + self.kv_nbytes] + .view(self.mem_manager.windowed_draft_kv.kv_buffer.dtype) + .view(self.shape) + ) + seq_len_page = page[self.seq_len_offset : self.seq_len_offset + 8].view(torch.int64) + window = self.shape[0] + indexes = torch.arange(req_idx * window, (req_idx + 1) * window, dtype=torch.int64, device=page.device) + for tp_index, mem in enumerate(dp_mems): + store = mem.windowed_draft_kv + page_io( + mem_indexes=indexes, + page_tensor=kv_page, + kv_buffer=store.kv_buffer, + tp_index=tp_index, + tp_world_size=len(dp_mems), + mode=mode, + ) + if mode == "read": + store.req_to_seq_len[req_idx : req_idx + 1].copy_(seq_len_page, non_blocking=True) + elif tp_index == 0: + seq_len_page.copy_(store.req_to_seq_len[req_idx : req_idx + 1], non_blocking=True) diff --git a/lightllm/common/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index 7f2e6ac9b5..cdac9e408d 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -3,6 +3,7 @@ from .linear_att import ReqManagerForMamba from .glm5_next import Glm5NextReqManager from .hybrid_base import HybridAttentionReqManager +from .windowed_mtp import ReqManagerForWindowedMTP from .req_sampling_params import ReqSamplingParamsManager __all__ = [ @@ -10,6 +11,7 @@ "DeepseekV4ReqManager", "HybridAttentionReqManager", "ReqManagerForMamba", + "ReqManagerForWindowedMTP", "Glm5NextReqManager", "ReqSamplingParamsManager", ] diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index c681203e3c..3b4611be92 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -57,15 +57,20 @@ def init_hybrid_attention_state(self, req: "InferReq"): 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 + if self.mem_manager.windowed_draft_kv is not None: + self.mem_manager.windowed_draft_kv.reset_req(req.req_idx) return def create_small_page_cache_manager(self, size: int): - self.small_page_buffers = LinearAttCacheManager(size=size, linear_config=self.linear_config) + self.small_page_buffers = LinearAttCacheManager( + size=size, linear_config=self.linear_config, window_config=self.mem_manager.window_state_config + ) return self.small_page_buffers def save_big_page_states(self, b_req_idx: torch.Tensor, reqs: List["InferReq"], buffer_indexes: List[int]): from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer + buffer_indexes_cpu = buffer_indexes buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) state_cache_manager = self.big_page_buffers copy_linear_att_state_to_kv_buffer( @@ -77,6 +82,12 @@ def save_big_page_states(self, b_req_idx: torch.Tensor, reqs: List["InferReq"], cpu_kv_ssm_state=state_cache_manager.ssm_state_cache.buffer, mtp_step=self.mtp_step, ) + if state_cache_manager.draft_window is not None: + for req, buffer_idx in zip(reqs, buffer_indexes_cpu): + if buffer_idx != -1: + self.mem_manager.windowed_draft_kv.save_checkpoint( + req.req_idx, state_cache_manager.draft_window, buffer_idx + ) return def save_state(self, req: "InferReq", buffer_idx: int, state_cache_manager: LinearAttCacheManager): @@ -88,6 +99,8 @@ def save_state(self, req: "InferReq", buffer_idx: int, state_cache_manager: Line 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) + if state_cache_manager.draft_window is not None: + self.mem_manager.windowed_draft_kv.save_checkpoint(req_idx, state_cache_manager.draft_window, buffer_idx) def get_mamba_cache(self, layer_idx_in_all: int): assert ( @@ -117,6 +130,10 @@ def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheMana 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 state_cache_manager.draft_window is not None: + self.mem_manager.windowed_draft_kv.restore_checkpoint( + req.req_idx, state_cache_manager.draft_window, buffer_idx + ) if self.req_to_mtp_state_index is not None: self.req_to_mtp_state_index[req.req_idx] = 0 return diff --git a/lightllm/common/req_manager/windowed_mtp.py b/lightllm/common/req_manager/windowed_mtp.py new file mode 100644 index 0000000000..fd543f9cce --- /dev/null +++ b/lightllm/common/req_manager/windowed_mtp.py @@ -0,0 +1,21 @@ +from lightllm.common.state_cache_manager import WindowStateCacheManager +from .hybrid_base import HybridAttentionReqManager + + +class ReqManagerForWindowedMTP(HybridAttentionReqManager): + """Full-attention target with request-local draft state in hybrid checkpoints.""" + + def create_small_page_cache_manager(self, size: int): + self.small_page_buffers = WindowStateCacheManager( + size=size, config=self.mem_manager.window_state_config, dtype=self.mem_manager.dtype + ) + return self.small_page_buffers + + def init_hybrid_attention_state(self, req): + self.mem_manager.windowed_draft_kv.reset_req(req.req_idx) + + def save_state(self, req, buffer_idx, state_cache_manager): + self.mem_manager.windowed_draft_kv.save_checkpoint(req.req_idx, state_cache_manager.draft_window, buffer_idx) + + def restore_state(self, req, state_cache_manager, buffer_idx): + self.mem_manager.windowed_draft_kv.restore_checkpoint(req.req_idx, state_cache_manager.draft_window, buffer_idx) diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py index 8e400e2d2c..faeb256963 100644 --- a/lightllm/common/state_cache_manager/__init__.py +++ b/lightllm/common/state_cache_manager/__init__.py @@ -3,20 +3,27 @@ from .layer_cache import LayerCache from .linear_att import LinearAttCacheConfig, LinearAttCacheManager from .glm5_next import Glm5NextCacheConfig +from .windowed_mtp import WindowStateCacheManager def get_hybrid_cache_config(): """Return the model-specific layout used by hybrid CPU/disk cache pages.""" from transformers.configuration_utils import PretrainedConfig from lightllm.utils.envs_utils import get_env_start_args + from .windowed_mtp import WindowedMTPCacheConfig args = get_env_start_args() model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir) model_type = model_cfg["model_type"] if model_type == "deepseek_v4": return DeepseekV4CacheConfig.from_model_config(model_cfg, args) + linear_config = None if model_type in ("glm5_next", "glm5_next_text"): - return Glm5NextCacheConfig.from_model_config(model_cfg, args) - if model_type in ("qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): - return LinearAttCacheConfig.load_from_args() + linear_config = Glm5NextCacheConfig.from_model_config(model_cfg, args) + elif model_type in ("qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): + linear_config = LinearAttCacheConfig.load_from_args() + if args.mtp_windowed: + return WindowedMTPCacheConfig.load_from_args(linear_config) + if linear_config is not None: + return linear_config raise ValueError("No hybrid state-cache layout registered for this model") diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py index 35a37dcd28..014f6d6c5a 100644 --- a/lightllm/common/state_cache_manager/linear_att.py +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -8,6 +8,7 @@ from .base import StateCacheManager from .layer_cache import LayerCache +from .windowed_mtp import WindowStateBuffers, WindowStateCacheConfig logger = init_logger(__name__) @@ -161,9 +162,13 @@ def __init__( size: int, linear_config: LinearAttCacheConfig, keep_num: int = 0, # 用于记录需要保留的缓存数量，用于支持含有 linear_att 的如qwen3.5 模型的cpu cache的碎页处理。 + window_config: WindowStateCacheConfig = None, ): super().__init__(size, keep_num) self.linear_config = linear_config + self.draft_window = ( + WindowStateBuffers(size, window_config, linear_config.full_att_dtype) if window_config is not None else None + ) # init the layer cache self.conv_state_cache = LayerCache( size=self.size, @@ -190,5 +195,7 @@ def get_state_cache(self, buffer_idx: int): def clear_to_init_state(self): self.conv_state_cache.buffer.zero_() self.ssm_state_cache.buffer.zero_() + if self.draft_window is not None: + self.draft_window.clear_to_init_state() super().clear_to_init_state() return diff --git a/lightllm/common/state_cache_manager/windowed_mtp.py b/lightllm/common/state_cache_manager/windowed_mtp.py new file mode 100644 index 0000000000..99f8512691 --- /dev/null +++ b/lightllm/common/state_cache_manager/windowed_mtp.py @@ -0,0 +1,173 @@ +"""Window state configuration and checkpoints shared by hybrid attention caches.""" + +from dataclasses import dataclass + +import torch + +from lightllm.utils.log_utils import init_logger +from lightllm.utils.config_utils import load_config +from .base import StateCacheManager +from .layer_cache import LayerCache + + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class WindowStateCacheConfig: + layers: int + kv_heads: int + head_dim: int + window: int + + @classmethod + def from_model_config(cls, config, window, tp_world_size): + return cls( + layers=config["n_layer"], + kv_heads=config["num_key_value_heads"] // tp_world_size, + head_dim=config.get("head_dim", config["n_embed"] // config["num_attention_heads"]), + window=window, + ) + + def get_state_shape(self): + return (self.window, 2 * self.kv_heads, self.head_dim) + + def get_state_bytes(self, element_size): + return self.layers * self.window * 2 * self.kv_heads * self.head_dim * element_size + 8 + + +def load_window_state_config(args): + config = load_config(args.mtp_draft_model_dir[0]) + config.setdefault("num_key_value_heads", config["num_attention_heads"]) + if config.get("model_type") in ("qwen3_5", "qwen3_5_text") or args.mtp_mode == "dflash2": + config.update(config.get("dflash_config", {})) + return WindowStateCacheConfig.from_model_config(config, args.mtp_window_size, args.tp // args.dp) + + +class WindowStateBuffers: + """Pinned CPU payload; the containing StateCacheManager owns all slot IDs.""" + + def __init__(self, size: int, config: WindowStateCacheConfig, dtype: torch.dtype): + self.kv_buffer = LayerCache( + size=size, + dtype=dtype, + shape=config.get_state_shape(), + layer_num=config.layers, + device="cpu", + size_first=True, + ).buffer + self.seq_lens = torch.zeros(size, dtype=torch.int64, device="cpu", pin_memory=True) + logger.info( + f"draft window checkpoint slots={size}, pinned_bytes={size * config.get_state_bytes(dtype.itemsize)}" + ) + + def get_state_cache(self, buffer_idx: int): + return self.kv_buffer[buffer_idx], self.seq_lens[buffer_idx : buffer_idx + 1] + + def clear_to_init_state(self): + self.kv_buffer.zero_() + self.seq_lens.zero_() + + +class WindowStateCacheManager(StateCacheManager): + def __init__(self, size: int, config: WindowStateCacheConfig, dtype: torch.dtype, keep_num: int = 0): + super().__init__(size, keep_num) + self.draft_window = WindowStateBuffers(size, config, dtype) + + def get_state_cache(self, buffer_idx: int): + return self.draft_window.get_state_cache(buffer_idx) + + def clear_to_init_state(self): + self.draft_window.clear_to_init_state() + super().clear_to_init_state() + + +@dataclass +class WindowedMTPCacheConfig: + """CPU pages contain target KV/state followed by one draft window per TP rank.""" + + window_config: WindowStateCacheConfig + dtype: torch.dtype + tp_world_size: int + full_att_all_num_kv_heads: int + full_att_head_dim: int + full_att_layer_num: int + linear_config: object = None + + @classmethod + def load_from_args(cls, linear_config=None): + from lightllm.utils.config_utils import get_num_key_value_heads, get_head_dim, get_layer_num + from lightllm.utils.envs_utils import get_env_start_args + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + args = get_env_start_args() + return cls( + window_config=load_window_state_config(args), + dtype=get_torch_dtype(args.data_type), + tp_world_size=args.tp // args.dp, + full_att_all_num_kv_heads=get_num_key_value_heads(args.model_dir), + full_att_head_dim=get_head_dim(args.model_dir), + full_att_layer_num=( + linear_config.get_full_att_kv_layer_num_with_draft_model() + if linear_config is not None + else get_layer_num(args.model_dir) + ), + linear_config=linear_config, + ) + + def get_cpu_cache_full_att_bytes(self): + if self.linear_config is not None: + return self.linear_config.get_cpu_cache_full_att_bytes() + from lightllm.utils.envs_utils import get_env_start_args + + args = get_env_start_args() + page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num + assert page_tokens == args.cpu_cache_token_page_size + return ( + page_tokens + * self.full_att_layer_num + * 2 + * self.full_att_all_num_kv_heads + * self.full_att_head_dim + * self.dtype.itemsize + ) + + def get_full_att_kv_view(self, cpu_cache_tensor): + from lightllm.utils.envs_utils import get_env_start_args + + page_tokens = get_env_start_args().cpu_cache_token_page_size + pages = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint8) + return ( + pages[:, : self.get_cpu_cache_full_att_bytes()] + .view(self.dtype) + .view( + pages.shape[0], + self.full_att_layer_num, + page_tokens, + 2 * self.full_att_all_num_kv_heads, + self.full_att_head_dim, + ) + ) + + def get_cpu_cache_window_offset(self): + if self.linear_config is not None: + return self.linear_config.get_cpu_cache_big_page_bytes() + return (self.get_cpu_cache_full_att_bytes() + 15) // 16 * 16 + + def get_cpu_cache_window_rank_bytes(self): + kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 8 + # Align the int64 sequence length and each rank's payload. + return (((kv_bytes + 7) // 8 * 8 + 8) + 15) // 16 * 16 + + def get_cpu_cache_big_page_bytes(self): + return self.get_cpu_cache_window_offset() + self.tp_world_size * self.get_cpu_cache_window_rank_bytes() + + def get_window_views(self, cpu_cache_tensor, tp_rank): + pages = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint8) + start = self.get_cpu_cache_window_offset() + tp_rank * self.get_cpu_cache_window_rank_bytes() + kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 8 + seq_len_offset = start + (kv_bytes + 7) // 8 * 8 + return ( + pages[:, start : start + kv_bytes].view(self.dtype), + pages[:, seq_len_offset : seq_len_offset + 8].view(torch.int64), + ) diff --git a/lightllm/models/qwen3_5_dflash/model.py b/lightllm/models/qwen3_5_dflash/model.py index 422ccf61cf..32fd283537 100644 --- a/lightllm/models/qwen3_5_dflash/model.py +++ b/lightllm/models/qwen3_5_dflash/model.py @@ -29,6 +29,8 @@ def _init_custom(self): self.mask_token_id = self.config["mask_token_id"] def _init_mem_manager(self): + if self.uses_windowed_draft_kv: + return super()._init_mem_manager() target_mem_manager = self.main_model.mem_manager draft_kv_shape = (self.config["num_key_value_heads"], self.config["head_dim"]) target_kv_shape = ( diff --git a/lightllm/models/qwen3_5_dspark/model.py b/lightllm/models/qwen3_5_dspark/model.py index 896949d1cf..0a513f66a7 100644 --- a/lightllm/models/qwen3_5_dspark/model.py +++ b/lightllm/models/qwen3_5_dspark/model.py @@ -26,6 +26,8 @@ def _init_custom(self): self.mask_token_id = self.config["mask_token_id"] def _init_mem_manager(self): + if self.uses_windowed_draft_kv: + return super()._init_mem_manager() target_mem_manager = self.main_model.mem_manager draft_kv_shape = (self.config["num_key_value_heads"], self.config["head_dim"]) target_kv_shape = ( diff --git a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py index 84ee0b4059..08fd0aa494 100644 --- a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py @@ -26,6 +26,11 @@ def context_forward( infer_state: Qwen3DFlashInferStateInfo, layer_weight: Qwen3DFlashTransformerLayerWeight, ) -> torch.Tensor: + cache_kv = self._get_context_kv(input_embdings, infer_state, layer_weight) + self._post_cache_kv(cache_kv, infer_state, layer_weight) + return input_embdings + + def _get_context_kv(self, input_embdings, infer_state, layer_weight): token_num, _ = input_embdings.shape cache_kv = layer_weight.kv_proj.mm(input_embdings, use_custom_tensor_mananger=False) qk_rmsnorm_forward( @@ -41,8 +46,30 @@ def context_forward( infer_state.position_sin, partial_rotary_factor=self.partial_rotary_factor, ) - self._post_cache_kv(cache_kv, infer_state, layer_weight) - return input_embdings + return cache_kv + + def token_attention_forward(self, input_embdings, infer_state, layer_weight): + att_state = infer_state.decode_att_state + if not att_state.backend.model.uses_windowed_draft_kv: + return super().token_attention_forward(input_embdings, infer_state, layer_weight) + + q, noise_kv = self._get_qkv(input_embdings, infer_state, layer_weight) + windowed_kv = att_state.prepare_kv(self.layer_num_, noise_kv, alloc_func=self.alloc_tensor) + output = self._token_attention_kernel(q, infer_state, layer_weight, windowed_kv=windowed_kv) + return self._get_o(output, infer_state, layer_weight) + + def _token_attention_kernel(self, q, infer_state, layer_weight, windowed_kv=None): + # Windowed KV is prepared by the backend; ordinary KV stays in the memory manager. + if windowed_kv is None: + return super()._token_attention_kernel(q, infer_state, layer_weight) + k, v = windowed_kv + output = infer_state.decode_att_state.decode_att( + q=q.view(-1, self.tp_q_head_num_, self.head_dim_), + k=k, + v=v, + alloc_func=self.alloc_tensor, + ) + return output.view(q.shape) def _get_qkv(self, input, infer_state: Qwen3DFlashInferStateInfo, layer_weight: Qwen3DFlashTransformerLayerWeight): q = layer_weight.q_proj.mm(input, use_custom_tensor_mananger=False) diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 9a1811a50a..79fcd98cb3 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -5,7 +5,9 @@ Fp8Fa3AttBackend, ) from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.attention.fa3.windowed_mtp import WindowedMTPAttBackend from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_prefill_params from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.qwen3_dflash.infer_struct import Qwen3DFlashInferStateInfo from lightllm.models.qwen3_dflash.layer_infer.post_layer_infer import Qwen3DFlashPostLayerInfer @@ -13,12 +15,14 @@ from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer from lightllm.models.qwen3_dflash.layer_weights.pre_and_post_layer_weight import Qwen3DFlashPreAndPostLayerWeight from lightllm.models.qwen3_dflash.layer_weights.transformer_layer_weight import Qwen3DFlashTransformerLayerWeight +from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager class Qwen3DFlashModel(LlamaTpPartModel): """Qwen3 DFlash draft model.""" is_mtp_draft_model = True + uses_windowed_draft_kv = False pre_and_post_weight_class = Qwen3DFlashPreAndPostLayerWeight transformer_weight_class = Qwen3DFlashTransformerLayerWeight @@ -41,6 +45,7 @@ def _verify_params(self): assert self.args.mtp_step <= self.config["block_size"] self.config["block_size"] = self.args.mtp_step + self.uses_windowed_draft_kv = getattr(self.args, "mtp_windowed", False) def _init_custom(self): self._cos_cached = self.main_model._cos_cached @@ -52,10 +57,20 @@ def _init_req_manager(self): self.req_manager = self.main_model.req_manager def _init_mem_manager(self): - # Draft KV uses target-owned cache slots. + # The shared manager owns target KV and, when enabled, draft window KV. self.mem_manager = self.main_model.mem_manager + if self.uses_windowed_draft_kv: + self.mem_manager.init_windowed_draft_kv(self.req_manager.max_request_num) def _init_att_backend(self): + if self.uses_windowed_draft_kv: + from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + if flash_attn_with_kvcache is None: + raise NotImplementedError("Windowed MTP requires FA3") + + self.prefill_att_backend = self.decode_att_backend = WindowedMTPAttBackend(model=self) + return super()._init_att_backend() # FA3 is currently the only backend that supports non-causal block decode. # TODO: Remove this restriction after Triton and FlashInfer support non-causal block decode. @@ -85,8 +100,57 @@ def _init_weights(self, start_layer_index=None): for i in range(self.config["n_layer"]) ] + @torch.no_grad() + def commit_window_kv(self, model_input, hidden, b_req_mtp_start_loc=None, accept_len=None): + """Commit accepted target features to the bounded draft KV store.""" + if hidden.numel() == 0: + return + store = self.mem_manager.windowed_draft_kv + if model_input.is_prefill: + b_req_idx = model_input.b_req_idx + b_start_loc = model_input.b_prefill_start_loc + b_ready_cache_len = model_input.b_ready_cache_len + b_commit_len, _, _, _, position_ids = gen_prefill_params( + hidden.shape[0], b_ready_cache_len, model_input.b_seq_len + ) + else: + b_start_loc = b_req_mtp_start_loc + b_req_idx = model_input.b_req_idx.index_select(0, b_start_loc) + b_ready_cache_len = model_input.b_seq_len.index_select(0, b_start_loc) - 1 + b_commit_len = accept_len + position_ids = model_input.b_seq_len - 1 + # Preserve the target's packed prefill or verify rows, including dynamic verify layouts. + mem_index = store.prepare_write_indices( + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_commit_len, + token_num=hidden.shape[0], + alloc_func=self.pre_infer.alloc_tensor, + ) + infer_state = self.infer_state_class() + infer_state.mtp_draft_input_hiddens = hidden + features = self.pre_infer.context_forward(None, infer_state, self.pre_post_weight) + infer_state.position_cos = self._cos_cached.index_select(0, position_ids) + infer_state.position_sin = self._sin_cached.index_select(0, position_ids) + for i, (layer, weight) in enumerate(zip(self.layers_infer, self.trans_layers_weight)): + kv = layer._get_context_kv(features, infer_state, weight) + store.write(i, mem_index, kv) + + def _select_mem_indexes(self, model_input: ModelInput): + if self.uses_windowed_draft_kv and not model_input.is_prefill: + # Windowed draft blocks do not write into the target's request pages. + hold_mem_index = self.mem_manager.HOLD_TOKEN_MEMINDEXES[0] + return g_pin_mem_manager.get_const_gpu_tensor( + key=f"window_mtp_mem_indexes_{hold_mem_index}_{model_input.b_req_idx.dtype}", + shape=model_input.b_req_idx.shape, + fill_value=hold_mem_index, + dtype=model_input.b_req_idx.dtype, + ) + return super()._select_mem_indexes(model_input) + def _decode(self, model_input: ModelInput) -> ModelOutput: - if model_input.mtp_draft_input_hiddens is None: + if self.uses_windowed_draft_kv or model_input.mtp_draft_input_hiddens is None: return super()._decode(model_input) assert model_input.mtp_draft_input_hiddens.shape[0] == model_input.batch_size diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index b03fd77cc2..6286c9269e 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -12,10 +12,15 @@ def __init__(self, layer_num, network_config): self.conv_group_size_ = int(network_config["conv_group_size"]) self.sliding_window_ = int(network_config.get("sliding_window", 0) or 0) - def _token_attention_kernel(self, q, infer_state, layer_weight): - k, v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) - k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) - v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) + def _token_attention_kernel(self, q, infer_state, layer_weight, windowed_kv=None): + if windowed_kv is None: + # Ordinary cache: reshape the shared storage layout to the draft model's KV heads. + k, v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) + v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) + else: + # Windowed cache: history + current block, already packed as [batch, seq, heads, dim]. + k, v = windowed_kv q = q.view(-1, self.tp_q_head_num_, self.head_dim_) use_sliding_window = self.sliding_window_ > 0 # DFlash2 attention is non-causal, so its local window must retain both diff --git a/lightllm/models/qwen3_dflash2/model.py b/lightllm/models/qwen3_dflash2/model.py index 3542b37188..63c3f80268 100644 --- a/lightllm/models/qwen3_dflash2/model.py +++ b/lightllm/models/qwen3_dflash2/model.py @@ -20,10 +20,7 @@ class Qwen3DFlash2Model(Qwen3DFlashModel): def _init_config(self): super()._init_config() - dflash_config = self.config.get("dflash_config", {}) - if not isinstance(dflash_config, dict): - raise ValueError("dflash_config must be an object in the DFlash2 checkpoint config") - self.config.update(dflash_config) + self.config.update(self.config.get("dflash_config", {})) rope_parameters = self.config.get("rope_parameters", {}) if "rope_theta" in rope_parameters and "rope_theta" not in self.config: @@ -54,6 +51,7 @@ def _verify_params(self): physical_block_size = self.args.mtp_step + 1 assert physical_block_size <= self.config["block_size"] self.config["block_size"] = physical_block_size + self.uses_windowed_draft_kv = getattr(self.args, "mtp_windowed", False) def _init_custom(self): # The released Qwen3.8 drafter uses its own Qwen3 rotary layout. @@ -62,6 +60,8 @@ def _init_custom(self): self.mask_token_id = int(self.config["mask_token_id"]) def _init_mem_manager(self): + if self.uses_windowed_draft_kv: + return super()._init_mem_manager() main_mem_manager = self.main_model.mem_manager draft_head_num = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) draft_width = draft_head_num * self.config["head_dim"] diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 02365c3594..a27201c8a5 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -854,6 +854,17 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""Number of additional draft tokens per request. For DSpark and DFlash this value is derived from the draft checkpoint block_size.""", ) + parser.add_argument( + "--mtp_windowed", + action="store_true", + help="""Use sliding-window draft KV storage for DFlash, DSpark or DFlash2.""", + ) + parser.add_argument( + "--mtp_window_size", + type=int, + default=512, + help="""Number of tokens retained in the draft KV cache when mtp_windowed is enabled.""", + ) parser.add_argument( "--mtp_dynamic_verify", action="store_true", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 3e2fb44b5b..97b28aaaf9 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -148,7 +148,7 @@ def _launch_subprocesses(args: StartArgs): args.graph_max_len_in_batch = args.max_req_total_len # mode setting check. - is_hybrid_model = is_hybrid_att_model(args.model_dir) + is_hybrid_model = is_hybrid_att_model(args.model_dir) or args.mtp_windowed if is_hybrid_model: # 混合注意力模型的状态缓存依赖动态缓存和分块预填充，与是否开启输出约束无关。 assert not args.disable_dynamic_prompt_cache, "hybrid attention models require dynamic prompt cache" @@ -217,6 +217,12 @@ def _launch_subprocesses(args: StartArgs): assert args.mtp_draft_model_dir is None assert args.mtp_step == 0 + if args.mtp_windowed: + assert args.mtp_window_size > 0 + assert args.mtp_mode in ("dflash", "dspark", "dflash2"), "windowed draft KV requires dflash, dspark or dflash2" + if args.llm_kv_type not in (None, "None"): + raise ValueError("windowed draft KV currently requires unquantized KV") + # automatically set visual_dp based on visual_tp and the per-node world size. # In visual proxy mode keep the caller-provided visual_dp / visual_tp. node_world_size = args.tp // args.nnodes diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index 71f1c2d5ae..fe2e83811f 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -213,7 +213,7 @@ def init( self.post_init() args = get_env_start_args() - if is_hybrid_att_model(args.model_dir): + if is_hybrid_att_model(args.model_dir) or getattr(args, "mtp_windowed", False): self._fill_hybrid_token_hash() if args.enable_cpu_cache: cpu_cache_hash_list, cpu_cache_page_len_list = self._calcu_hybrid_cpu_cache_page_len_list() diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index baaefae2e3..43412c2ce0 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -215,6 +215,8 @@ class StartArgs: }, ) mtp_draft_model_dir: Optional[List[str]] = field(default=None) + mtp_windowed: bool = field(default=False) + mtp_window_size: int = field(default=512) mtp_step: int = field(default=0) mtp_dynamic_verify: bool = field(default=False) kv_quant_calibration_config_path: Optional[str] = field(default=None) 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 6c621b6adb..1e90e430fc 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -218,7 +218,9 @@ def init_model(self, kvargs): [rank for rank in range(self.global_world_size)], backend="nccl" ) - if self.args.run_mode in ["prefill", "decode"] or self.args.enable_dp_prompt_cache_fetch: + if not self.args.mtp_windowed and ( + self.args.run_mode in ["prefill", "decode"] or self.args.enable_dp_prompt_cache_fetch + ): # 如果存在需要跨进程使用mem manger的特性，则将mem manager写入到 shm中，方便 # 读取 self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager) @@ -251,6 +253,11 @@ def init_model(self, kvargs): self.init_mtp_draft_model(model_kvargs) self.init_spec_engine() + # window PD 需要在 draft 分配独立缓存后再导出。 + if self.args.mtp_windowed and self.is_pd_mode: + self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager) + dist.barrier(group=self.node_nccl_group) + if self.args.enable_cpu_cache: self.multi_level_cache_module = MultiLevelKvCacheModule(self) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py index 9c8347864a..c287d00028 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py @@ -39,6 +39,9 @@ def fill_draft_model_kv_state( # DFlash prefill 直接复用 target prompt 的 token 布局，并注入 target # hidden 初始化唯一 draft model 的 KV。使用浅副本避免在 target 输入上 # 保留 draft 专用状态。 + if getattr(self.backend.draft_models[0], "uses_windowed_draft_kv", False): + self.backend.draft_models[0].commit_window_kv(target_model_input, target_hidden) + return draft_input = copy.copy(target_model_input) draft_input.mtp_draft_input_hiddens = target_hidden self.backend.draft_models[0].forward(draft_input) @@ -78,12 +81,16 @@ def propose_next( accepted_tail_rows = (b_req_mtp_start_loc + accept_len - 1).long() - # target verify 的行布局和 mem_indexes 对应本轮所有被验证 token。 - # 附加 target hidden 后执行一次 draft forward，将这些行提交到 DFlash - # KV cache；浅副本保证 target_model_input 本身保持不变。 - verify_draft_input = copy.copy(target_model_input) - verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden - draft_model.forward(verify_draft_input) + # full 模式通过 draft forward 写入 verify 行对应的 KV；windowed 模式 + # 沿用 verify 行进行投影，仅将接受且仍在窗口内的行写入独立缓存。 + if getattr(draft_model, "uses_windowed_draft_kv", False): + draft_model.commit_window_kv( + target_model_input, target_model_output.mtp_collector.spec_hidden, b_req_mtp_start_loc, accept_len + ) + else: + verify_draft_input = copy.copy(target_model_input) + verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden + draft_model.forward(verify_draft_input) # 每个请求始终展开完整 block，未被本轮 proposal 返回的 block 尾部仍会 # 参与 parallel forward；KV 位置由请求页表和序列位置确定。 diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py index 1d7ca20395..f5c72c793d 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py @@ -40,6 +40,9 @@ def fill_draft_model_kv_state( # DSpark prefill 直接使用 target prompt 的 token 布局和 hidden，将 prompt # KV 写入唯一的 parallel-block draft model。使用浅副本，避免把 draft # 专用 hidden 挂到后续流程仍可能读取的 target ModelInput 上。 + if getattr(self.backend.draft_models[0], "uses_windowed_draft_kv", False): + self.backend.draft_models[0].commit_window_kv(target_model_input, target_hidden) + return draft_input = copy.copy(target_model_input) draft_input.mtp_draft_input_hiddens = target_hidden self.backend.draft_models[0].forward(draft_input) @@ -81,12 +84,16 @@ def propose_next( accepted_tail_rows = (b_req_mtp_start_loc + accept_len - 1).long() - # target verify 的行布局和 mem_indexes 已经对应本轮所有被验证 token。 - # 仅附加 target hidden 后执行一次 draft forward，即可把这些行提交到 - # DSpark KV cache；浅副本保证 target_model_input 本身不被修改。 - verify_draft_input = copy.copy(target_model_input) - verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden - draft_model.forward(verify_draft_input) + # full 模式通过 draft forward 写入 verify 行对应的 KV；windowed 模式 + # 沿用 verify 行进行投影，仅将接受且仍在窗口内的行写入独立缓存。 + if getattr(draft_model, "uses_windowed_draft_kv", False): + draft_model.commit_window_kv( + target_model_input, target_model_output.mtp_collector.spec_hidden, b_req_mtp_start_loc, accept_len + ) + else: + verify_draft_input = copy.copy(target_model_input) + verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden + draft_model.forward(verify_draft_input) # DSpark 每个请求固定展开一个完整 block。block 第一行是 # accepted-tail anchor，其余行使用 mask token，由 parallel backbone diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 2cdbc07dce..174ed0933d 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -3,6 +3,7 @@ from typing import Any, Dict, Optional, List from functools import lru_cache from .envs_utils import get_env_start_args +from lightllm.common.build_utils import repair_config from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -14,6 +15,15 @@ def get_config_json(model_path: str): return json_obj +def load_config(weight_dir_): + """Read model configuration and normalize the common architecture aliases.""" + config = get_config_json(weight_dir_) + repair_config(config, same_names=["num_attention_heads", "n_head"]) + repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(config, same_names=["num_hidden_layers", "n_layer"]) + return config + + def get_generation_config_dict(model_path: str) -> Dict[str, Any]: from transformers import GenerationConfig diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index d3ef067c9b..d82c2d83f3 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -270,6 +270,8 @@ def enable_cpu_cache_numa_interleave() -> bool: @lru_cache(maxsize=None) def get_added_mtp_kv_layer_num() -> int: args = get_env_start_args() + if getattr(args, "mtp_windowed", False): + return 0 mtp_mode = args.mtp_mode if mtp_mode is None: @@ -305,6 +307,8 @@ def get_mtp_weight_layer_num() -> int: return args.mtp_step if mtp_mode == "eagle_no_att": return 1 + if args.mtp_windowed: + return _get_mtp_draft_backbone_layer_num(args.mtp_draft_model_dir[0]) return get_added_mtp_kv_layer_num() diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index b3047b488e..c3533e2c8f 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -62,7 +62,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": args = get_env_start_args() assert args.enable_cpu_cache - is_hybrid_model = is_hybrid_att_model(args.model_dir) + is_hybrid_model = is_hybrid_att_model(args.model_dir) or args.mtp_windowed mem_manager_class = None if is_hybrid_model else select_mem_manager_class() if is_hybrid_model: hybrid_config = get_hybrid_cache_config() diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index 5ec2a9145e..1694dd4ae1 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -1,11 +1,10 @@ import os import gc -import json import torch from contextlib import contextmanager from transformers import AutoModelForCausalLM import argparse -from lightllm.common.build_utils import repair_config +from lightllm.utils.config_utils import load_config from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_mtp_weight_layer_num @@ -74,6 +73,14 @@ def profile_mtp_weight_memory(model): weight_memory_before = torch.cuda.memory_allocated() yield target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before + if model.args.mtp_windowed: + from lightllm.common.state_cache_manager.windowed_mtp import load_window_state_config + + config = load_window_state_config(model.args) + draft_cache_bytes = (model.max_req_num + 1) * config.get_state_bytes(model.data_type.itemsize) + total_gpu_bytes = torch.cuda.get_device_properties(get_current_device_id()).total_memory + # Include the window reservation in the existing per-rank budget reduction. + model.mem_fraction -= draft_cache_bytes / total_gpu_bytes model.mem_fraction = get_mtp_adjusted_mem_fraction( mem_fraction=model.mem_fraction, target_weight_bytes=target_weight_bytes, @@ -82,22 +89,6 @@ def profile_mtp_weight_memory(model): ) -def load_config(weight_dir_): - """ - Load model configuration from the specified directory - Args: - weight_dir_: Path to the weight directory - Returns: - config: Model configuration - """ - with open(os.path.join(weight_dir_, "config.json"), "r") as json_file: - config = json.load(json_file) - repair_config(config, same_names=["num_attention_heads", "n_head"]) - repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(config, same_names=["num_hidden_layers", "n_layer"]) - return config - - def load_model(model_dir, tp_size, data_type): """ Load the model and apply parallelization diff --git a/unit_tests/common/test_windowed_mtp.py b/unit_tests/common/test_windowed_mtp.py new file mode 100644 index 0000000000..734b92772f --- /dev/null +++ b/unit_tests/common/test_windowed_mtp.py @@ -0,0 +1,245 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.kv_cache_mem_manager.windowed_mtp import WindowKVStore, WindowKVPageHelper +from lightllm.common.state_cache_manager.windowed_mtp import WindowedMTPCacheConfig, WindowStateCacheConfig + + +@pytest.fixture +def tensor_cache(): + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + g_cache_manager.cache_env_in() + try: + yield g_cache_manager.empty + finally: + g_cache_manager.cache_env_out() + + +def test_checkpoint_and_cpu_rank_layout(): + from lightllm.common.req_manager.windowed_mtp import ReqManagerForWindowedMTP + + store = WindowKVStore(3, 2, 1, 3, torch.float32, "cpu", 4) + expected = torch.arange(48, dtype=torch.float32).view(2, 4, 2, 3) + store.kv_buffer[:, :4].copy_(expected) + store.req_to_seq_len[0] = 9 + buffers = SimpleNamespace(kv_buffer=torch.zeros(2, 2, 4, 2, 3), seq_lens=torch.zeros(2, dtype=torch.int64)) + manager = ReqManagerForWindowedMTP.__new__(ReqManagerForWindowedMTP) + manager.mem_manager = SimpleNamespace( + windowed_draft_kv=store, big_page_buffers=SimpleNamespace(draft_window=buffers) + ) + manager.save_big_page_states(torch.tensor([0]), [SimpleNamespace(req_idx=0)], [1]) + store.restore_checkpoint(2, buffers, 1) + torch.testing.assert_close(store.kv_buffer[:, 8:12], expected, rtol=0, atol=0) + assert store.req_to_seq_len.tolist() == [9, 0, 9] + + config = WindowStateCacheConfig(2, 1, 3, 4) + target = SimpleNamespace(get_cpu_cache_big_page_bytes=lambda: 32) + layout = WindowedMTPCacheConfig(config, torch.float16, 2, 0, 0, 0, target) + pages = torch.zeros(2, layout.get_cpu_cache_big_page_bytes(), dtype=torch.uint8) + for rank in range(2): + kv, seq_lens = layout.get_window_views(pages, rank) + kv.fill_(rank + 1) + seq_lens.fill_(2 ** 35 + rank) + for rank in range(2): + kv, seq_lens = layout.get_window_views(pages, rank) + assert kv.eq(rank + 1).all() + assert seq_lens.eq(2 ** 35 + rank).all() + assert not pages[:, :32].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_write_indices_preserve_source_layout(tensor_cache): + store = WindowKVStore(3, 1, 1, 3, torch.float32, "cuda", 4) + reqs = torch.tensor([2, 1, 0], dtype=torch.int32, device="cuda") + q_lengths = torch.tensor([7, 0, 2], dtype=torch.int32, device="cuda") + ready = torch.tensor([5, 9, 7], dtype=torch.int32, device="cuda") + starts = torch.tensor([0, 7, 7], dtype=torch.int32, device="cuda") + mem_index = store.prepare_write_indices(reqs, starts, ready, q_lengths, token_num=9, alloc_func=tensor_cache) + # Retain four rows from the long query, skip the empty query, then keep both short-query rows. + assert mem_index.tolist() == [-1, -1, -1, 8, 9, 10, 11, 3, 0] + assert store.req_to_seq_len.tolist() == [9, 9, 12] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("window", [1, 4, 8, 257]) +def test_commit_and_pack(window, tensor_cache): + store = WindowKVStore(3, 2, 1, 3, torch.float32, "cuda", window) + reqs = torch.tensor([0, 2], dtype=torch.int32, device="cuda") + first_len = max(7, window + 3) + hidden = torch.arange((first_len + 3) * 24, device="cuda", dtype=torch.float32).view(-1, 12)[::2, ::2] + starts = torch.tensor([0, first_len], dtype=torch.int32, device="cuda") + first = torch.zeros(2, dtype=torch.int32, device="cuda") + lengths = torch.tensor([first_len, 3], dtype=torch.int32, device="cuda") + mem_index = store.prepare_write_indices( + reqs, starts, first, lengths, token_num=hidden.shape[0], alloc_func=tensor_cache + ) + for layer in range(2): + store.write(layer, mem_index, hidden.view(-1, 2, 3) + layer) + + accepted = torch.arange(5 * 6, device="cuda", dtype=torch.float32).view(5, 6) + 1000 + starts = torch.tensor([0, 3], dtype=torch.int32, device="cuda") + first = torch.tensor([first_len, 3], dtype=torch.int32, device="cuda") + lengths = torch.tensor([2, 1], dtype=torch.int32, device="cuda") + mem_index = store.prepare_write_indices( + reqs, starts, first, lengths, token_num=accepted.shape[0], alloc_func=tensor_cache + ) + assert mem_index.shape == (5,) + assert mem_index[[2, 4]].eq(-1).all() + for layer in range(2): + store.write(layer, mem_index, accepted.view(-1, 2, 3) + layer) + history = [torch.cat((hidden[:first_len], accepted[:2])), torch.cat((hidden[first_len:], accepted[3:4]))] + noise = torch.randn(4, 2, 6, device="cuda")[:, :, ::2] + for layer in range(2): + packed = store.pack(layer, reqs, noise, 2, alloc_func=tensor_cache) + for batch, values in enumerate(history): + retained = values[-window:].view(-1, 2, 3) + layer + expected = torch.cat((retained, noise[batch * 2 : (batch + 1) * 2])) + torch.testing.assert_close(packed[batch, : len(expected)], expected, rtol=0, atol=0) + assert store.req_to_seq_len.tolist() == [first_len + 2, 0, 4] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("sliding_window", [0, 4]) +def test_fa3_graph_refreshes_request_state(tensor_cache, sliding_window): + from lightllm.common.basemodel.attention.fa3.windowed_mtp import WindowedMTPAttBackend + from lightllm.common.basemodel.cuda_graph import CudaGraph + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer + from lightllm.models.qwen3_dflash2.layer_infer.transformer_layer_infer import Qwen3DFlash2TransformerLayerInfer + from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + if flash_attn_with_kvcache is None: + pytest.skip("FA3 required") + torch.manual_seed(0) + window, block_size, head_dim = 4, 2, 64 + store = WindowKVStore(3, 1, 1, head_dim, torch.bfloat16, "cuda", window) + store.kv_buffer.normal_() + store.req_to_seq_len.copy_(torch.tensor([7, 0, 2], device="cuda")) + backend = SimpleNamespace( + model=SimpleNamespace(block_size=block_size, draft_layer_start=0, uses_windowed_draft_kv=True) + ) + + def make_state(req_ids): + state = InferStateInfo() + state.is_prefill = False + state.input_ids = torch.zeros(4, dtype=torch.int64, device="cuda") + state.b_req_idx = torch.tensor(req_ids, dtype=torch.int32, device="cuda").repeat_interleave(block_size) + state.mem_manager = SimpleNamespace(windowed_draft_kv=store) + state.query = torch.randn(4, 2, head_dim, dtype=torch.bfloat16, device="cuda") + state.noise_kv = torch.randn(4, 2, head_dim, dtype=torch.bfloat16, device="cuda") + state.decode_att_state = WindowedMTPAttBackend.create_att_decode_state(backend, state) + state.init_att_state() + return state + + layer_class = Qwen3DFlash2TransformerLayerInfer if sliding_window else Qwen3DFlashTransformerLayerInfer + layer = layer_class.__new__(layer_class) + layer.sliding_window_ = sliding_window + layer.layer_num_, layer.tp_q_head_num_, layer.tp_k_head_num_, layer.head_dim_ = 0, 2, 1, head_dim + layer._get_qkv = lambda input, state, weight: (state.query, state.noise_kv) + layer._get_o = lambda output, state, weight: output.flatten(1) + weight = SimpleNamespace() + captured = make_state([0, 2]) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + layer.token_attention_forward(captured.query.flatten(1), captured, weight) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = layer.token_attention_forward(captured.query.flatten(1), captured, weight) + runner = CudaGraph.__new__(CudaGraph) + runner.graph = {4: (graph, captured, output)} + + for req_ids, lengths in [([2, 1], [7, 6, 0]), ([1, 0], [2, 1, 0])]: + store.kv_buffer.normal_() + store.req_to_seq_len.copy_(torch.tensor(lengths, device="cuda")) + runtime = make_state(req_ids) + att_state = runtime.decode_att_state + expected_lengths = [min(lengths[req], window) + block_size for req in req_ids] + assert att_state.cache_seqlens.tolist() == expected_lengths + expected = [] + for batch, req in enumerate(req_ids): + positions = torch.arange(max(0, lengths[req] - window), lengths[req], device="cuda") + history = store.kv_buffer[0, req * window + positions % window] + noise = runtime.noise_kv[batch * block_size : (batch + 1) * block_size] + kv = torch.cat((history, noise)).float() + q = runtime.query[batch * block_size : (batch + 1) * block_size].float() + scores = torch.einsum("qhd,kd->hqk", q, kv[:, 0]) * head_dim ** -0.5 + if sliding_window: + q_pos = torch.arange(block_size, device="cuda") + len(history) + k_pos = torch.arange(len(kv), device="cuda") + mask = (q_pos[:, None] - k_pos[None, :]).abs() >= sliding_window + scores.masked_fill_(mask[None], float("-inf")) + expected.append(torch.einsum("hqk,kd->qhd", scores.softmax(-1), kv[:, 1])) + expected = torch.cat(expected).flatten(1).to(runtime.query.dtype) + actual = runner._replay(runtime) + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (1, 2), (2, 1)]) +def test_pd_page_round_trip(monkeypatch, prefill_tp, decode_tp): + def make_mems(tp): + monkeypatch.setattr("lightllm.utils.envs_utils.get_env_start_args", lambda: SimpleNamespace(tp=tp, dp=1)) + mems = [ + SimpleNamespace( + windowed_draft_kv=WindowKVStore(3, 2, 2 // tp, 8, torch.float32, "cuda", 4), + kv_move_buffer=torch.zeros(1, 2048, dtype=torch.uint8, device="cuda"), + ) + for _ in range(tp) + ] + helper = WindowKVPageHelper(mems[0], offset=16) + helper.assert_page_size() + return mems, helper + + source, writer = make_mems(prefill_tp) + destination, reader = make_mems(decode_tp) + global_kv = torch.randn(2, 4, 4, 8, device="cuda") + + def shard_heads(rank, tp): + start = rank * (2 // tp) + return list(range(start, start + 2 // tp)) + list(range(2 + start, 2 + start + 2 // tp)) + + for rank, mem in enumerate(source): + mem.windowed_draft_kv.kv_buffer[:, :4].copy_(global_kv[:, :, shard_heads(rank, prefill_tp)]) + mem.windowed_draft_kv.req_to_seq_len[0] = 7 + source[0].kv_move_buffer[:, :16].fill_(42) + writer.copy_req_page(0, 0, source, "write") + destination[0].kv_move_buffer.copy_(source[0].kv_move_buffer) + for mem in destination: + mem.windowed_draft_kv.kv_buffer.zero_() + reader.copy_req_page(0, 2, destination, "read") + for rank, mem in enumerate(destination): + store = mem.windowed_draft_kv + torch.testing.assert_close(store.kv_buffer[:, 8:12], global_kv[:, :, shard_heads(rank, decode_tp)]) + assert store.req_to_seq_len.tolist() == [0, 0, 7] + assert not store.kv_buffer[:, :8].any() + assert destination[0].kv_move_buffer[:, :16].eq(42).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_tail_checkpoint_copy_preserves_stream_order(): + from lightllm.common.kv_cache_mem_manager.operator.windowed_mtp import WindowedMTPMemOperator + from lightllm.common.state_cache_manager.windowed_mtp import WindowStateBuffers + + config = WindowStateCacheConfig(2, 1, 8, 4) + source = SimpleNamespace(draft_window=WindowStateBuffers(2, config, torch.bfloat16)) + destination = SimpleNamespace(draft_window=WindowStateBuffers(2, config, torch.bfloat16)) + operator = WindowedMTPMemOperator.__new__(WindowedMTPMemOperator) + operator.mem_manager = SimpleNamespace(big_page_buffers=destination) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + torch.cuda._sleep(1_000_000) + expected = torch.randn(source.draft_window.kv_buffer[1].shape, dtype=torch.bfloat16, device="cuda") + seq_len = torch.tensor([2 ** 35 + 7], dtype=torch.int64, device="cuda") + source.draft_window.kv_buffer[1].copy_(expected, non_blocking=True) + source.draft_window.seq_lens[1:2].copy_(seq_len, non_blocking=True) + operator._copy_tail_state(source, source_id=1, destination_id=0) + stream.synchronize() + torch.testing.assert_close(destination.draft_window.kv_buffer[0], expected.cpu(), rtol=0, atol=0) + assert destination.draft_window.seq_lens.tolist() == [2 ** 35 + 7, 0] + assert not destination.draft_window.kv_buffer[1].any()