From 92a4a45a158b10d773b8987e7548b87884153878 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 22 Sep 2026 13:17:41 +0800 Subject: [PATCH 1/2] perf: batch KV page allocation and fuse token and KV index gathering --- docs/kv_cache_page_size.md | 45 -- lightllm/common/basemodel/attention/fa3/fp.py | 1 - .../common/basemodel/attention/fa3/mla.py | 2 +- lightllm/common/basemodel/basemodel.py | 215 +++---- lightllm/common/basemodel/batch_objs.py | 5 +- lightllm/common/basemodel/cuda_graph.py | 19 +- .../common/basemodel/prefill_cuda_graph.py | 3 - .../triton_kernel/copy_kv_index_to_req.py | 275 +++------ .../triton_kernel/dynamic_mtp_utils.py | 88 +-- .../basemodel/triton_kernel/fa3_utils.py | 33 -- .../triton_kernel/gather_token_id.py | 322 ---------- lightllm/common/infer_utils.py | 15 - .../common/kv_cache_mem_manager/allocator.py | 143 ++--- .../deepseek2_mem_manager.py | 4 +- .../kv_cache_mem_manager/mem_manager.py | 24 +- .../ppl_int4kv_mem_manager.py | 4 - .../ppl_int8kv_mem_manager.py | 4 - .../qwen3next_mem_manager.py | 9 - lightllm/common/req_manager/base.py | 214 +++++-- .../common/req_manager/req_sampling_params.py | 13 +- lightllm/models/qwen3_dflash/model.py | 3 +- lightllm/server/core/objs/req.py | 1 - lightllm/server/httpserver/manager.py | 2 +- lightllm/server/router/batch.py | 3 +- lightllm/server/router/manager.py | 3 +- .../server/router/model_infer/infer_batch.py | 370 ++++-------- .../model_infer/mode_backend/base_backend.py | 192 +++--- .../mode_backend/diverse_backend/impl.py | 34 +- .../dp_backend/dp_shared_kv_trans.py | 15 +- .../mode_backend/generic_pre_process.py | 61 +- .../mode_backend/multi_level_kv_cache.py | 14 +- .../pd/decode_node_impl/decode_impl.py | 24 +- .../pd/prefill_node_impl/prefill_impl.py | 12 +- .../mtp_speculative/proposers/base.py | 2 +- .../server/router/req_queue/base_queue.py | 3 +- .../server/router/req_queue/dp_base_queue.py | 1 - lightllm/utils/custom_kernel_utis.py | 22 - skills/test_model/SKILL.md | 40 ++ ...t_kv_cache_calib_per_head_qwen2.5_14b.json | 499 +++++++++++++++- .../static_inference/static_benchmark.py | 20 +- .../basemodel/attention/test_fp8_page_size.py | 45 -- .../common/basemodel/test_model_output.py | 2 +- .../common/basemodel/test_overlap_utils.py | 25 +- .../gqa/test_flash_decoding_autotune.py | 1 + .../test_mtp_fused_recurrent_autotune.py | 2 + .../test_mtp_fused_recurrent_equiv.py | 63 +- .../basemodel/triton_kernel/test_fa3_utils.py | 14 +- .../triton_kernel/test_repack_kv_index.py | 31 - .../triton_kernel/test_routing_capture.py | 33 -- .../test_select_kv_index_from_req.py | 75 --- .../triton_kernel/test_sp_pad_kernel.py | 9 +- unit_tests/common/test_paged_mem_manager.py | 32 - .../test_pinned_memory_non_blocking_copy.py | 39 -- .../test_quantized_paged_mem_manager.py | 33 -- unit_tests/common/test_req_manager_page.py | 549 ------------------ .../models/test_qwen3_dspark_model_output.py | 11 +- unit_tests/server/core/objs/test_req.py | 14 +- .../test_running_request_lifecycle.py | 9 - .../router/dynamic_prompt/test_radix_cache.py | 14 +- .../mode_backend/test_diverse_backend.py | 74 --- .../mode_backend/test_generic_pre_process.py | 31 - .../mode_backend/test_multi_level_kv_cache.py | 29 +- .../test_paged_side_allocations.py | 145 ----- .../mode_backend/test_pd_dynamic_split.py | 8 +- .../mode_backend/test_prompt_cache_preload.py | 52 -- .../mode_backend/test_rl_backend_ops.py | 33 -- .../structured_output/test_state.py | 1 + .../model_infer/test_multimodal_chunking.py | 29 +- unit_tests/server/test_dp_request_capacity.py | 6 +- unit_tests/server/test_pd_start_args.py | 75 --- unit_tests/test_static_benchmark_paged_kv.py | 153 ----- .../utils/test_config_utils_fixed_kv.py | 33 -- unit_tests/utils/test_envs_utils.py | 12 +- unit_tests/utils/test_sgl_utils.py | 1 - unit_tests/utils/test_speculative_utils.py | 2 +- 75 files changed, 1365 insertions(+), 3079 deletions(-) delete mode 100644 docs/kv_cache_page_size.md delete mode 100644 lightllm/common/infer_utils.py delete mode 100644 unit_tests/common/basemodel/attention/test_fp8_page_size.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py delete mode 100644 unit_tests/common/test_paged_mem_manager.py delete mode 100644 unit_tests/common/test_pinned_memory_non_blocking_copy.py delete mode 100644 unit_tests/common/test_quantized_paged_mem_manager.py delete mode 100644 unit_tests/common/test_req_manager_page.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_diverse_backend.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py delete mode 100644 unit_tests/server/test_pd_start_args.py delete mode 100644 unit_tests/test_static_benchmark_paged_kv.py delete mode 100644 unit_tests/utils/test_config_utils_fixed_kv.py diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md deleted file mode 100644 index 76e99a88d7..0000000000 --- a/docs/kv_cache_page_size.md +++ /dev/null @@ -1,45 +0,0 @@ -# KV Cache 多 Token 页设计 - -## 目标与不变量 - -启动参数 `--page_size N` 控制注意力 KV Cache 的物理页大小，默认值为 `1`。`N > 1` 时保持 token 级的 -KV 存储和请求表，但资源的申请、缓存和释放以完整物理页为单位。 - -实现依赖以下不变量： - -1. `cur_kv_len` 是请求已经写入有效 KV 的逻辑 token 数；`hold_kv_len` 是请求拥有的物理容量，始终满足 - `cur_kv_len <= hold_kv_len` 且 `hold_kv_len % page_size == 0`。 -2. 一个物理页内的 KV 槽连续，页首索引可被 `page_size` 整除。请求表保存已持有页的全部 token 索引， - 包括尚未使用的尾部槽位。 -3. `req_to_token_indexs` 保存请求拥有的完整页；创建 `InferState` 时根据请求索引和序列位置直接聚合本轮 - 真实参与计算的 token，`ModelInput` 不携带物理 KV 索引，`InferState.mem_index` 不包含预留尾部。 -4. Radix Cache 只插入、拆分、命中和淘汰完整页；不足一页的请求尾部在请求结束或暂停时整页回收。 -5. `page_size=1` 与多 token 页使用相同的调度期预留和模型执行期索引选择路径。 - -页容量计算和整页申请由调度层在请求获准进入 Prefill/Decode batch 时完成；输入构造层只构造本轮输入， -真实推理索引由模型执行层从请求表选取。 -`ReqManager` 只保留请求 ID、请求表及其原有生命周期职责，不提供 page allocator 接口。 - -## 生命周期 - -- Prefill：请求通过调度后，按目标 KV 长度将容量向上对齐，新页一次申请并完整写入请求表。模型执行时再选择 - `[cur_kv_len, target_kv_len)` 对应的真实索引。 -- Decode：请求通过调度后，若尾页仍有预留槽位则不再访问 allocator；容量不足时先补齐完整页再执行。 -- Prefix Cache 命中：命中长度天然是整页，`cur_kv_len` 与 `hold_kv_len` 同时初始化为命中长度。 -- 完成/暂停：完整逻辑页可进入 Radix Cache；重复页和未完成尾页按物理页展开后释放。 -- Attention：FA3/FlashInfer 页表每项由物理 token 页首索引除以 `page_size` 得到，KV buffer 视为 - `[page_count, page_size, ...]`。 - -## 当前兼容范围 - -当前支持普通与分块 Prefill、Decode、动态 Prompt Cache、FA3/FlashInfer（含 MLA Decode）、MTP、 -PD 分离、CPU KV Cache、DP Prompt Cache 拉取、DP Prefill Balance 和混合线性注意力。 -diverse mode 仍只支持 `page_size=1`，启动阶段会对其他取值明确报错。 - -## 边界处理 - -- KV 总容量和请求表宽度分别向下、向上对齐，防止物理页越界。 -- CUDA Graph 的 HOLD request 映射到额外保留的一整个物理页；请求表按 - `[HOLD, HOLD+1, ..., HOLD+page_size-1]` 循环填充，而不是重复同一个 token 索引。 -- Radix 子节点用“首个完整 token 页”作为键，避免不同序列仅首 token 相同造成页级分支冲突。 -- 非法的 `page_size < 1` 以及尚未支持的功能组合在模型加载前失败。 diff --git a/lightllm/common/basemodel/attention/fa3/fp.py b/lightllm/common/basemodel/attention/fa3/fp.py index 6fcee4312a..70294c5e6e 100644 --- a/lightllm/common/basemodel/attention/fa3/fp.py +++ b/lightllm/common/basemodel/attention/fa3/fp.py @@ -3,7 +3,6 @@ import torch import triton from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, TYPE_CHECKING from lightllm.utils.sgl_utils import flash_attn_with_kvcache, flash_attn_with_kvcache_autotune from lightllm.common.basemodel.triton_kernel.fa3_utils import ( build_dynamic_spec_fa3_decode_params, diff --git a/lightllm/common/basemodel/attention/fa3/mla.py b/lightllm/common/basemodel/attention/fa3/mla.py index 244df6be12..e4ce1f299c 100644 --- a/lightllm/common/basemodel/attention/fa3/mla.py +++ b/lightllm/common/basemodel/attention/fa3/mla.py @@ -3,7 +3,7 @@ import torch import triton from ..base_att import BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Optional, TYPE_CHECKING, Tuple +from typing import Tuple from lightllm.utils.sgl_utils import flash_attn_with_kvcache from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params, page_table_copy from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 34a1c9c71c..a8ae1044a6 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -17,15 +17,11 @@ 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.basemodel.triton_kernel.copy_kv_index_to_req import ( - select_kv_index_from_req, - select_kv_index_from_req_prefill, -) +from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import build_kv_indexes_and_input_ids from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager from lightllm.common.basemodel.cuda_graph import CudaGraph from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph 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.dist_utils import get_dp_world_size from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory @@ -329,24 +325,17 @@ def forward(self, model_input: ModelInput): else: return self._decode(model_input) - def _select_mem_indexes(self, model_input: ModelInput): - if model_input.is_prefill: - return select_kv_index_from_req_prefill( - req_to_token_indexs=self.req_manager.req_to_token_indexs, - b_req_idx=model_input.b_req_idx, - b_seq_len=model_input.b_seq_len, - b_ready_cache_len=model_input.b_ready_cache_len, - b_start_loc=model_input.b_prefill_start_loc, - max_q_seq_len=model_input.max_q_seq_len, - token_num=model_input.input_ids.shape[0], - ) - return select_kv_index_from_req( - req_to_token_indexs=self.req_manager.req_to_token_indexs, - b_req_idx=model_input.b_req_idx, - b_seq_len=model_input.b_seq_len, - ) - def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0): + need_build_input_ids = model_input.input_ids is None or ( + model_input.is_prefill and self.args.enable_prefill_decode_mixed + ) + model_input.mem_indexes, model_input.input_ids = build_kv_indexes_and_input_ids( + model_input, + self.req_manager.req_to_token_indexs, + next_token_ids=self.req_manager.req_sampling_params_manager.req_to_next_token_ids + if need_build_input_ids + else None, + ) infer_state = self.infer_state_class() infer_state.hidden_collector = self.hidden_collector_prototype.new_instance() infer_state.input_ids = model_input.input_ids @@ -377,7 +366,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.mem_manager = self.mem_manager infer_state.req_manager = self.req_manager - infer_state.mem_index = self._select_mem_indexes(model_input) + infer_state.mem_index = model_input.mem_indexes infer_state.microbatch_index = microbatch_index infer_state.dist_group = dist_group_manager.get_group(microbatch_index) @@ -399,6 +388,36 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) return infer_state + def _create_prefill_inferstate(self, model_input: ModelInput, handle_token_num: int, microbatch_index: int = 0): + padded_input = self._create_padded_prefill_model_input(model_input, handle_token_num) + infer_state = self._create_inferstate(padded_input, microbatch_index) + if padded_input is not model_input: + token_num = model_input.input_ids.numel() + # Expose the real-token views for prompt logprobs and subsequent draft inputs. + model_input.mem_indexes = infer_state.mem_index[:token_num] + model_input.input_ids = infer_state.input_ids[:token_num] + infer_state.init_some_extra_state(self) + infer_state.init_att_state() + return infer_state + + def _create_decode_inferstate( + self, + model_input: ModelInput, + batch_size: int, + need_capture: bool, + microbatch_index: int = 0, + *, + use_cuda_graph: bool = False, + ): + model_input = self._create_padded_decode_model_input(model_input, batch_size) + infer_state = self._create_inferstate(model_input, microbatch_index) + # Attention backends need both graph flags before initializing their state. + infer_state.is_cuda_graph_capture = need_capture + infer_state.is_cuda_graph_replay = use_cuda_graph and not need_capture + infer_state.init_some_extra_state(self) + infer_state.init_att_state() + return infer_state + def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_size: int): if model_input.batch_size == new_batch_size: return model_input @@ -410,7 +429,8 @@ def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_s new_model_input.batch_size = new_batch_size new_model_input.total_token_num += padded_batch_size * 2 new_model_input.max_kv_seq_len = max(2, model_input.max_kv_seq_len) - new_model_input.input_ids = F.pad(new_model_input.input_ids, (0, padded_batch_size), mode="constant", value=1) + if new_model_input.input_ids is not None: + new_model_input.input_ids = F.pad(new_model_input.input_ids, (0, padded_batch_size), value=1) new_model_input.b_req_idx = F.pad( new_model_input.b_req_idx, (0, padded_batch_size), mode="constant", value=self.req_manager.HOLD_REQUEST_ID ) @@ -466,14 +486,12 @@ def _create_padded_prefill_model_input(self, model_input: ModelInput, new_handle new_model_input.b_seq_len = F.pad(new_model_input.b_seq_len, (0, 1), mode="constant", value=padded_token_num) new_model_input.b_ready_cache_len = F.pad(new_model_input.b_ready_cache_len, (0, 1), mode="constant", value=0) new_model_input.b_is_decode_req = F.pad(new_model_input.b_is_decode_req, (0, 1), mode="constant", value=False) - b_q_seq_len = new_model_input.b_seq_len - new_model_input.b_ready_cache_len - new_model_input.b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len + new_model_input.b_prefill_start_loc = F.pad( + model_input.b_prefill_start_loc.to(torch.int32), (0, 1), value=handle_token_num + ) # 构建新的list, 使用 append 可能会让外面使用的数组引用发生变化，导致错误。 - new_model_input.b_prefill_has_output_cpu = [e for e in new_model_input.b_prefill_has_output_cpu] + [False] - - new_model_input.multimodal_params = [e for e in new_model_input.multimodal_params] + [ - {"images": [], "audios": []} - ] + new_model_input.b_prefill_has_output_cpu = list(model_input.b_prefill_has_output_cpu) + [False] + new_model_input.multimodal_params = list(model_input.multimodal_params) + [{"images": [], "audios": []}] # 特殊模型，特殊模式的特殊变量的特殊 padding if new_model_input.mtp_draft_input_hiddens is not None: @@ -530,16 +548,6 @@ def _prefill( self, model_input: ModelInput, ): - if self.args.enable_prefill_decode_mixed and model_input.input_ids.shape[0] > 0: - gather_token_prefill_decode_mixed( - input_ids=model_input.input_ids, - req_to_next_token_ids=self.req_manager.req_sampling_params_manager.req_to_next_token_ids, - b_req_idx=model_input.b_req_idx, - b_mtp_index=model_input.b_mtp_index, - b_is_decode_req=model_input.b_is_decode_req, - b_prefill_start_loc=model_input.b_prefill_start_loc, - ) - origin_handle_token_num = model_input.input_ids.shape[0] origin_batch_size = model_input.batch_size @@ -554,16 +562,7 @@ def _prefill( handle_token_num=infer_handle_token_num ) - model_input = self._create_padded_prefill_model_input( - model_input=model_input, new_handle_token_num=infer_handle_token_num - ) - - infer_state = self._create_inferstate(model_input) - prefill_mem_indexes_ready_event = torch.cuda.Event() - prefill_mem_indexes_ready_event.record() - - infer_state.init_some_extra_state(self) - infer_state.init_att_state() + infer_state = self._create_prefill_inferstate(model_input, infer_handle_token_num) model_output = self._context_forward(infer_state=infer_state) model_output = self._create_unpad_prefill_model_output( @@ -571,29 +570,12 @@ def _prefill( origin_handle_token_num=origin_handle_token_num, origin_batch_size=origin_batch_size, ) - model_output.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event return model_output def _decode( self, model_input: ModelInput, ) -> ModelOutput: - if model_input.input_ids is None: - if model_input.batch_size > 0: - model_input.input_ids = gather_token( - req_to_next_token_ids=(self.req_manager.req_sampling_params_manager.req_to_next_token_ids), - b_req_idx=model_input.b_req_idx, - b_mtp_index=model_input.b_mtp_index, - ) - else: - # 空 DP rank 不启动 gather kernel，但仍将 input_ids 规范化为 - # CUDA 空 tensor；后续内部 padding 会为 dummy request 填入 token id 1。 - model_input.input_ids = torch.empty( - (0,), - dtype=torch.int64, - device=model_input.b_req_idx.device, - ) - origin_batch_size = model_input.batch_size # 空 DP rank 也需要完整的 dummy verify group；TP/SP 切分不能破坏 MTP 分组。 infer_batch_size = self._align_decode_batch_size(max(1, origin_batch_size)) @@ -610,14 +592,9 @@ def _decode( infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size) need_capture = self.graph.need_capture(infer_batch_size) - model_input = self._create_padded_decode_model_input(model_input=model_input, new_batch_size=infer_batch_size) - infer_state = self._create_inferstate(model_input) - # attention backend 会根据执行方式准备 capture、replay 或 eager 状态， - # 因此必须在 init_att_state 之前完成赋值。 - infer_state.is_cuda_graph_capture = need_capture - infer_state.is_cuda_graph_replay = use_cuda_graph and not need_capture - infer_state.init_some_extra_state(self) - infer_state.init_att_state() + infer_state = self._create_decode_inferstate( + model_input, infer_batch_size, need_capture, use_cuda_graph=use_cuda_graph + ) if use_cuda_graph: if need_capture: @@ -631,7 +608,6 @@ def _decode( @final def _context_forward(self, infer_state: InferStateInfo): - input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" @@ -729,15 +705,6 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod for model_input in (model_input0, model_input1): model_input.to_cuda() - if self.args.enable_prefill_decode_mixed and model_input.input_ids.shape[0] > 0: - gather_token_prefill_decode_mixed( - input_ids=model_input.input_ids, - req_to_next_token_ids=(self.req_manager.req_sampling_params_manager.req_to_next_token_ids), - b_req_idx=model_input.b_req_idx, - b_mtp_index=model_input.b_mtp_index, - b_is_decode_req=model_input.b_is_decode_req, - b_prefill_start_loc=model_input.b_prefill_start_loc, - ) return self._microbatch_overlap_prefill_cuda(model_input0, model_input1) def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input1: ModelInput): @@ -755,24 +722,8 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input origin_batch_size0 = model_input0.batch_size origin_batch_size1 = model_input1.batch_size - model_input0 = self._create_padded_prefill_model_input( - model_input=model_input0, new_handle_token_num=infer_handle_token_num0 - ) - model_input1 = self._create_padded_prefill_model_input( - model_input=model_input1, new_handle_token_num=infer_handle_token_num1 - ) - - infer_state0 = self._create_inferstate(model_input0, 0) - infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() - - infer_state1 = self._create_inferstate(model_input1, 1) - infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() - - prefill_mem_indexes_ready_event = torch.cuda.Event() - prefill_mem_indexes_ready_event.record() - + infer_state0 = self._create_prefill_inferstate(model_input0, infer_handle_token_num0, 0) + infer_state1 = self._create_prefill_inferstate(model_input1, infer_handle_token_num1, 1) model_output0, model_output1 = self._overlap_tpsp_context_forward(infer_state0, infer_state1=infer_state1) model_output0 = self._create_unpad_prefill_model_output( @@ -788,8 +739,6 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input # 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候 # 该调用没有实际意义 dist_group_manager.clear_deepep_buffer() - model_output0.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event - model_output1.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event return model_output0, model_output1 @torch.no_grad() @@ -798,19 +747,6 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode for model_input in (model_input0, model_input1): model_input.to_cuda() - if model_input.input_ids is None: - if model_input.batch_size > 0: - model_input.input_ids = gather_token( - req_to_next_token_ids=(self.req_manager.req_sampling_params_manager.req_to_next_token_ids), - b_req_idx=model_input.b_req_idx, - b_mtp_index=model_input.b_mtp_index, - ) - else: - model_input.input_ids = torch.empty( - (0,), - dtype=torch.int64, - device=model_input.b_req_idx.device, - ) return self._microbatch_overlap_decode_cuda(model_input0, model_input1) def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1: ModelInput): @@ -821,23 +757,20 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 max_len_in_batch = max(2, model_input0.max_kv_seq_len, model_input1.max_kv_seq_len) infer_batch_size = self._align_decode_batch_size(max(1, origin_batch_size0, origin_batch_size1)) - if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch): + use_cuda_graph = self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch) + need_capture = False + if use_cuda_graph: infer_batch_size = self.graph.find_closest_graph_batch_size(infer_batch_size) need_capture = self.graph.need_capture(infer_batch_size) - padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) - padded_model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) - infer_state0 = self._create_inferstate(padded_model_input0, 0) - infer_state0.is_cuda_graph_capture = need_capture - infer_state0.is_cuda_graph_replay = not need_capture - infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() - - infer_state1 = self._create_inferstate(padded_model_input1, 1) - infer_state1.is_cuda_graph_capture = need_capture - infer_state1.is_cuda_graph_replay = not need_capture - infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() + infer_state0 = self._create_decode_inferstate( + model_input0, infer_batch_size, need_capture, 0, use_cuda_graph=use_cuda_graph + ) + infer_state1 = self._create_decode_inferstate( + model_input1, infer_batch_size, need_capture, 1, use_cuda_graph=use_cuda_graph + ) + + if use_cuda_graph: if need_capture: model_output0, model_output1 = self.graph.capture_decode( self._overlap_tpsp_token_forward, @@ -850,23 +783,11 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 infer_state1=infer_state1, ) - model_output0 = self._create_unpad_decode_model_output(model_output0, origin_batch_size=origin_batch_size0) - model_output1 = self._create_unpad_decode_model_output(model_output1, origin_batch_size=origin_batch_size1) else: - model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) - model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) - infer_state0 = self._create_inferstate(model_input0, 0) - infer_state0.init_some_extra_state(self) - infer_state0.init_att_state() - - infer_state1 = self._create_inferstate(model_input1, 1) - infer_state1.init_some_extra_state(self) - infer_state1.init_att_state() - model_output0, model_output1 = self._overlap_tpsp_token_forward(infer_state0, infer_state1=infer_state1) - model_output0 = self._create_unpad_decode_model_output(model_output0, origin_batch_size=origin_batch_size0) - model_output1 = self._create_unpad_decode_model_output(model_output1, origin_batch_size=origin_batch_size1) + model_output0 = self._create_unpad_decode_model_output(model_output0, origin_batch_size=origin_batch_size0) + model_output1 = self._create_unpad_decode_model_output(model_output1, origin_batch_size=origin_batch_size1) return model_output0, model_output1 @final diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index d8f2329b24..447916fafc 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -32,6 +32,8 @@ class ModelInput: # Decode 逐行携带的 radix node 标识。相同 id 表示请求引用同一个共享 # radix node；该 id 只用于重建 diverse attention 的 b_mark_shared_group。 b_shared_radix_node_id: torch.Tensor = None + # GPU write slots are resolved after padding, from the final row positions. + mem_indexes: torch.Tensor = None is_prefill: bool = False b_ready_cache_len: torch.Tensor = None # Request/row-aligned MRoPE position offset. It is decode-only; prefill @@ -189,9 +191,6 @@ class ModelOutput: # attached here. ModelOutput therefore owns a stable output view instead # of the mutable collector used while the forward is still running. mtp_collector: Optional[ModelMtpOutputCollector] = None - # 用于判断 mem_indexes 是否成功写入 req manager 中的事件对象。 - prefill_mem_indexes_ready_event: torch.Event = None - # prompt_logics 用于在开启 return_all_prompt_logics 模式（如 enable_prompt_logprobs）时， # 保存整个 prefill 阶段每一个 token 位置对应的 logits（而非仅最后一个位置的 logits）。 # 此时 logits 依然只保存每个请求最后一个位置的 logits，prompt_logics 为可选项，仅在 diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index ee888c1a00..8e28ea03b0 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -1,4 +1,3 @@ -import os import torch import torch.distributed as dist import copy @@ -8,13 +7,10 @@ from typing import Optional from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args -from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import build_kv_indexes_and_input_ids from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -from lightllm.utils.torch_memory_saver_utils import ( - TorchMemorySaverWrapper, - MemoryTag, -) +from lightllm.utils.torch_memory_saver_utils import TorchMemorySaverWrapper from .infer_struct import InferStateInfo @@ -286,6 +282,12 @@ def warmup(self, model): multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], **model._gen_special_model_input(batch_size), ) + # Warm the ordinary decode gather path as well as the explicit-ID path below. + _, model_input.input_ids = build_kv_indexes_and_input_ids( + model_input, + model.req_manager.req_to_token_indexs, + model.req_manager.req_sampling_params_manager.req_to_next_token_ids, + ) model_output: ModelOutput = model.forward(model_input) del model_output del input_ids @@ -344,6 +346,11 @@ def warmup_overlap(self, model): multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], **model._gen_special_model_input(batch_size), ) + _, micro_batch.input_ids = build_kv_indexes_and_input_ids( + micro_batch, + model.req_manager.req_to_token_indexs, + model.req_manager.req_sampling_params_manager.req_to_next_token_ids, + ) decode_batches.append(micro_batch) del micro_batch diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index 28d44fefbf..1eb7e3000b 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -1,6 +1,4 @@ -import os import torch -import copy import bisect import triton from typing import List, Tuple @@ -8,7 +6,6 @@ from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor -from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from .infer_struct import InferStateInfo from .cuda_graph import CudaGraph diff --git a/lightllm/common/basemodel/triton_kernel/copy_kv_index_to_req.py b/lightllm/common/basemodel/triton_kernel/copy_kv_index_to_req.py index efe563d764..f992e25b02 100644 --- a/lightllm/common/basemodel/triton_kernel/copy_kv_index_to_req.py +++ b/lightllm/common/basemodel/triton_kernel/copy_kv_index_to_req.py @@ -1,204 +1,103 @@ import torch - import triton import triton.language as tl +from lightllm.common.basemodel.batch_objs import ModelInput @triton.jit -def _fwd_kernel_copy_kv_index_to_req( - req_to_token_indexs, b_req_idx, b_seq_len, memindex, stride_req_to_token_b, stride_req_to_token_s -): - cur_index = tl.program_id(0) - cur_req_idx = tl.load(b_req_idx + cur_index) - cur_token_index = tl.load(memindex + cur_index) - cur_seq_len = tl.load(b_seq_len + cur_index) - dest_offset = req_to_token_indexs + cur_req_idx * stride_req_to_token_b + (cur_seq_len - 1) * stride_req_to_token_s - tl.store(dest_offset, cur_token_index) - return - - -@torch.no_grad() -def copy_kv_index_to_req(req_to_token_indexs, b_req_idx, b_seq_len, memindex): - seq_len = b_seq_len.shape[0] - assert b_seq_len.shape[0] == memindex.shape[0] and b_req_idx.shape[0] == b_seq_len.shape[0] - grid = (seq_len,) - num_warps = 1 - - _fwd_kernel_copy_kv_index_to_req[grid]( - req_to_token_indexs, - b_req_idx, - b_seq_len, - memindex, - req_to_token_indexs.stride(0), - req_to_token_indexs.stride(1), - num_warps=num_warps, - num_stages=1, - ) - return - - -@triton.jit -def _fwd_kernel_select_kv_index_from_req( - req_to_token_indexs, - b_req_idx, - b_seq_len, - out_memindex, - stride_req_to_token_b, - stride_req_to_token_s, -): - cur_index = tl.program_id(0) - cur_req_idx = tl.load(b_req_idx + cur_index) - cur_seq_len = tl.load(b_seq_len + cur_index) - src_offset = req_to_token_indexs + cur_req_idx * stride_req_to_token_b + (cur_seq_len - 1) * stride_req_to_token_s - tl.store(out_memindex + cur_index, tl.load(src_offset)) - - -@torch.no_grad() -def select_kv_index_from_req(req_to_token_indexs, b_req_idx, b_seq_len): - """Select the one logical KV slot consumed by each decode row.""" - out_memindex = torch.empty_like(b_req_idx) - _fwd_kernel_select_kv_index_from_req[(b_req_idx.shape[0],)]( - req_to_token_indexs, - b_req_idx, - b_seq_len, - out_memindex, - req_to_token_indexs.stride(0), - req_to_token_indexs.stride(1), - num_warps=1, - num_stages=1, - ) - return out_memindex - - -@triton.jit -def _fwd_kernel_copy_kv_index_to_req_prefill( - req_to_token_indexs, - b_req_idx, - b_seq_len, - b_ready_cache_len, - b_start_loc, - memindex, - stride_req_to_token_b, - stride_req_to_token_s, +def _build_kv_indexes_and_input_ids( + Table, + Req, + Ends, + Ready, + Starts, + Out, + Tokens, + Next, + Mtp, + Mixed, + SB: tl.constexpr, + SS: tl.constexpr, + NEXT_STRIDE: tl.constexpr, + BATCH: tl.constexpr, + PREFILL: tl.constexpr, + GATHER: tl.constexpr, BLOCK: tl.constexpr, ): - - block_index = tl.program_id(0) - batch_index = tl.program_id(1) - cur_req_idx = tl.load(b_req_idx + batch_index) - cur_seq_len = tl.load(b_seq_len + batch_index) - cur_ready_cache_len = tl.load(b_ready_cache_len + batch_index) - cur_start_loc = tl.load(b_start_loc + batch_index) - copy_len = cur_seq_len - cur_ready_cache_len - - block_range = block_index * BLOCK + tl.arange(0, BLOCK) - block_mask = block_range < copy_len - cur_token_index = tl.load(memindex + cur_start_loc + block_range, mask=block_mask) - dest_offset = ( - req_to_token_indexs - + cur_req_idx * stride_req_to_token_b - + (cur_ready_cache_len + block_range) * stride_req_to_token_s - ) - tl.store(dest_offset, cur_token_index, mask=block_mask) - - return - - -def get_triton_config(max_q_seq_len: int) -> tuple[int, int]: - if max_q_seq_len <= 512: - return 256, 2 - elif max_q_seq_len <= 4096: - return 512, 4 + pos = tl.arange(0, BLOCK) + if PREFILL: + batch = tl.program_id(0) + req, end = tl.load(Req + batch), tl.load(Ends + batch) + start, out_start = tl.load(Ready + batch), tl.load(Starts + batch) + pos += tl.program_id(1) * BLOCK + index = tl.load(Table + req * SB + (start + pos) * SS, pos < end - start, other=0) + tl.store(Out + out_start + pos, index, pos < end - start) + if GATHER: + if tl.program_id(1) == 0 and tl.load(Mixed + batch): + token = tl.load(Next + req * NEXT_STRIDE + tl.load(Mtp + batch)) + tl.store(Tokens + out_start, token) else: - return 1024, 8 - - -@torch.no_grad() -def copy_kv_index_to_req_prefill( - req_to_token_indexs: torch.Tensor, - b_req_idx: torch.Tensor, - b_seq_len: torch.Tensor, - b_ready_cache_len: torch.Tensor, - b_start_loc: torch.Tensor, - memindex: torch.Tensor, - max_q_seq_len: int, -): - batch_size = b_req_idx.shape[0] - BLOCK, num_warps = get_triton_config(max_q_seq_len) - grid = (triton.cdiv(max_q_seq_len, BLOCK), batch_size) - num_warps = 1 - - _fwd_kernel_copy_kv_index_to_req_prefill[grid]( + row = tl.program_id(0) * BLOCK + pos + valid = row < BATCH + req, end = tl.load(Req + row, valid, other=0), tl.load(Ends + row, valid, other=1) + tl.store(Out + row, tl.load(Table + req * SB + (end - 1) * SS, valid, other=0), valid) + if GATHER: + step = tl.load(Mtp + row, valid, other=0) + tl.store(Tokens + row, tl.load(Next + req * NEXT_STRIDE + step, valid, other=0), valid) + + +def build_kv_indexes_and_input_ids(model_input: ModelInput, req_to_token_indexs, next_token_ids=None): + """Return KV write indexes and input IDs, gathering into supplied or newly allocated IDs in one launch.""" + prefill = model_input.is_prefill + batch = model_input.batch_size + b_req_idx = model_input.b_req_idx + input_ids = model_input.input_ids + max_q_seq_len = model_input.max_q_seq_len + if input_ids is None and next_token_ids is not None: + input_ids = torch.empty_like(b_req_idx, dtype=torch.int64) + out = ( + torch.empty(input_ids.numel(), dtype=torch.int32, device=req_to_token_indexs.device) + if prefill + else torch.empty_like(b_req_idx, dtype=torch.int32) + ) + block = (256 if max_q_seq_len <= 512 else 512 if max_q_seq_len <= 4096 else 1024) if prefill else 256 + grid = (batch, triton.cdiv(max_q_seq_len, block)) if prefill else (triton.cdiv(batch, block),) + _build_kv_indexes_and_input_ids[grid]( req_to_token_indexs, b_req_idx, - b_seq_len, - b_ready_cache_len, - b_start_loc, - memindex, - req_to_token_indexs.stride(0), - req_to_token_indexs.stride(1), - BLOCK=BLOCK, - num_warps=num_warps, - num_stages=1, + model_input.b_seq_len, + model_input.b_ready_cache_len if prefill else None, + model_input.b_prefill_start_loc if prefill else None, + out, + input_ids, + next_token_ids, + model_input.b_mtp_index, + model_input.b_is_decode_req, + *req_to_token_indexs.stride(), + 0 if next_token_ids is None else next_token_ids.stride(0), + 0 if prefill else batch, + prefill, + next_token_ids is not None, + block, + num_warps=block // 128 if prefill else 1, ) + return out, input_ids @triton.jit -def _fwd_kernel_select_kv_index_from_req_prefill( - req_to_token_indexs, - b_req_idx, - b_seq_len, - b_ready_cache_len, - b_start_loc, - out_memindex, - stride_req_to_token_b, - stride_req_to_token_s, - BLOCK: tl.constexpr, -): - block_index = tl.program_id(0) - batch_index = tl.program_id(1) - cur_req_idx = tl.load(b_req_idx + batch_index) - cur_seq_len = tl.load(b_seq_len + batch_index) - cur_ready_cache_len = tl.load(b_ready_cache_len + batch_index) - cur_start_loc = tl.load(b_start_loc + batch_index) - copy_len = cur_seq_len - cur_ready_cache_len - - block_range = block_index * BLOCK + tl.arange(0, BLOCK) - block_mask = block_range < copy_len - src_offset = ( - req_to_token_indexs - + cur_req_idx * stride_req_to_token_b - + (cur_ready_cache_len + block_range) * stride_req_to_token_s - ) - memindex = tl.load(src_offset, mask=block_mask) - tl.store(out_memindex + cur_start_loc + block_range, memindex, mask=block_mask) - - -@torch.no_grad() -def select_kv_index_from_req_prefill( - req_to_token_indexs, - b_req_idx, - b_seq_len, - b_ready_cache_len, - b_start_loc, - max_q_seq_len, - token_num, -): - """Select the logical KV slots consumed by the current prefill step.""" - out_memindex = torch.empty((token_num,), dtype=torch.int32, device=req_to_token_indexs.device) - block, num_warps = get_triton_config(max_q_seq_len) - grid = (triton.cdiv(max_q_seq_len, block), b_req_idx.shape[0]) - _fwd_kernel_select_kv_index_from_req_prefill[grid]( - req_to_token_indexs, - b_req_idx, - b_seq_len, - b_ready_cache_len, - b_start_loc, - out_memindex, - req_to_token_indexs.stride(0), - req_to_token_indexs.stride(1), - BLOCK=block, - num_warps=num_warps, - num_stages=1, +def _update_req_token_indexes(Table, Packet, STRIDE: tl.constexpr, BLOCK: tl.constexpr, PAGE_SIZE: tl.constexpr): + row, block = tl.program_id(0), tl.program_id(1) + req = tl.load(Packet + 4 * row) + start = tl.load(Packet + 4 * row + 1) + size = tl.load(Packet + 4 * row + 2) + # Metadata offsets use token slots; the payload follows all four-column headers. + offset = 4 * tl.num_programs(0) + tl.load(Packet + 4 * row + 3) // PAGE_SIZE + pos = block * BLOCK + tl.arange(0, BLOCK) + value = tl.load(Packet + offset + pos // PAGE_SIZE, pos < size, other=0) + pos % PAGE_SIZE + tl.store(Table + req * STRIDE + start + pos, value, pos < size) + + +def update_req_token_indexes(table, packet, count, max_size, page_size=1): + _update_req_token_indexes[(count, triton.cdiv(max_size, 256))]( + table, packet, table.stride(0), 256, page_size, num_warps=4 ) - return out_memindex diff --git a/lightllm/common/basemodel/triton_kernel/dynamic_mtp_utils.py b/lightllm/common/basemodel/triton_kernel/dynamic_mtp_utils.py index ad995b2b5d..6dd8890a2e 100644 --- a/lightllm/common/basemodel/triton_kernel/dynamic_mtp_utils.py +++ b/lightllm/common/basemodel/triton_kernel/dynamic_mtp_utils.py @@ -209,86 +209,40 @@ def _compact_decode_model_input( ) -> ModelInput: assert not model_input.is_prefill assert selected_row_mask.is_cuda - assert model_input.b_req_idx.is_cuda - assert model_input.b_mtp_index.is_cuda - assert model_input.b_seq_len.is_cuda - assert model_input.b_shared_seq_len.is_cuda - assert model_input.b_shared_radix_node_id.is_cuda - # Dynamic scheduling guarantees exactly dynamic_batch_size selected rows. selected_row_mask = selected_row_mask.to(torch.int32) old_batch_size = model_input.b_req_idx.shape[0] selected_dst_pos = torch.empty((old_batch_size,), dtype=torch.int32, device=model_input.b_req_idx.device) - - out_input_ids = None - if model_input.input_ids is not None: - assert model_input.input_ids.is_cuda - out_input_ids = torch.empty( - (dynamic_batch_size,), dtype=model_input.input_ids.dtype, device=model_input.input_ids.device - ) - - out_b_shared_seq_len = torch.empty( - (dynamic_batch_size,), dtype=model_input.b_shared_seq_len.dtype, device=model_input.b_shared_seq_len.device - ) - out_b_shared_radix_node_id = torch.empty( - (dynamic_batch_size,), - dtype=model_input.b_shared_radix_node_id.dtype, - device=model_input.b_shared_radix_node_id.device, - ) - - out_b_req_idx = torch.empty( - (dynamic_batch_size,), dtype=model_input.b_req_idx.dtype, device=model_input.b_req_idx.device + fields = ( + "input_ids", + "b_req_idx", + "b_mtp_index", + "b_seq_len", + "b_position_delta", + "b_shared_seq_len", + "b_shared_radix_node_id", ) - out_b_mtp_index = torch.empty( - (dynamic_batch_size,), dtype=model_input.b_mtp_index.dtype, device=model_input.b_mtp_index.device - ) - out_b_seq_len = torch.empty( - (dynamic_batch_size,), dtype=model_input.b_seq_len.dtype, device=model_input.b_seq_len.device - ) - out_b_position_delta = None - if model_input.b_position_delta is not None: - assert model_input.b_position_delta.is_cuda - out_b_position_delta = torch.empty( - (dynamic_batch_size,), - dtype=model_input.b_position_delta.dtype, - device=model_input.b_position_delta.device, - ) - - dummy_1d = model_input.b_req_idx - BLOCK_SIZE = triton.next_power_of_2(old_batch_size) - grid = (1,) - _fwd_kernel_compact_dynamic_mtp_model_input[grid]( - input_ids=model_input.input_ids if model_input.input_ids is not None else dummy_1d, - out_input_ids=out_input_ids if out_input_ids is not None else dummy_1d, - b_req_idx=model_input.b_req_idx, - out_b_req_idx=out_b_req_idx, - b_mtp_index=model_input.b_mtp_index, - out_b_mtp_index=out_b_mtp_index, - b_seq_len=model_input.b_seq_len, - out_b_seq_len=out_b_seq_len, - b_position_delta=model_input.b_position_delta if model_input.b_position_delta is not None else dummy_1d, - out_b_position_delta=out_b_position_delta if out_b_position_delta is not None else dummy_1d, - b_shared_seq_len=model_input.b_shared_seq_len, - out_b_shared_seq_len=out_b_shared_seq_len, - b_shared_radix_node_id=model_input.b_shared_radix_node_id, - out_b_shared_radix_node_id=out_b_shared_radix_node_id, + tensors = {} + for name in fields: + value = getattr(model_input, name) + if value is not None: + assert value.is_cuda + tensors[name] = value + tensors["out_" + name] = value.new_empty((dynamic_batch_size,)) if value is not None else None + + _fwd_kernel_compact_dynamic_mtp_model_input[(1,)]( + **tensors, selected_mask=selected_row_mask, selected_dst_pos=selected_dst_pos, batch_size=old_batch_size, HAS_INPUT_IDS=model_input.input_ids is not None, HAS_B_POSITION_DELTA=model_input.b_position_delta is not None, - BLOCK_SIZE=BLOCK_SIZE, + BLOCK_SIZE=triton.next_power_of_2(old_batch_size), num_warps=8, num_stages=1, ) - - model_input.input_ids = out_input_ids - model_input.b_req_idx = out_b_req_idx - model_input.b_mtp_index = out_b_mtp_index - model_input.b_seq_len = out_b_seq_len - model_input.b_position_delta = out_b_position_delta - model_input.b_shared_seq_len = out_b_shared_seq_len - model_input.b_shared_radix_node_id = out_b_shared_radix_node_id + for name in fields: + setattr(model_input, name, tensors["out_" + name]) if model_input.mtp_draft_input_hiddens is not None: assert model_input.mtp_draft_input_hiddens.is_cuda diff --git a/lightllm/common/basemodel/triton_kernel/fa3_utils.py b/lightllm/common/basemodel/triton_kernel/fa3_utils.py index 23d8b0f818..b9081e7776 100644 --- a/lightllm/common/basemodel/triton_kernel/fa3_utils.py +++ b/lightllm/common/basemodel/triton_kernel/fa3_utils.py @@ -275,36 +275,3 @@ def build_dynamic_spec_fa3_decode_params( num_stages=1, ) return b_q_seq_len, b_kv_seq_len, b_att_req_idx, b_att_seq_len - - -def test_page_table_copy(): - import torch - - batch_size, seq_len = 2, 8 - - req_to_token_indexs = torch.arange(batch_size * seq_len, dtype=torch.int32).reshape(batch_size, seq_len).cuda() - - page_table = torch.full((batch_size, seq_len), -1, dtype=torch.int32, device="cuda") - - b_req_idx = torch.tensor([0, 2, 1, 3], dtype=torch.int32, device="cuda")[::2] - print(b_req_idx.stride()) - - page_table_copy(page_table, req_to_token_indexs, b_req_idx) - - print("req_to_token_indexs:") - print(req_to_token_indexs.cpu().numpy()) - print("b_req_idx:", b_req_idx.cpu().numpy()) - print("page_table:") - print(page_table.cpu().numpy()) - - for batch in range(batch_size): - src_idx = b_req_idx[batch].item() - expected = req_to_token_indexs[src_idx].cpu().numpy() - got = page_table[batch].cpu().numpy() - assert (expected == got).all(), f"Batch {batch} mismatch: expected {expected}, got {got}" - - print("✅ Test passed!") - - -if __name__ == "__main__": - test_page_table_copy() diff --git a/lightllm/common/basemodel/triton_kernel/gather_token_id.py b/lightllm/common/basemodel/triton_kernel/gather_token_id.py index a5dfc6d226..c1794ec00e 100644 --- a/lightllm/common/basemodel/triton_kernel/gather_token_id.py +++ b/lightllm/common/basemodel/triton_kernel/gather_token_id.py @@ -86,325 +86,3 @@ def scatter_token( num_stages=1, ) return - - -@triton.jit -def _fwd_kernel_gather( - req_to_next_token_ids, - req_to_next_token_ids_stride, - req_to_next_token_ids_stride_1, - output, - b_req_idx, - b_mtp_index, - num_size, - BLOCK: tl.constexpr, -): - block_index = tl.program_id(0) - block_range = block_index * BLOCK + tl.arange(0, BLOCK) - block_mask = block_range < num_size - cur_req_idx = tl.load(b_req_idx + block_range, mask=block_mask) - cur_mtp_index = tl.load(b_mtp_index + block_range, mask=block_mask) - cur_next_token_id = tl.load( - req_to_next_token_ids + cur_req_idx * req_to_next_token_ids_stride + cur_mtp_index, mask=block_mask - ) - tl.store(output + block_range, cur_next_token_id, mask=block_mask) - return - - -def gather_token(req_to_next_token_ids: torch.Tensor, b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor): - """ - This function is used to gather the token_info(CPU tensor) to the token_info(GPU tensor). - Args: - req_to_token_info: (max_req_num, max_mtp_step) - b_req_idx: (batch_size,) - b_mtp_index: (batch_size,) - Returns: - output: (batch_size,) - """ - batch_size = b_req_idx.shape[0] - output = torch.empty(batch_size, dtype=req_to_next_token_ids.dtype, device="cuda") - BLOCK = 256 - grid = (triton.cdiv(batch_size, BLOCK),) - num_warps = 1 - _fwd_kernel_gather[grid]( - req_to_next_token_ids=req_to_next_token_ids, - req_to_next_token_ids_stride=req_to_next_token_ids.stride(0), - req_to_next_token_ids_stride_1=req_to_next_token_ids.stride(1), - output=output, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - num_size=batch_size, - BLOCK=BLOCK, - num_warps=num_warps, - num_stages=1, - ) - return output - - -@triton.jit -def _fwd_kernel_gather_prefill_decode_mixed( - input_ids, - req_to_next_token_ids, - req_to_next_token_ids_stride, - req_to_next_token_ids_stride_1, - b_req_idx, - b_mtp_index, - b_is_decode_req, - b_prefill_start_loc, - num_size, - BLOCK: tl.constexpr, -): - block_index = tl.program_id(0) - block_range = block_index * BLOCK + tl.arange(0, BLOCK) - block_mask = block_range < num_size - cur_req_idx = tl.load(b_req_idx + block_range, mask=block_mask) - cur_mtp_index = tl.load(b_mtp_index + block_range, mask=block_mask) - cur_next_token_id = tl.load( - req_to_next_token_ids + cur_req_idx * req_to_next_token_ids_stride + cur_mtp_index, mask=block_mask - ) - cur_is_decode_req = tl.load(b_is_decode_req + block_range, mask=block_mask, other=False) - cur_prefill_start_loc = tl.load(b_prefill_start_loc + block_range, mask=block_mask, other=-1) - - tl.store(input_ids + cur_prefill_start_loc, cur_next_token_id, mask=block_mask & cur_is_decode_req) - return - - -def gather_token_prefill_decode_mixed( - input_ids: torch.Tensor, - req_to_next_token_ids: torch.Tensor, - b_req_idx: torch.Tensor, - b_mtp_index: torch.Tensor, - b_is_decode_req: torch.Tensor, - b_prefill_start_loc: torch.Tensor, -): - """ - This function is used to gather the token_info(CPU tensor) to the token_info(GPU tensor). - Args: - input_ids: (batch_size,) - req_to_next_token_ids: (max_req_num, max_mtp_step) - b_req_idx: (batch_size,) - b_mtp_index: (batch_size,) - b_is_decode_req: (batch_size,) - b_prefill_start_loc: (batch_size,) - Returns: - input_ids: - """ - batch_size = b_req_idx.shape[0] - BLOCK = 256 - grid = (triton.cdiv(batch_size, BLOCK),) - num_warps = 1 - _fwd_kernel_gather_prefill_decode_mixed[grid]( - input_ids=input_ids, - req_to_next_token_ids=req_to_next_token_ids, - req_to_next_token_ids_stride=req_to_next_token_ids.stride(0), - req_to_next_token_ids_stride_1=req_to_next_token_ids.stride(1), - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - b_is_decode_req=b_is_decode_req, - b_prefill_start_loc=b_prefill_start_loc, - num_size=batch_size, - BLOCK=BLOCK, - num_warps=num_warps, - num_stages=1, - ) - return input_ids - - -def test_scatter_token_to_cpu(): - batch_size = 30 - req_to_token_info = torch.zeros((1000, 1), dtype=torch.float32, pin_memory=True) - token_info = torch.randn((batch_size,)).cuda() - req_ids = torch.arange(20, 20 + batch_size, dtype=torch.int32).cuda() - mtp_index = torch.zeros((batch_size,), dtype=torch.int32).cuda() - scatter_token(token_info, req_to_token_info, req_ids, mtp_index) - diff = (req_to_token_info[20 : 20 + batch_size].cuda().view(-1) - token_info).abs().max() - assert diff < 1e-6 - print("test_scatter_token_to_cpu passed") - - -def test_gather_token(): - batch_size = 30 - req_to_token_info = torch.zeros((1000, 1), dtype=torch.float32, pin_memory=True) - token_info = torch.randn((batch_size,)).cuda() - req_ids = torch.arange(20, 20 + batch_size, dtype=torch.int32).cuda() - mtp_index = torch.zeros((batch_size,), dtype=torch.int32).cuda() - scatter_token(token_info, req_to_token_info, req_ids, mtp_index) - output = gather_token( - req_to_next_token_ids=req_to_token_info, - b_req_idx=req_ids, - b_mtp_index=mtp_index, - ) - diff = (token_info - output).abs().max() - assert diff < 1e-6 - print("test_gather_token passed") - - -def _ref_gather_token_prefill_decode_mixed( - input_ids: torch.Tensor, - req_to_next_token_ids: torch.Tensor, - b_req_idx: torch.Tensor, - b_mtp_index: torch.Tensor, - b_is_decode_req: torch.Tensor, - b_prefill_start_loc: torch.Tensor, -) -> torch.Tensor: - out = input_ids.clone() - table = req_to_next_token_ids.detach().cpu() - req_idx_cpu = b_req_idx.detach().cpu() - mtp_cpu = b_mtp_index.detach().cpu() - is_decode_cpu = b_is_decode_req.detach().cpu() - start_loc_cpu = b_prefill_start_loc.detach().cpu() - for i in range(req_idx_cpu.shape[0]): - if is_decode_cpu[i].item(): - rid = int(req_idx_cpu[i].item()) - mid = int(mtp_cpu[i].item()) - loc = int(start_loc_cpu[i].item()) - out[loc] = table[rid, mid] - return out - - -def _run_gather_token_prefill_decode_mixed_case( - input_ids: torch.Tensor, - req_to_next_token_ids: torch.Tensor, - b_req_idx: torch.Tensor, - b_mtp_index: torch.Tensor, - b_is_decode_req: torch.Tensor, - b_prefill_start_loc: torch.Tensor, -): - input_cuda = input_ids.clone().cuda() - req_table = req_to_next_token_ids.cuda() - b_req_idx_cuda = b_req_idx.cuda() - b_mtp_index_cuda = b_mtp_index.cuda() - b_is_decode_cuda = b_is_decode_req.cuda() - b_start_loc_cuda = b_prefill_start_loc.cuda() - - expected = _ref_gather_token_prefill_decode_mixed( - input_cuda, - req_table, - b_req_idx_cuda, - b_mtp_index_cuda, - b_is_decode_cuda, - b_start_loc_cuda, - ) - gather_token_prefill_decode_mixed( - input_cuda, - req_table, - b_req_idx_cuda, - b_mtp_index_cuda, - b_is_decode_cuda, - b_start_loc_cuda, - ) - diff = (input_cuda - expected).abs().max() - assert diff < 1e-6, f"max diff {diff.item()}" - - -def test_gather_token_prefill_decode_mixed_decode_only(): - """仅 decode 行：按 b_prefill_start_loc 写入 req_to_next_token_ids 中的 next token。""" - req_to_next_token_ids = torch.zeros((32, 4), dtype=torch.int64, device="cuda") - req_to_next_token_ids[3, 0] = 42 - req_to_next_token_ids[7, 0] = 99 - req_to_next_token_ids[11, 2] = 17 - - input_ids = torch.tensor([0, 0, 0], dtype=torch.int64, device="cuda") - b_req_idx = torch.tensor([3, 7, 11], dtype=torch.int32, device="cuda") - b_mtp_index = torch.tensor([0, 0, 2], dtype=torch.int32, device="cuda") - b_is_decode_req = torch.tensor([True, True, True], dtype=torch.bool, device="cuda") - b_prefill_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda") - - _run_gather_token_prefill_decode_mixed_case( - input_ids, req_to_next_token_ids, b_req_idx, b_mtp_index, b_is_decode_req, b_prefill_start_loc - ) - print("test_gather_token_prefill_decode_mixed_decode_only passed") - - -def test_gather_token_prefill_decode_mixed_mixed_batch(): - """prefill + decode 混合：仅 decode 位置被覆盖，prefill token 保持不变。""" - req_to_next_token_ids = torch.zeros((16, 2), dtype=torch.int64, device="cuda") - req_to_next_token_ids[5, 0] = 9001 - - # prefill [10,11,12] | decode placeholder | prefill [20,21] - input_ids = torch.tensor([10, 11, 12, -1, 20, 21], dtype=torch.int64, device="cuda") - b_req_idx = torch.tensor([0, 5, 1], dtype=torch.int32, device="cuda") - b_mtp_index = torch.tensor([0, 0, 0], dtype=torch.int32, device="cuda") - b_is_decode_req = torch.tensor([False, True, False], dtype=torch.bool, device="cuda") - b_q_seq_len = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") - b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len - - _run_gather_token_prefill_decode_mixed_case( - input_ids, req_to_next_token_ids, b_req_idx, b_mtp_index, b_is_decode_req, b_prefill_start_loc - ) - print("test_gather_token_prefill_decode_mixed_mixed_batch passed") - - -def test_gather_token_prefill_decode_mixed_prefill_only_unchanged(): - """无 decode 行时 input_ids 不应被修改。""" - req_to_next_token_ids = torch.full((8, 1), 777, dtype=torch.int64, device="cuda") - input_ids = torch.tensor([1, 2, 3, 4], dtype=torch.int64, device="cuda") - b_req_idx = torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda") - b_mtp_index = torch.zeros(3, dtype=torch.int32, device="cuda") - b_is_decode_req = torch.zeros(3, dtype=torch.bool, device="cuda") - b_q_seq_len = torch.tensor([2, 1, 1], dtype=torch.int32, device="cuda") - b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len - - before = input_ids.clone() - gather_token_prefill_decode_mixed( - input_ids, - req_to_next_token_ids, - b_req_idx, - b_mtp_index, - b_is_decode_req, - b_prefill_start_loc, - ) - assert torch.equal(input_ids, before) - print("test_gather_token_prefill_decode_mixed_prefill_only_unchanged passed") - - -def test_gather_token_prefill_decode_mixed_large_batch(): - """batch_size > 256，覆盖多 block 的 triton grid。""" - batch_size = 300 - max_req = 400 - req_to_next_token_ids = torch.arange(max_req * 2, dtype=torch.int64, device="cuda").view(max_req, 2) - input_ids = torch.zeros(batch_size, dtype=torch.int64, device="cuda") - b_req_idx = torch.arange(10, 10 + batch_size, dtype=torch.int32, device="cuda") - b_mtp_index = (b_req_idx % 2).to(torch.int32) - b_is_decode_req = torch.ones(batch_size, dtype=torch.bool, device="cuda") - b_prefill_start_loc = torch.arange(batch_size, dtype=torch.int32, device="cuda") - - _run_gather_token_prefill_decode_mixed_case( - input_ids, req_to_next_token_ids, b_req_idx, b_mtp_index, b_is_decode_req, b_prefill_start_loc - ) - print("test_gather_token_prefill_decode_mixed_large_batch passed") - - -def test_gather_token_prefill_decode_mixed_roundtrip_with_scatter(): - """scatter_token 写入后，mixed gather 能读回同一 next token。""" - batch_size = 16 - req_to_next_token_ids = torch.zeros((64, 3), dtype=torch.float32, pin_memory=True) - token_info = torch.arange(100, 100 + batch_size, dtype=torch.float32, device="cuda") - b_req_idx = torch.arange(4, 4 + batch_size, dtype=torch.int32, device="cuda") - b_mtp_index = torch.zeros(batch_size, dtype=torch.int32, device="cuda") - scatter_token(token_info, req_to_next_token_ids, b_req_idx, b_mtp_index) - - input_ids = torch.zeros(batch_size, dtype=torch.int64, device="cuda") - b_is_decode_req = torch.ones(batch_size, dtype=torch.bool, device="cuda") - b_prefill_start_loc = torch.arange(batch_size, dtype=torch.int32, device="cuda") - - gather_token_prefill_decode_mixed( - input_ids, - req_to_next_token_ids, - b_req_idx, - b_mtp_index, - b_is_decode_req, - b_prefill_start_loc, - ) - assert torch.equal(input_ids, token_info.to(torch.int64)) - print("test_gather_token_prefill_decode_mixed_roundtrip_with_scatter passed") - - -if __name__ == "__main__": - test_scatter_token_to_cpu() - test_gather_token() - test_gather_token_prefill_decode_mixed_decode_only() - test_gather_token_prefill_decode_mixed_mixed_batch() - test_gather_token_prefill_decode_mixed_prefill_only_unchanged() - test_gather_token_prefill_decode_mixed_large_batch() - test_gather_token_prefill_decode_mixed_roundtrip_with_scatter() diff --git a/lightllm/common/infer_utils.py b/lightllm/common/infer_utils.py deleted file mode 100644 index e1b9cc3830..0000000000 --- a/lightllm/common/infer_utils.py +++ /dev/null @@ -1,15 +0,0 @@ -from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import copy_kv_index_to_req_prefill - - -def init_req_to_token_indexes( - req_to_token_indexs, b_req_idx, b_seq_len, b_ready_cache_len, b_start_loc, alloc_mem_index, max_q_seq_len -): - copy_kv_index_to_req_prefill( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_ready_cache_len=b_ready_cache_len, - b_start_loc=b_start_loc, - memindex=alloc_mem_index, - max_q_seq_len=max_q_seq_len, - ) diff --git a/lightllm/common/kv_cache_mem_manager/allocator.py b/lightllm/common/kv_cache_mem_manager/allocator.py index 1331943b1a..60e37f91e1 100644 --- a/lightllm/common/kv_cache_mem_manager/allocator.py +++ b/lightllm/common/kv_cache_mem_manager/allocator.py @@ -1,103 +1,74 @@ +import numpy as np import torch from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt from lightllm.utils.dist_utils import get_current_rank_in_node -from lightllm.utils.log_utils import init_logger -from typing import Union, List - -logger = init_logger(__name__) class KvCacheAllocator: - def __init__(self, size: int) -> None: - self.size = size - self.mem_state = torch.arange( - 0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self.mark_start = 0 - self.mark_end = self.size - - self._mem_state_return = torch.arange( - 0, self.size * 3, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self._return_start = 0 - - self.can_use_mem_size = self.size - - rank_in_node = get_current_rank_in_node() - # 用共享内存进行共享，router 模块读取进行精确的调度估计；基础层会统一添加服务前缀以防止实例冲突。 - self.shared_can_use_token_num = SharedInt(f"mem_manger_can_use_token_num_{rank_in_node}") - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - return - - def alloc(self, need_size) -> torch.Tensor: - if need_size > self.mark_end - self.mark_start: - logger.error(f"warn no enough cache need_size {need_size} left_size {self.can_use_mem_size}") - assert False, "error alloc state" - - start = self.mark_start - end = self.mark_start + need_size - self.mark_start += need_size - + """Allocate token indexes; consecutive pages use compressed storage when page_size > 1.""" + + def __init__(self, size: int, page_size: int = 1) -> None: + self.page_size = page_size + rank = get_current_rank_in_node() + self.shared_can_use_token_num = SharedInt(f"mem_manger_can_use_token_num_{rank}") + self.resize(size) + + def _alloc_token_index_blocks(self, need_size: int): + """Borrow index storage until the next mutation; P > 1 stores each block's first token index.""" + assert 0 <= need_size <= self.can_use_mem_size, "error alloc state" + start, size = self.size - self.can_use_mem_size, need_size + if self.page_size > 1: + assert need_size % self.page_size == 0 + start, size = start // self.page_size, size // self.page_size self.can_use_mem_size -= need_size self.shared_can_use_token_num.set_value(self.can_use_mem_size) + return self._free_token_index_blocks[start : start + size] - # 利用缓冲区返回，避免异步情况下的内存竞争 - if self._return_start + need_size > self._mem_state_return.shape[0]: - self._return_start = 0 - ans = self._mem_state_return[self._return_start : self._return_start + need_size] - ans.copy_(self.mem_state[start:end]) - self._return_start += need_size - return ans - - def free(self, free_index: Union[torch.Tensor, List[int]]): - """_summary_ - - Args: - free_index (torch.Tensor): _description_ - """ - - end = self.mark_start - start = self.mark_start - len(free_index) - assert start >= 0, f"error free state start: {self.mark_start} free len {len(free_index)}" - - if isinstance(free_index, list): - self.mem_state.numpy()[start:end] = free_index + def alloc(self, need_size: int) -> torch.Tensor: + index_blocks = self._alloc_token_index_blocks(need_size) + out = torch.empty(need_size, dtype=torch.int32, device="cpu", pin_memory=True) + if self.page_size > 1: + self.expand_token_index_blocks(index_blocks, out.numpy()) else: - # 从 gpu 到 cpu 的拷贝操作是流内阻塞操作 - self.mem_state[start:end] = free_index - - self.mark_start -= len(free_index) - - self.can_use_mem_size += len(free_index) + out.numpy()[:] = index_blocks + return out + + def expand_token_index_blocks(self, index_blocks, out=None): + assert self.page_size > 1 + if out is None: + out = np.empty(len(index_blocks) * self.page_size, dtype=np.int32) + # Short pages expand faster along the page-base axis than in short token runs. + np.add( + index_blocks, + self.page_offsets[:, None], + out=out.reshape(-1, self.page_size).T, + order="C" if self.page_size < 8 else "F", + ) + return out + + def free(self, free_index): + size = len(free_index) + assert size % self.page_size == 0 + values = free_index.cpu().numpy() if isinstance(free_index, torch.Tensor) else np.asarray(free_index) + self.free_token_index_blocks(values[:: self.page_size]) + + def free_token_index_blocks(self, index_blocks): + size = len(index_blocks) * self.page_size + assert size <= self.size - self.can_use_mem_size, "error free state" + end = (self.size - self.can_use_mem_size) // self.page_size + self._free_token_index_blocks[end - len(index_blocks) : end] = index_blocks + self.can_use_mem_size += size self.shared_can_use_token_num.set_value(self.can_use_mem_size) - if self.can_use_mem_size == len(self.mem_state): - logger.debug(f"freed all gpu mem size {self.can_use_mem_size}") - return - def free_all(self): - self.mem_state.numpy()[:] = list(range(0, len(self.mem_state))) - self.mark_start = 0 - self.mark_end = len(self.mem_state) - self.can_use_mem_size = len(self.mem_state) + self._free_token_index_blocks[:] = np.arange(0, self.size, self.page_size, dtype=np.int32) + self.can_use_mem_size = self.size self.shared_can_use_token_num.set_value(self.can_use_mem_size) - return def resize(self, new_size: int) -> None: - """ - just for test code - """ + assert new_size % self.page_size == 0 self.size = new_size - self.mem_state = torch.arange( - 0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self.mark_start = 0 - self.mark_end = self.size - - self._mem_state_return = torch.arange( - 0, self.size * 3, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self._return_start = 0 - - self.can_use_mem_size = self.size - self.shared_can_use_token_num.set_value(self.can_use_mem_size) + self._free_token_index_blocks = np.arange(0, new_size, self.page_size, dtype=np.int32) + self.page_offsets = np.arange(self.page_size, dtype=np.int32) + self.can_use_mem_size = new_size + self.shared_can_use_token_num.set_value(new_size) diff --git a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py index 4be820b753..96ed9bc668 100644 --- a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py @@ -1,8 +1,6 @@ import torch -import os -import torch.distributed as dist from .mem_manager import MemoryManager -from typing import List, Union, Any +from typing import List, Any from lightllm.utils.log_utils import init_logger from lightllm.common.kv_trans_kernel.nixl_kv_trans import mla_page_io from .operator import Deepseek2MemOperator diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 38c870ce4e..8dfd7f1721 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -1,6 +1,4 @@ import math -import os -import re import torch import torch.distributed as dist import torch.multiprocessing as mp @@ -47,7 +45,7 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False self.page_size = get_env_start_args().page_size self.size = self.size // self.page_size * self.page_size - self.allocator = KvCacheAllocator(self.size) + self.allocator = KvCacheAllocator(self.size, self.page_size) self._init_buffers( self.size, @@ -194,9 +192,6 @@ def read_page_kv_move_buffer_to_mem( # logger.info(f"dst token tensor {self.kv_buffer[:, mem_indexes[0], 0, 0]}") # logger.info(f"dst page token tensor {cur_page[0, :, 0, 0]}") - def _free_buffers(self): - self.kv_buffer = None - def alloc(self, need_size) -> torch.Tensor: assert ( need_size % self.page_size == 0 @@ -209,23 +204,6 @@ def free(self, free_index: Union[torch.Tensor, List[int]]) -> None: def free_all(self): self.allocator.free_all() - def resize_mem(self, new_size): - """ - just for test code - """ - size = new_size // self.page_size * self.page_size - dtype = self.dtype - head_num = self.head_num - head_dim = self.head_dim - layer_num = self.layer_num - - self.size = size - self.allocator.resize(size) - self._free_buffers() - self._init_buffers(size, dtype, head_num, head_dim, layer_num) - self.HOLD_TOKEN_MEMINDEXES = tuple(range(self.size, self.size + self.page_size)) - return - def get_index_kv_buffer(self, index): return {"kv_buffer": self.kv_buffer[:, index]} diff --git a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py index 6b02cbed63..39c08eb613 100755 --- a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py @@ -38,10 +38,6 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): device="cuda", ) - def _free_buffers(self): - self.kv_buffer = None - self.scale_buffer = None - def get_index_kv_buffer(self, index): return {"kv_buffer": self.kv_buffer[:, index], "scale_buffer": self.scale_buffer[:, index]} diff --git a/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py index 0f5e8e9370..6e64051081 100755 --- a/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py @@ -37,10 +37,6 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): device="cuda", ) - def _free_buffers(self): - self.kv_buffer = None - self.scale_buffer = None - def get_index_kv_buffer(self, index): return {"kv_buffer": self.kv_buffer[:, index], "scale_buffer": self.scale_buffer[:, index]} 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..713c6d0890 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -61,15 +61,6 @@ def _init_linear_att_buffers(self): 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 diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index cfe5d716ee..a5f1ea10bb 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -1,6 +1,9 @@ from typing import List +import numpy as np + import torch +from triton.runtime import driver from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.utils.envs_utils import get_env_start_args @@ -10,45 +13,31 @@ logger = init_logger("lightllm.common.req_manager") -class _ReqNode: - def __init__(self, index): - self.index = index - self.next: "_ReqNode" = None - - -class _ReqLinkedList: +class _ReqIndexPool: def __init__(self, max_request_num): - self.nodes = [_ReqNode(i) for i in range(max_request_num)] - self.marks = [0 for _ in range(max_request_num)] - self.root_node = _ReqNode(-1) - for i in range(0, max_request_num - 1): - self.nodes[i].next = self.nodes[i + 1] - self.root_node.next = self.nodes[0] - self.can_alloc_size = max_request_num - return + self.free_indexes = list(range(max_request_num - 1, -1, -1)) + self.marks = [0] * max_request_num + + @property + def can_alloc_size(self): + return len(self.free_indexes) def alloc(self): - if self.root_node.next is None: + if not self.free_indexes: logger.warning("alloc req index fail") return None - get_node = self.root_node.next - self.root_node.next = self.root_node.next.next - assert self.marks[get_node.index] == 0 - self.marks[get_node.index] = 1 - self.can_alloc_size -= 1 - return get_node.index + index = self.free_indexes.pop() + assert self.marks[index] == 0 + self.marks[index] = 1 + return index def free(self, index): assert self.marks[index] == 1 - node = self.nodes[index] - node.next = self.root_node.next - self.root_node.next = node self.marks[index] = 0 - self.can_alloc_size += 1 - return + self.free_indexes.append(index) def is_all_free(self): - return self.can_alloc_size == len(self.marks) + return len(self.free_indexes) == len(self.marks) class ReqManager: @@ -59,12 +48,16 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana # 的这个请求管理 id， 主要是为了兼容 DP 运行模式下，让各个 DP 能 padding 到 DP 中最大 # 的那个batch size 进行运行，所有 padding 的请求都会使用预留的这个请求管理 id 进行处理 # 这样让 DP 的实现更为简化一些。 - self.req_list = _ReqLinkedList(max_request_num) + self.req_list = _ReqIndexPool(max_request_num) page_size = get_env_start_args().page_size max_sequence_length = (max_sequence_length + page_size - 1) // page_size * page_size self.req_to_token_indexs = torch.zeros( (max_request_num + 1, max_sequence_length), dtype=torch.int32, device="cuda" ) + # P=1 stores every token index; P>1 stores the first token index of each consecutive block. + self._cpu_token_index_blocks = np.empty((max_request_num + 1, max_sequence_length // page_size), dtype=np.int32) + self.token_indexes_ready = torch.cuda.Event() + self._token_indexes_stream = None self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num) self.max_request_num = max_request_num self.HOLD_REQUEST_ID = max_request_num @@ -86,25 +79,163 @@ def init_hold_request_indexs(self): # HOLD_REQUEST_ID 对应的请求行供 DP padding、overlap microbatch 等占位请求使用。将该行 # 按 page_size 划分后，每一页都映射到 mem_manager 额外保留的同一个物理页；这样占位请求 # 无论访问哪一个逻辑位置，都会落到合法且不会参与正常分配的 KV cache 地址上。 - hold_row = self.req_to_token_indexs[self.HOLD_REQUEST_ID] - hold_page = torch.tensor( - self.mem_manager.HOLD_TOKEN_MEMINDEXES, - dtype=hold_row.dtype, - device=hold_row.device, + self._cpu_token_index_blocks[self.HOLD_REQUEST_ID].fill(self.mem_manager.HOLD_TOKEN_MEMINDEXES[0]) + self.req_to_token_indexs[self.HOLD_REQUEST_ID].copy_( + self.get_cpu_token_indexes(self.HOLD_REQUEST_ID, 0, self.req_to_token_indexs.shape[1]) ) - hold_row.view(-1, self.mem_manager.page_size).copy_(hold_page) - return + + def get_cpu_token_index_blocks(self, req_idx, start, end): + """Borrow CPU index storage for [start, end); P>1 requires whole consecutive blocks.""" + page_size = self.mem_manager.page_size + if page_size > 1: + assert start % page_size == end % page_size == 0 + return self._cpu_token_index_blocks[req_idx, start // page_size : end // page_size] + return self._cpu_token_index_blocks[req_idx, start:end] + + def get_cpu_token_indexes(self, req_idx, start, end): + """Return owned token indexes for cache insertion and asynchronous PD metadata.""" + page_size = self.mem_manager.page_size + if page_size > 1: + first, offset = divmod(start, page_size) + index_blocks = self._cpu_token_index_blocks[req_idx, first : (end + page_size - 1) // page_size] + indexes = self.mem_manager.allocator.expand_token_index_blocks(index_blocks) + return torch.from_numpy(indexes[offset : offset + end - start]) + return torch.from_numpy(self._cpu_token_index_blocks[req_idx, start:end].copy()) + + def write_token_index_prefix(self, req_idx, start, indexes, req_prefix_lens): + """Update CPU prefix mappings and collect pending GPU copies.""" + page_size = self.mem_manager.page_size + assert start == 0 + size = indexes.numel() + assert size % page_size == 0 + self._cpu_token_index_blocks[req_idx, : size // page_size] = indexes.numpy()[::page_size] + if size: + req_prefix_lens[req_idx] = max(size, req_prefix_lens.get(req_idx, 0)) + + def copy_token_index_prefixes_to_gpu(self, req_prefix_lens): + """Copy the collected CPU token-index prefixes on the current CUDA stream.""" + if req_prefix_lens: + updates, offset = [], 0 + for row, size in req_prefix_lens.items(): + updates.append((row, 0, size, offset)) + offset += size + self._copy_token_indexes_to_gpu(updates) + + def copy_token_index_prefix(self, source_req_idx, dest_req_idx, size): + """Copy a published prefix on the GPU and retain its CPU page mapping.""" + self.get_cpu_token_index_blocks(dest_req_idx, 0, size)[:] = self.get_cpu_token_index_blocks( + source_req_idx, 0, size + ) + stream = self.wait_token_indexes() + table = self.req_to_token_indexs + table[dest_req_idx, :size].copy_(table[source_req_idx, :size], non_blocking=True) + self.token_indexes_ready.record(stream) + self._token_indexes_stream = stream + + def alloc_req_token_indexes(self, allocations, radix_cache=None): + """Allocate token capacity once and publish the same indexes to the CPU and GPU tables.""" + need = sum(size for _, size in allocations) + if radix_cache is not None: + radix_cache.free_radix_cache_to_get_enough_token(need) + page_size = self.mem_manager.page_size + index_blocks = self.mem_manager.allocator._alloc_token_index_blocks(need) + updates, offset = [], 0 + for req, size in allocations: + start = req.hold_kv_len + block_offset = offset // page_size + if size == page_size: + self._cpu_token_index_blocks[req.req_idx, start // page_size] = index_blocks[block_offset] + else: + self._cpu_token_index_blocks[ + req.req_idx, start // page_size : (start + size) // page_size + ] = index_blocks[block_offset : block_offset + size // page_size] + updates.append((req.req_idx, start, size, offset)) + req.hold_kv_len += size + offset += size + self._copy_token_indexes_to_gpu(updates, index_blocks) + + def alloc_token_indexes(self, req, end, radix_cache=None): + """Return write slots for [cur_kv_len, end), reusing any existing page tail.""" + size = self.get_need_alloc_token_num(req, end) + if size: + self.alloc_req_token_indexes([(req, size)], radix_cache) + return self.get_cpu_token_indexes(req.req_idx, req.cur_kv_len, end) + + def _copy_token_indexes_to_gpu(self, updates, index_blocks=None): + """Publish token ranges using one owned snapshot; expand compressed P>1 indexes only as needed.""" + from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import update_req_token_indexes + + page_size = self.mem_manager.page_size + if index_blocks is None: + parts = [ + self._cpu_token_index_blocks[row, start // page_size : (start + size) // page_size] + for row, start, size, _ in updates + ] + index_blocks = parts[0] if len(parts) == 1 else np.concatenate(parts) + if not len(index_blocks): + return + packed = len(updates) > 4 + header = 4 * len(updates) if packed else 0 + packet = torch.empty( + header + len(index_blocks) * (1 if packed else page_size), dtype=torch.int32, device="cpu", pin_memory=True + ) + data = packet.numpy() + if packed: + meta = data[:header].reshape(-1, 4) + meta[:] = updates + data[header:] = index_blocks + elif page_size > 1: + self.mem_manager.allocator.expand_token_index_blocks(index_blocks, data) + else: + data[:] = index_blocks + stream = self.wait_token_indexes() + if packed: + update_req_token_indexes( + self.req_to_token_indexs, + packet.cuda(non_blocking=True), + len(updates), + max(x[2] for x in updates), + page_size, + ) + else: + table = self.req_to_token_indexs + row_stride, token_stride = table.stride() + base = table.storage_offset() + for row, start, size, offset in updates: + table.as_strided((size,), (token_stride,), base + row * row_stride + start * token_stride).copy_( + packet if len(updates) == 1 else packet[offset : offset + size], + non_blocking=True, + ) + self.token_indexes_ready.record(stream) + self._token_indexes_stream = stream + + def get_need_alloc_token_num(self, req, target): + """Return additional page-aligned KV capacity, measured in token slots.""" + page_size = self.mem_manager.page_size + if page_size > 1: + size = max(0, (target + page_size - 1) // page_size * page_size - req.hold_kv_len) + assert size % page_size == 0 + return size + return max(0, target - req.hold_kv_len) + + def wait_token_indexes(self): + # Reuse the writer's Python Stream; query only the raw handle on the common path. + stream = self._token_indexes_stream + if stream is None or driver.active.get_current_stream(stream.device_index) != stream.cuda_stream: + stream = torch.cuda.current_stream() + self.token_indexes_ready.wait(stream) + return stream def alloc(self): return self.req_list.alloc() - def free(self, free_req_indexes: List[int], free_token_index): + def free(self, free_req_indexes: List[int], free_token_index_blocks): for req_index in free_req_indexes: self.req_list.free(req_index) if self.req_list.is_all_free(): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") - self.mem_manager.free(free_token_index) + self.mem_manager.allocator.free_token_index_blocks(free_token_index_blocks) def free_req(self, free_req_index: int): self.req_list.free(free_req_index) @@ -112,10 +243,9 @@ def free_req(self, free_req_index: int): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") return - def free_token(self, free_token_index): - self.mem_manager.free(free_token_index) - return + def free_token_index_blocks(self, free_token_index_blocks): + self.mem_manager.allocator.free_token_index_blocks(free_token_index_blocks) def free_all(self): - self.req_list = _ReqLinkedList(self.max_request_num) + self.req_list = _ReqIndexPool(self.max_request_num) return diff --git a/lightllm/common/req_manager/req_sampling_params.py b/lightllm/common/req_manager/req_sampling_params.py index 3774f38a2a..eb727aa952 100644 --- a/lightllm/common/req_manager/req_sampling_params.py +++ b/lightllm/common/req_manager/req_sampling_params.py @@ -37,6 +37,7 @@ def __init__(self, max_request_num): dtype=torch.int64, device="cuda", ) + self.req_to_next_token_ids[max_request_num].fill_(1) # HOLD request used by padded decode. # 常驻 CPU 快照，由 MTP scatter 后的异步拷贝更新；读取前复用已有 CUDA event 等待完成。 self.req_to_next_token_ids_cpu = torch.zeros( (max_request_num + 1, self.mtp_verify_width), @@ -105,15 +106,13 @@ def init_req_sampling_params(self, req: "InferReq"): else: self.req_to_out_token_id_counter[req.req_idx].fill_(0) if req.sampling_param.shm_param.input_penalty and req.need_out_token_id_statistics: - prompt_ids = g_pin_mem_manager.gen_from_list( - key="prompt_ids_for_penalty", - data=req.shm_req.get_prompt_ids_numpy(), - dtype=torch.int32, - ).cuda(non_blocking=True) + prompt_tokens = req.shm_req.get_prompt_ids_numpy() + prompt_ids = torch.empty(len(prompt_tokens), dtype=torch.int32, device="cpu", pin_memory=True) + prompt_ids.numpy()[:] = prompt_tokens token_id_counter( - prompt_ids=prompt_ids, out_token_id_counter=self.req_to_out_token_id_counter[req.req_idx] + prompt_ids=prompt_ids.cuda(non_blocking=True), + out_token_id_counter=self.req_to_out_token_id_counter[req.req_idx], ) - torch.cuda.current_stream().synchronize() return diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 9a1811a50a..e12039dcf8 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -6,6 +6,7 @@ ) from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import build_kv_indexes_and_input_ids 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 @@ -101,7 +102,7 @@ def _decode(self, model_input: ModelInput) -> ModelOutput: infer_state.position_cos = torch.index_select(self._cos_cached, 0, position_ids) infer_state.position_sin = torch.index_select(self._sin_cached, 0, position_ids) infer_state.mem_manager = self.mem_manager - infer_state.mem_index = self._select_mem_indexes(model_input) + infer_state.mem_index, _ = build_kv_indexes_and_input_ids(model_input, self.req_manager.req_to_token_indexs) # 这里调用的是 DFlash 覆写后的 context_forward：pre layer 只投影 target hidden， # transformer layer 只生成带 RoPE 的 K/V 并写入 draft cache，不会计算 attention、FFN 和 logits。 diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index 71f1c2d5ae..1e8c4605f0 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -1,4 +1,3 @@ -import os import math import ctypes import asyncio diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 43b06f02ef..0fc7f70be6 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -30,7 +30,7 @@ from lightllm.server.core.objs.out_token_circlequeue import LIGHTLLM_OUT_TOKEN_QUEUE_SIZE from lightllm.server.core.objs.io_objs import GroupReqObjs from lightllm.server.core.objs.shm_req_manager import ShmReqManager -from lightllm.server.core.objs.atomic_array_lock import AtomicShmArrayLock, AsyncLock, AtomicLockItem +from lightllm.server.core.objs.atomic_array_lock import AtomicShmArrayLock, AsyncLock from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt from lightllm.utils.log_utils import init_logger from lightllm.server.metrics.manager import MetricClient diff --git a/lightllm/server/router/batch.py b/lightllm/server/router/batch.py index e8d1dca23b..551afd77de 100644 --- a/lightllm/server/router/batch.py +++ b/lightllm/server/router/batch.py @@ -1,6 +1,5 @@ import time -from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple, Union +from typing import List from lightllm.server.core.objs import ShmReqManager, Req from lightllm.utils.log_utils import init_logger from .stats import RouterStatics diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py index 189d3b6b78..942b60751f 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -10,10 +10,9 @@ import zmq.asyncio import torch.multiprocessing as mp import torch.distributed as dist -import multiprocessing from typing import List from .batch import Batch, Req -from .model_infer.model_rpc import start_model_process, ModelRpcClient +from .model_infer.model_rpc import start_model_process from .req_queue import build_req_queue from lightllm.server.core.objs.io_objs import ( GroupReqIndexes, diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 35fb8acda4..307b1a2103 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -6,20 +6,17 @@ import pickle from sortedcontainers import SortedDict -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union from lightllm.common.req_manager import ReqManager, HybridAttentionReqManager -from lightllm.utils.infer_utils import mark_start, mark_end -from lightllm.server.core.objs import Req, SamplingParams, FinishStatus, ShmReqManager +from lightllm.server.core.objs import Req, FinishStatus, ShmReqManager from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import ( HybridAttPagedRadixCache, HybridAttPagedTreeNode, ) from lightllm.utils.log_utils import init_logger -from lightllm.server.httpserver.req_id_generator_server import convert_sub_id_to_group_id from lightllm.server.multimodal_params import MultimodalParams -from lightllm.utils.custom_kernel_utis import custom_cat from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.multimodal_chunk_utils import adjust_mm_chunk_end, get_mm_input_spans from lightllm.server.pd_io_struct import PDDecodeNodeInfo @@ -113,6 +110,7 @@ def get_cpu_kv_cache_stream(self) -> torch.cuda.Stream: def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: bool = True) -> List["InferReq"]: req_objs = [] request_ids = [] + req_prefix_lens = {} for r in requests: r_id, r_index, multimodal_params, dp_rank_in_node = r assert r_id not in self.requests_mapping.keys() @@ -124,11 +122,13 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: multimodal_params=multimodal_params, vocab_size=self.vocab_size, init_prefix_cache=init_prefix_cache, + req_prefix_lens=req_prefix_lens, ) self.requests_mapping[r_id] = r_obj request_ids.append(r_id) req_objs.append(r_obj) + self.req_manager.copy_token_index_prefixes_to_gpu(req_prefix_lens) self.infer_req_ids.extend(request_ids) # diverse mode 下，建立一组请求间的主从关系 @@ -151,25 +151,25 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: return req_objs - def free_a_req_mem(self, free_token_index: List, req: "InferReq"): + def free_a_req_mem(self, free_token_index_blocks: List, req: "InferReq"): if self.radix_cache is None: - free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][0 : req.hold_kv_len]) + free_token_index_blocks.append(self.req_manager.get_cpu_token_index_blocks(req.req_idx, 0, req.hold_kv_len)) elif CacheTier.GPU not in req.cache_tiers: - self._free_req_mem_without_radix_insert(free_token_index=free_token_index, req=req) + self._free_req_mem_without_radix_insert(free_token_index_blocks=free_token_index_blocks, req=req) + elif not self.is_hybrid_att_model: + self._full_att_free_req(free_token_index_blocks=free_token_index_blocks, req=req) else: - if not self.is_hybrid_att_model: - self._full_att_free_req(free_token_index=free_token_index, req=req) - else: - self._hybrid_att_free_req(free_token_index=free_token_index, req=req) - assert len(req.hybrid_len_to_big_page_id) == 0 + self._hybrid_att_free_req(free_token_index_blocks=free_token_index_blocks, req=req) + assert len(req.hybrid_len_to_big_page_id) == 0 req.cur_kv_len = 0 req.hold_kv_len = 0 req.shm_req.shm_cur_kv_len = req.cur_kv_len - return - def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "InferReq"): + def _free_req_mem_without_radix_insert(self, free_token_index_blocks: List, req: "InferReq"): shared_kv_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len - free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][shared_kv_len : req.hold_kv_len]) + free_token_index_blocks.append( + self.req_manager.get_cpu_token_index_blocks(req.req_idx, shared_kv_len, req.hold_kv_len) + ) if self.is_hybrid_att_model: # 释放请求尾部尚未移交给 radix cache 的 hybrid attention 小页状态。 @@ -185,29 +185,32 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer if req.shared_kv_node is not None: self.radix_cache.dec_node_ref_counter(req.shared_kv_node) req.shared_kv_node = None - return - def _full_att_free_req(self, free_token_index: List, req: "InferReq"): + def _full_att_free_req(self, free_token_index_blocks: List, req: "InferReq"): page_size = self.args.page_size cache_kv_len = req.cur_kv_len // page_size * page_size - # radix cache 只保存完整页，因此不足一页的有效 KV 和页内预留空间都不会参与插入。 - # 先释放截断后的 [cache_kv_len, hold_kv_len) 区间，再处理可缓存前缀的插入和去重。 - free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][cache_kv_len : req.hold_kv_len]) - input_token_ids = req.get_input_token_ids() - key = torch.tensor(input_token_ids[0:cache_kv_len], dtype=torch.int64, device="cpu") - # .cpu() 是 流内阻塞操作 - value = self.req_manager.req_to_token_indexs[req.req_idx][:cache_kv_len].detach().cpu() - - prefix_len, _ = self.radix_cache.insert(key, value) + self._cache_req_prefix(free_token_index_blocks, req, cache_kv_len) + + def _cache_req_prefix(self, free_token_index_blocks: List, req: "InferReq", cache_kv_len: int, **cache_metadata): + """Cache complete KV pages and release the tail, duplicate pages and shared reference.""" + indexes = self.req_manager.get_cpu_token_indexes(req.req_idx, 0, cache_kv_len) + free_token_index_blocks.append( + self.req_manager.get_cpu_token_index_blocks(req.req_idx, cache_kv_len, req.hold_kv_len) + ) + key = torch.tensor(req.get_input_token_ids()[:cache_kv_len], dtype=torch.int64, device="cpu") + # The request row may be reused while the cached prefix is still alive. + prefix_len, _ = self.radix_cache.insert(key, indexes, **cache_metadata) old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len - free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][old_prefix_len:prefix_len]) + free_token_index_blocks.append( + self.req_manager.get_cpu_token_index_blocks(req.req_idx, old_prefix_len, prefix_len) + ) if req.shared_kv_node is not None: assert req.shared_kv_node.node_prefix_total_len <= prefix_len self.radix_cache.dec_node_ref_counter(req.shared_kv_node) req.shared_kv_node = None return - def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): + def _hybrid_att_free_req(self, free_token_index_blocks: List, req: "InferReq"): assert g_infer_context.is_hybrid_att_model is True args = get_env_start_args() hash_page_size = args.linear_att_hash_page_size @@ -227,83 +230,32 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): if req.hybrid_cache_len <= req.cur_kv_len and req.tail_small_page_buffer_id is not None: # 只有小页可以有 tail_small_page_buffer_id，然后进行小页插入。 assert page_num % big_page_num != 0 - free_token_index.append( - self.req_manager.req_to_token_indexs[req.req_idx][req.hybrid_cache_len : req.hold_kv_len] - ) - req.cur_kv_len = req.hybrid_cache_len - input_token_ids = req.get_input_token_ids() - key = torch.tensor(input_token_ids[0 : req.cur_kv_len], dtype=torch.int64, device="cpu") - value = self.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].detach().cpu() - block_hashs = req.shm_req.hybrid_token_hash_list.get_all()[:page_num] - state_idxs = [None for _ in range(page_num)] - state_idxs[-1] = req.tail_small_page_buffer_id - req.tail_small_page_buffer_id = None - prefix_len, _ = self.radix_cache.insert( - key, - value, - block_hashs=block_hashs, - block_state_idxs=state_idxs, - len_to_big_page_id=req.hybrid_len_to_big_page_id, - ) - old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len - free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][old_prefix_len:prefix_len]) - if req.shared_kv_node is not None: - assert req.shared_kv_node.node_prefix_total_len <= prefix_len - self.radix_cache.dec_node_ref_counter(req.shared_kv_node) - req.shared_kv_node = None - return - - if shared_kv_len < tail_big_page_token_num <= req.cur_kv_len: - free_token_index.append( - self.req_manager.req_to_token_indexs[req.req_idx][tail_big_page_token_num : req.hold_kv_len] - ) - req.cur_kv_len = tail_big_page_token_num - + cache_kv_len = req.hybrid_cache_len + elif shared_kv_len < tail_big_page_token_num <= req.cur_kv_len: assert req.tail_small_page_buffer_id is None - input_token_ids = req.get_input_token_ids() - key = torch.tensor(input_token_ids[0 : req.cur_kv_len], dtype=torch.int64, device="cpu") - value = self.req_manager.req_to_token_indexs[req.req_idx][: req.cur_kv_len].detach().cpu() - cur_page_num = tail_big_page_token_num // hash_page_size assert tail_big_page_token_num % hash_page_size == 0 - block_hashs = req.shm_req.hybrid_token_hash_list.get_all()[:cur_page_num] - state_idxs = [None for _ in range(cur_page_num)] - prefix_len, _ = self.radix_cache.insert( - key, - value, - block_hashs=block_hashs, - block_state_idxs=state_idxs, - len_to_big_page_id=req.hybrid_len_to_big_page_id, - ) - old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len - free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][old_prefix_len:prefix_len]) - if req.shared_kv_node is not None: - assert req.shared_kv_node.node_prefix_total_len <= prefix_len - self.radix_cache.dec_node_ref_counter(req.shared_kv_node) - req.shared_kv_node = None - return - - if shared_kv_len <= req.cur_kv_len: - free_token_index.append(self.req_manager.req_to_token_indexs[req.req_idx][shared_kv_len : req.hold_kv_len]) - # 该分支不会把 prefill 阶段累积的 big page id 插入 radix cache（典型为 pause/abort - # 在 prefill 跨过 big page 边界后、到达末尾前触发），需在此显式释放，避免泄漏。 - - # 释放本请求 prefill 阶段在 big page 边界上申请、但尚未插入 radix cache 的 big page - # state buffer。仅当请求未走 insert 分支(小页/大页插入)就被释放时才会有残留，典型场景： - # big page 模式下请求在 prefill 跨过 big page 边界后、到达末尾前被 pause / abort。 - # 若不释放，会泄漏 big page state slot，并触发 free_a_req_mem 中 dict 为空的断言。 - if req.hybrid_len_to_big_page_id: - self.radix_cache.big_page_buffers.free_state_cache(list(req.hybrid_len_to_big_page_id.values())) - req.hybrid_len_to_big_page_id.clear() - - req.cur_kv_len = shared_kv_len + cache_kv_len = tail_big_page_token_num + elif shared_kv_len <= req.cur_kv_len: assert req.tail_small_page_buffer_id is None - if req.shared_kv_node is not None: - assert req.shared_kv_node.node_prefix_total_len == req.cur_kv_len - self.radix_cache.dec_node_ref_counter(req.shared_kv_node) - req.shared_kv_node = None + req.cur_kv_len = shared_kv_len + self._free_req_mem_without_radix_insert(free_token_index_blocks, req) return - - assert False, f"error state: cur_kv_len: {req.cur_kv_len}" + else: + assert False, f"error state: cur_kv_len: {req.cur_kv_len}" + + req.cur_kv_len = cache_kv_len + page_num = cache_kv_len // hash_page_size + state_idxs = [None] * page_num + state_idxs[-1] = req.tail_small_page_buffer_id + req.tail_small_page_buffer_id = None + self._cache_req_prefix( + free_token_index_blocks, + req, + cache_kv_len, + block_hashs=req.shm_req.hybrid_token_hash_list.get_all()[:page_num], + block_state_idxs=state_idxs, + len_to_big_page_id=req.hybrid_len_to_big_page_id, + ) def _save_promptcache_kvbuffer(self): """ @@ -327,7 +279,7 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool should_modify_shm = modify_shm_finish_state and self.backend.is_master_in_dp free_req_index = [] - free_token_index = [] + free_token_index_blocks = [] for request_id in finished_request_ids: req: InferReq = self.requests_mapping.pop(request_id) if self.args.diverse_mode: @@ -336,7 +288,7 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool if should_modify_shm: req.final_token_metadata.dump() - self.free_a_req_mem(free_token_index, req) + self.free_a_req_mem(free_token_index_blocks, req) free_req_index.append(req.req_idx) # logger.info(f"infer release req id {req.shm_req.request_id}") @@ -346,9 +298,9 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool req.shm_req.shm_infer_released = True self.shm_req_manager.put_back_req_obj(req.shm_req) - if free_token_index: - free_token_index = custom_cat(free_token_index) - self.req_manager.free(free_req_index, free_token_index) + if free_token_index_blocks: + free_token_index_blocks = np.concatenate(free_token_index_blocks) + self.req_manager.free(free_req_index, free_token_index_blocks) finished_req_ids_set = set(finished_request_ids) self.infer_req_ids = [_id for _id in self.infer_req_ids if _id not in finished_req_ids_set] @@ -362,8 +314,6 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool f"mem manager total size {self.req_manager.mem_manager.allocator.size}\n" ) - return - def filter_reqs(self, finished_reqs: List["InferReq"]): if finished_reqs: self._filter([req.req_id for req in finished_reqs]) @@ -372,13 +322,12 @@ def filter_reqs(self, finished_reqs: List["InferReq"]): @torch.no_grad() def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): if pause_reqs: - - free_token_index = [] + free_token_index_blocks = [] for req in pause_reqs: if self.args.diverse_mode: # 发生暂停的时候，需要清除 diverse 模式下的主从关系 req.clear_master_slave_state() - self.free_a_req_mem(free_token_index, req) + self.free_a_req_mem(free_token_index_blocks, req) assert req.wait_pause is True req.wait_pause = False req.paused = True @@ -386,27 +335,28 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): req.shm_req.is_paused = True logger.debug(f"infer paused req id {req.req_id}") - if len(free_token_index) != 0: - free_token_index = custom_cat(free_token_index) - self.req_manager.free_token(free_token_index) + if free_token_index_blocks: + free_token_index_blocks = np.concatenate(free_token_index_blocks) + self.req_manager.free_token_index_blocks(free_token_index_blocks) return self def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool): if paused_reqs: # pause_reqs 可能刚刚释放了 KV，因此恢复前重新读取实时可用容量。 can_alloc_token_num = self.get_can_alloc_token_num() + req_prefix_lens = {} for req in paused_reqs: # 暂停恢复保持原有的保守语义：只有当前完整序列所需的 KV 页面都有足够空间时才恢复， # 才允许执行 radix cache 匹配，避免在容量不足时提前改变引用或重建 hybrid cache 状态。 - alloc_token_num = req._kv_cache_alloc_need(req.get_cur_total_len()) + alloc_token_num = self.req_manager.get_need_alloc_token_num(req, req.get_cur_total_len()) if alloc_token_num > can_alloc_token_num: break if g_infer_context.is_hybrid_att_model: - req._hybrid_match_radix_cache() + req._hybrid_match_radix_cache(req_prefix_lens=req_prefix_lens) else: - req._match_radix_cache() + req._match_radix_cache(req_prefix_lens=req_prefix_lens) assert req.paused is True req.paused = False @@ -414,6 +364,7 @@ def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bo req.shm_req.is_paused = False logger.debug(f"infer recover paused req id {req.req_id}") can_alloc_token_num -= alloc_token_num + self.req_manager.copy_token_index_prefixes_to_gpu(req_prefix_lens) return def get_can_alloc_token_num(self): @@ -488,7 +439,7 @@ def __init__( vocab_size: int, ) -> None: self.shm_param = shm_req.sample_params - self.disable_prompt_cache = self.shm_param.disable_prompt_cache + self.disable_prompt_cache = self.shm_param.disable_prompt_cache or self.shm_param.prompt_logprobs >= 0 if self.shm_param.top_k == -1: self.shm_param.top_k = vocab_size @@ -535,6 +486,8 @@ def __init__( multimodal_params: MultimodalParams, vocab_size: int = -1, init_prefix_cache: bool = True, + *, + req_prefix_lens, ): self.args = get_env_start_args() self.req_id = req_id @@ -607,9 +560,9 @@ def __init__( if init_prefix_cache: if g_infer_context.is_hybrid_att_model: - self._hybrid_match_radix_cache() + self._hybrid_match_radix_cache(req_prefix_lens=req_prefix_lens) else: - self._match_radix_cache() + self._match_radix_cache(req_prefix_lens=req_prefix_lens) return def _init_all_state(self): @@ -650,7 +603,14 @@ def _init_all_state(self): return - def _match_radix_cache(self): + def _set_cache_prefix(self, shared_node, indexes, req_prefix_lens): + g_infer_context.req_manager.write_token_index_prefix(self.req_idx, 0, indexes, req_prefix_lens=req_prefix_lens) + self.shared_kv_node = shared_node + self.cur_kv_len = self.hold_kv_len = indexes.numel() + assert self.hold_kv_len % self.args.page_size == 0 + self.shm_req.prompt_cache_len = self.cur_kv_len + + def _match_radix_cache(self, req_prefix_lens): assert ( g_infer_context.is_hybrid_att_model is False ), "current _match_radix_cache does not support hybrid attention models, to do..." @@ -663,19 +623,12 @@ def _match_radix_cache(self): key, update_refs=True, mm_input_spans=self.mm_input_spans ) if share_node is not None: - self.shared_kv_node = share_node - ready_cache_len = share_node.node_prefix_total_len - # 从 cpu 到 gpu 是流内阻塞操作 - g_infer_context.req_manager.req_to_token_indexs[self.req_idx, 0:ready_cache_len] = value_tensor - self.cur_kv_len = int(ready_cache_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换 - self.hold_kv_len = self.cur_kv_len - assert self.hold_kv_len % self.args.page_size == 0 - self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 + self._set_cache_prefix(share_node, value_tensor, req_prefix_lens) self.shm_req.shm_cur_kv_len = self.cur_kv_len return - def _hybrid_match_radix_cache(self): + def _hybrid_match_radix_cache(self, req_prefix_lens): assert ( g_infer_context.is_hybrid_att_model is True ), "current _hybrid_match_radix_cache only supports hybrid attention models, to do..." @@ -697,96 +650,43 @@ def _hybrid_match_radix_cache(self): share_node, kv_len, value_tensor = g_infer_context.radix_cache.match_prefix( key, block_hashs=block_hashs, update_refs=True, mm_input_spans=self.mm_input_spans ) + radix_cache = g_infer_context.radix_cache + req_manager = g_infer_context.req_manager + if share_node is not None and not share_node.is_big_page_node() and not big_page_is_disable: + shared_kv_len = share_node.node_prefix_total_len + cur_big_page_tokens = shared_kv_len // big_page_token_num * big_page_token_num + need_tokens = shared_kv_len - cur_big_page_tokens + if g_infer_context.get_can_alloc_token_num() <= need_tokens: + # 无法重建小页尾部时，只恢复最近且不在图片内部的大页 checkpoint。 + share_node = radix_cache.deref_to_first_big_page_node( + node=share_node, mm_input_spans=self.mm_input_spans + ) + if share_node is not None: + assert share_node.is_big_page_node() + value_tensor = value_tensor[: share_node.node_prefix_total_len] + if share_node is not None: assert self.tail_small_page_buffer_id is None - if share_node.is_big_page_node(): - # 大页匹配 - self.shared_kv_node = share_node - ready_cache_len = share_node.node_prefix_total_len - # 从 cpu 到 gpu 是流内阻塞操作 - g_infer_context.req_manager.req_to_token_indexs[self.req_idx, 0:ready_cache_len] = value_tensor - self.cur_kv_len = int(ready_cache_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换 - self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 - assert self.tail_small_page_buffer_id is None - # 恢复 hybrid checkpoint - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=share_node.big_page_buffer_idx, req=self - ) - else: - # 小页匹配 - if big_page_is_disable: - # 如果 大页本质是被禁用的，可以直接使用小页的匹配结果 - self.shared_kv_node = share_node - ready_cache_len = share_node.node_prefix_total_len - # 从 cpu 到 gpu 是流内阻塞操作 - g_infer_context.req_manager.req_to_token_indexs[self.req_idx, 0:ready_cache_len] = value_tensor - self.cur_kv_len = int(ready_cache_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换 - self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 - assert self.tail_small_page_buffer_id is None - # 恢复 hybrid checkpoint - g_infer_context.req_manager.restore_small_page_state( - req=self, - ) + if share_node.is_big_page_node() or big_page_is_disable: + self._set_cache_prefix(share_node, value_tensor, req_prefix_lens) + if share_node.is_big_page_node(): + req_manager.restore_big_page_state(big_page_buffer_idx=share_node.big_page_buffer_idx, req=self) else: - # 如果 大页本质是被启用的，则需要使用小页的匹配结果, 将小页的kv 复制到的新申请的kv位置，同时释放 - # 对应的小页对应的节点，递归找到对应最近的大叶节点进行返回,然后赋值到req.shared_node 对象上 - shared_kv_len = share_node.node_prefix_total_len - cur_big_page_tokens = (shared_kv_len // big_page_token_num) * big_page_token_num - need_tokens = shared_kv_len - cur_big_page_tokens - radix_cache = g_infer_context.radix_cache - if g_infer_context.get_can_alloc_token_num() > need_tokens: - # 有充足的token 容量时 - assert need_tokens % self.args.page_size == 0 - radix_cache.free_radix_cache_to_get_enough_token(need_token_num=need_tokens) - tail_mems = radix_cache.mem_manager.alloc(need_size=need_tokens) - g_infer_context.req_manager.req_to_token_indexs[ - self.req_idx, 0:cur_big_page_tokens - ] = value_tensor[0:cur_big_page_tokens] - g_infer_context.req_manager.req_to_token_indexs[ - self.req_idx, cur_big_page_tokens:shared_kv_len - ] = tail_mems - - # 将 对应的 value_tensors 中的 kv 数据 拷贝到 tail_mems 中对应的数据去 - radix_cache.mem_manager.operator.copy_mem_to_mem( - value_tensor[cur_big_page_tokens:shared_kv_len], tail_mems - ) - # 尾部 KV 换到新 mem 后，同步拷贝已捕获的 top-k prompt logprobs。 - self.prompt_selected_logprobs.copy_capture_slots_if_needed( - source_indexes=value_tensor[cur_big_page_tokens:shared_kv_len], - destination_indexes=tail_mems, - ) - - self.shared_kv_node = share_node # 只是为了保证 restore_small_page_state 正确调用 - g_infer_context.req_manager.restore_small_page_state( - req=self, - ) - self.shared_kv_node = None - - big_page_shared_node = radix_cache.deref_to_first_big_page_node(node=share_node) - self.shared_kv_node = big_page_shared_node - self.cur_kv_len = int(shared_kv_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换 - self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 - else: - # 没有充足的token 容量时， 直接找到最接近的大页，进行大页恢复 - share_node = radix_cache.deref_to_first_big_page_node( - node=share_node, mm_input_spans=self.mm_input_spans - ) - if share_node is not None: - assert share_node.is_big_page_node() - # 大页匹配 - self.shared_kv_node = share_node - ready_cache_len = share_node.node_prefix_total_len - # 从 cpu 到 gpu 是流内阻塞操作 - g_infer_context.req_manager.req_to_token_indexs[ - self.req_idx, 0:ready_cache_len - ] = value_tensor[0:ready_cache_len] - self.cur_kv_len = int(ready_cache_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换 - self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 - assert self.tail_small_page_buffer_id is None - # 恢复 hybrid checkpoint - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=share_node.big_page_buffer_idx, req=self - ) + req_manager.restore_small_page_state(req=self) + else: + # 小页尾部复制到本请求独占的页；先恢复 checkpoint，再释放小页引用。 + assert need_tokens % self.args.page_size == 0 + radix_cache.free_radix_cache_to_get_enough_token(need_token_num=need_tokens) + tail_mems = radix_cache.mem_manager.alloc(need_size=need_tokens) + indexes = torch.cat([value_tensor[:cur_big_page_tokens], tail_mems]) + radix_cache.mem_manager.operator.copy_mem_to_mem(value_tensor[cur_big_page_tokens:], tail_mems) + self.prompt_selected_logprobs.copy_capture_slots_if_needed( + source_indexes=value_tensor[cur_big_page_tokens:], destination_indexes=tail_mems + ) + self.shared_kv_node = share_node + req_manager.restore_small_page_state(req=self) + share_node = radix_cache.deref_to_first_big_page_node(node=share_node) + self._set_cache_prefix(share_node, indexes, req_prefix_lens) self.shm_req.shm_cur_kv_len = self.cur_kv_len @@ -954,39 +854,17 @@ def _stop_sequences_matched(self, output_len: int): return True return False - def prefill_need_token_num(self, is_chuncked_prefill: bool) -> Tuple[int, int]: - """返回 (本轮需要计算的 token 数, 需要额外分配的 KV token 容量)。 - - 第一个值为目标 KV 长度减去已计算的 cur_kv_len，用于限制 batch 的计算量。 - 第二个值先将目标 KV 长度按 page_size 向上对齐，再扣除已持有的 hold_kv_len， - 用于检查和分配 KV 空间。已有页的剩余容量足够时，即使仍需计算 token，第二个值也可以为 0。 - """ + def prefill_need_token_num(self, is_chuncked_prefill: bool) -> int: + """本轮计算量；调度通过计算预算检查后再检查 KV 容量。""" if is_chuncked_prefill: target_kv_len = self.get_chuncked_input_token_len() else: target_kv_len = self.get_cur_total_len() - token_num = target_kv_len - self.cur_kv_len - alloc_token_num = self._kv_cache_alloc_need(target_kv_len) - return token_num, alloc_token_num - - def decode_need_token_num(self) -> Tuple[int, int]: - """返回 (decode 需要的 token 数, 需要额外分配的 KV token 容量)。 + return target_kv_len - self.cur_kv_len - 第一个值沿用分页分配前的需求估算：普通 decode 为 1，MTP 为 2 * (mtp_step + 1)， - 其中 MTP 的两倍系数包含原有的预留窗口，并非本轮实际接受的输出 token 数。 - 第二个值将 cur_kv_len 加上上述需求后按 page_size 向上对齐，再扣除已持有的 hold_kv_len， - 用于检查和分配 KV 空间。已有页的剩余容量足够时，第二个值为 0，仍可继续 decode。 - """ - decode_token_num = 1 if self.mtp_step == 0 else 2 * (1 + self.mtp_step) - alloc_token_num = self._kv_cache_alloc_need(self.cur_kv_len + decode_token_num) - return decode_token_num, alloc_token_num - - def _kv_cache_alloc_need(self, target_kv_len: int) -> int: - page_size = self.args.page_size - target_hold_len = (target_kv_len + page_size - 1) // page_size * page_size - alloc_token_num = max(target_hold_len - self.hold_kv_len, 0) - assert alloc_token_num % page_size == 0 - return alloc_token_num + def decode_need_token_num(self) -> int: + """普通 decode 为 1；MTP 保留原有的双窗口容量。""" + return 1 if self.mtp_step == 0 else 2 * (1 + self.mtp_step) class InferReqUpdatePack: 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..1ff50a82ac 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -5,7 +5,7 @@ import time import threading import torch.distributed as dist -from typing import List, Tuple, Callable, Optional, Union +from typing import List, Tuple, Callable, Optional from transformers.configuration_utils import PretrainedConfig from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.log_utils import init_logger @@ -413,11 +413,11 @@ def _capture_prompt_logprobs_if_needed( return mgr = PromptLogprobsCaptureManager.get_instance() - mem_indexes = self.model._select_mem_indexes(model_input) + mem_indexes = model_input.mem_indexes start_loc = 0 for req_obj in run_reqs: - q_len, _ = req_obj.prefill_need_token_num(is_chuncked_prefill=not self.disable_chunked_prefill) + q_len = req_obj.prefill_need_token_num(is_chuncked_prefill=not self.disable_chunked_prefill) topk = req_obj.sampling_param.shm_param.prompt_logprobs capture_count = min(q_len, req_obj.shm_req.input_len - req_obj.cur_kv_len - 1) if capture_count > 0 and topk == 0 and self.is_master_in_dp: @@ -557,9 +557,7 @@ def _read_pd_trans_io_buffer_and_update_req_status(self): req_to_next_token_ids = ( self.model.req_manager.req_sampling_params_manager.req_to_next_token_ids ) - # to do 这个地方是否需要加流同步 req_to_next_token_ids[req.req_idx, 0:1].fill_(obj.first_gen_token_id) - torch.cuda.current_stream().synchronize() InferReqUpdatePack(req_obj=req, output_len=req.cur_output_len).handle( next_token_id=obj.first_gen_token_id, next_token_logprob=obj.first_gen_token_logprob, @@ -616,29 +614,48 @@ def _timer_merge_radix_tree(self): ) return - # 一些可以复用的通用功能函数 - def _alloc_req_kv_mem( - self, - req_obj: InferReq, - alloc_token_num: int, - no_blcoking_copy: bool = False, - ) -> Optional[torch.Tensor]: - if alloc_token_num == 0: - return None - - assert alloc_token_num > 0 and alloc_token_num % self.args.page_size == 0 - if g_infer_context.radix_cache is not None: - g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(alloc_token_num) - - old_hold_kv_len = req_obj.hold_kv_len - new_hold_kv_len = old_hold_kv_len + alloc_token_num - mem_indexes = g_infer_context.req_manager.mem_manager.alloc(alloc_token_num) - # 高频调度路径允许异步写入请求索引表，其他调用方默认保持原有的同步拷贝语义。 - g_infer_context.req_manager.req_to_token_indexs[req_obj.req_idx, old_hold_kv_len:new_hold_kv_len].copy_( - mem_indexes, non_blocking=no_blcoking_copy - ) - req_obj.hold_kv_len = new_hold_kv_len - return mem_indexes + + + def _classify_reqs_and_alloc_kv(self, candidates, available): + """按优先级顺序接纳请求，再一次性分配所有必需的 token 容量。 + + prefill 和 decode 两组请求都会跨调度轮次保留已分配的 token 容量。 + decode 的可选额外预留仅使用满足所有必要分配后剩余的空闲容量。 + """ + req_manager = g_infer_context.req_manager + batch_max_tokens = self.batch_max_tokens + prefill, decode, rejected, allocations, headroom = [], [], [], [], [] + prefill_tokens = required = 0 + page_size = req_manager.mem_manager.page_size + for req, is_decode, tokens in candidates: + if not is_decode and prefill_tokens + tokens > batch_max_tokens: + continue + + target = req.cur_kv_len + tokens + size = req_manager.get_need_alloc_token_num(req, target) if target > req.hold_kv_len else 0 + if size > available: + rejected.append((req, is_decode)) + continue + + (decode if is_decode else prefill).append(req) + available -= size + if not is_decode: + prefill_tokens += tokens + if size: + allocation = [req, size] + allocations.append(allocation) + required += size + if is_decode and page_size < 8: + headroom.append(allocation) + + extra_budget = max(0, req_manager.mem_manager.allocator.can_use_mem_size - required) + for allocation in headroom: + extra = min(8, extra_budget) // page_size * page_size + allocation[1] += extra + extra_budget -= extra + if allocations: + req_manager.alloc_req_token_indexes(allocations, g_infer_context.radix_cache) + return prefill, decode, rejected, prefill_tokens @staticmethod def _is_decode_req(req: InferReq, no_decode: bool, strict_prefill: bool) -> bool: @@ -710,93 +727,65 @@ def _get_classed_reqs( wait_pause_reqs = [] paused_reqs = [] finished_reqs = [] - prefill_reqs = [] - decode_reqs = [] # 单轮最多处理少量因 token 容量不足而无法继续的请求，避免一次性影响大量请求。 # 普通 Decode 请求进入暂停队列等待恢复；PD Decode 请求则强制提前结束并进入清理流程。 pause_max_req_num = 2 - wait_pause_count = 0 - prefill_tokens = 0 can_alloc_token_num = g_infer_context.get_can_alloc_token_num() - for req_obj in ready_reqs: + def candidates(): + for req_obj in ready_reqs: + if req_obj.filter_mark: + finished_reqs.append(req_obj) + continue - if req_obj.filter_mark: - finished_reqs.append(req_obj) - continue + if req_obj.wait_pause: + wait_pause_reqs.append(req_obj) + continue - if req_obj.wait_pause: - wait_pause_reqs.append(req_obj) - continue + if req_obj.paused: + paused_reqs.append(req_obj) + continue - if req_obj.paused: - paused_reqs.append(req_obj) - continue + if req_obj.infer_aborted or req_obj.finish_status.is_finished(): + if support_overlap: + # 延迟处理 + req_obj.filter_mark = True + continue + else: + finished_reqs.append(req_obj) + continue + + is_decode = self._is_decode_req(req_obj, no_decode=no_decode, strict_prefill=strict_prefill) + + if is_decode: + token_num = req_obj.decode_need_token_num() + else: + # 在 diverse mode 模式下，prefill 只会使用 master 状态的请求，slave 请求依靠后续 + # 的推理代码中将master请求的状态复制到slave请求中去， 所以这里 slave 状态的请求，不 + # 放入到 prefill reqs 队列中，在其他模式下，所有请求都是 master状态，所以也不受影响 + if req_obj.is_slave_req(): + continue + + token_num = req_obj.prefill_need_token_num(is_chuncked_prefill=not self.disable_chunked_prefill) + yield req_obj, is_decode, token_num - if req_obj.infer_aborted or req_obj.finish_status.is_finished(): + prefill_reqs, decode_reqs, rejected, prefill_tokens = self._classify_reqs_and_alloc_kv( + candidates(), can_alloc_token_num + ) + for req_obj, is_decode in rejected[:pause_max_req_num]: + if is_decode and self.args.run_mode == "decode": + req_obj.finished_by_pd_decode_capacity = True if support_overlap: - # 延迟处理 req_obj.filter_mark = True - continue else: finished_reqs.append(req_obj) - continue - - if self._is_decode_req(req_obj, no_decode=no_decode, strict_prefill=strict_prefill): - # KV 容量检查使用额外分配量，已有页的剩余容量可以覆盖部分或全部 decode 需求。 - _, alloc_token_num = req_obj.decode_need_token_num() - # page_size 较小时，decode 会频繁触发 KV 内存分配。此处额外预申请不超过 8 个 token， - # 并将数量向下对齐到 page_size 的整数倍，以减少 alloc 调用次数并保持分页分配约束。 - if alloc_token_num > 0 and self.args.page_size < 8: - alloc_token_num += 8 // self.args.page_size * self.args.page_size - if alloc_token_num <= can_alloc_token_num: - self._alloc_req_kv_mem(req_obj, alloc_token_num, no_blcoking_copy=True) - decode_reqs.append(req_obj) - can_alloc_token_num -= alloc_token_num - else: - if wait_pause_count < pause_max_req_num: - if self.args.run_mode == "decode": - # PD Decode 节点的 token 容量不足时，强制当前请求提前结束以释放资源。 - # 单轮只处理 pause_max_req_num 个请求，避免所有资源不足的请求同时退出。 - wait_pause_count += 1 - setattr(req_obj, "finished_by_pd_decode_capacity", True) - if support_overlap: - # overlap 模式可能仍有异步计算在访问请求，先标记，下一轮再安全清理。 - req_obj.filter_mark = True - else: - # 非 overlap 模式没有在途的异步计算，可以在本轮直接清理。 - finished_reqs.append(req_obj) - self.logger.info( - f"force early finish for PD decode req_id={req_obj.req_id} " - f"because token capacity is insufficient" - ) - else: - req_obj.wait_pause = True - wait_pause_count += 1 - else: - # 在 diverse mode 模式下，prefill 只会使用 master 状态的请求，slave 请求依靠后续 - # 的推理代码中将master请求的状态复制到slave请求中去， 所以这里 slave 状态的请求，不 - # 放入到 prefill reqs 队列中，在其他模式下，所有请求都是 master状态，所以也不受影响 - if req_obj.is_slave_req(): - continue - - # 计算预算按本轮实际处理的 token 数累计，KV 预算按需要额外分配的页容量扣减。 - token_num, alloc_token_num = req_obj.prefill_need_token_num( - is_chuncked_prefill=not self.disable_chunked_prefill + self.logger.info( + f"force early finish for PD decode req_id={req_obj.req_id} because token capacity is insufficient" ) - if prefill_tokens + token_num > self.batch_max_tokens: - continue - if alloc_token_num <= can_alloc_token_num: - self._alloc_req_kv_mem(req_obj, alloc_token_num, no_blcoking_copy=True) - prefill_tokens += token_num - prefill_reqs.append(req_obj) - can_alloc_token_num -= alloc_token_num - else: - if wait_pause_count < pause_max_req_num: - req_obj.wait_pause = True - wait_pause_count += 1 + else: + req_obj.wait_pause = True # 先由控制器确定请求需要写入的缓存层级，再按是否包含 CPU cache 决定是否发起 offload。 cache_controller = g_infer_context.cache_placement_controller @@ -823,10 +812,7 @@ def _get_classed_reqs( g_infer_context.pause_reqs(wait_pause_reqs, is_master_in_dp=self.is_master_in_dp) if recover_paused: - g_infer_context.recover_paused_reqs( - paused_reqs=paused_reqs, - is_master_in_dp=self.is_master_in_dp, - ) + g_infer_context.recover_paused_reqs(paused_reqs=paused_reqs, is_master_in_dp=self.is_master_in_dp) # 在 enable_prefill_decode_mixed 模式下，如果存在 prefill 请求和 decode 请求， # 并且 prefill 请求需要的 token 数量 + decode 请求需要的 token 数量小于等于 batch_max_tokens， diff --git a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py index 46586959f1..7ff3302d96 100644 --- a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py @@ -1,5 +1,4 @@ import torch -from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.server.router.model_infer.infer_batch import ( g_infer_context, InferReq, @@ -40,7 +39,6 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq ) with torch.cuda.stream(g_infer_context.get_overlap_stream()): - model_output = self.model.forward(model_input) logits = model_output.logits @@ -89,8 +87,8 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() - model_output.prefill_mem_indexes_ready_event.synchronize() - update_packs = self._diverse_pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + with torch.cuda.stream(g_infer_context.get_overlap_stream()): + update_packs = self._diverse_pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) # 第三阶段 event_pack.notify_forward_and_wait_post_handle() sync_event.synchronize() @@ -129,6 +127,7 @@ def _diverse_copy( def _diverse_pre_post_handle(self, run_reqs: List[InferReq], is_chuncked_mode: bool) -> List[InferReqUpdatePack]: update_func_objs: List[InferReqUpdatePack] = [] + req_prefix_lens = {} # 通用状态预先填充 is_master_in_dp = self.is_master_in_dp pre_master_req_pack = None @@ -161,30 +160,28 @@ def _diverse_pre_post_handle(self, run_reqs: List[InferReq], is_chuncked_mode: b if req_obj.slave_reqs: # 存在 slave reqs 的 master req 需要将自己的 kv 信息写入到 radix cache 中 # 方便 slave req 进行 kv 的复用 - self._master_req_to_radix_cache(master_req=req_obj) + self._master_req_to_radix_cache(master_req=req_obj, req_prefix_lens=req_prefix_lens) else: # slave req 直接复用 master req 的更新包。 assert pre_master_req_pack is not None assert pre_master_req_pack.req_obj.shm_req.group_req_id == req_obj.shm_req.group_req_id - self._copy_master_req_to_slave_req(slave_req=req_obj) + self._copy_master_req_to_slave_req(slave_req=req_obj, req_prefix_lens=req_prefix_lens) # 在拷贝后，请求独立了，与 master_req 的关系解除 req_obj.remove_master_req() pack = InferReqUpdatePack(req_obj=req_obj, output_len=pre_master_req_pack.output_len) update_func_objs.append(pack) - torch.cuda.current_stream().synchronize() + self.model.req_manager.copy_token_index_prefixes_to_gpu(req_prefix_lens) return update_func_objs - def _master_req_to_radix_cache(self, master_req: InferReq): + def _master_req_to_radix_cache(self, master_req: InferReq, req_prefix_lens): key = master_req.get_input_token_ids()[0 : master_req.cur_kv_len] key = torch.tensor(key, dtype=torch.int64, device="cpu") - value = self.model.req_manager.req_to_token_indexs[master_req.req_idx][: master_req.cur_kv_len].detach().cpu() + value = self.model.req_manager.get_cpu_token_indexes(master_req.req_idx, 0, master_req.cur_kv_len) prefix_len, new_shared_kv_node = self.radix_cache.insert(key, value) old_prefix_len = 0 if master_req.shared_kv_node is None else master_req.shared_kv_node.node_prefix_total_len assert old_prefix_len <= master_req.cur_kv_len - self.model.mem_manager.free( - self.model.req_manager.req_to_token_indexs[master_req.req_idx][old_prefix_len:prefix_len] - ) + self.model.mem_manager.free(value[old_prefix_len:prefix_len]) # 将原有共享节点替换为新共享节点，新共享节点对应的长度为当前的cur_kv_len self.radix_cache.dec_node_ref_counter(master_req.shared_kv_node) @@ -196,10 +193,10 @@ def _master_req_to_radix_cache(self, master_req: InferReq): share_node, kv_len, value = self.radix_cache.match_prefix(key, update_refs=False) assert share_node == new_shared_kv_node and kv_len == master_req.cur_kv_len - self.model.req_manager.req_to_token_indexs[master_req.req_idx][0 : master_req.cur_kv_len] = value + self.model.req_manager.write_token_index_prefix(master_req.req_idx, 0, value, req_prefix_lens=req_prefix_lens) return - def _copy_master_req_to_slave_req(self, slave_req: InferReq): + def _copy_master_req_to_slave_req(self, slave_req: InferReq, req_prefix_lens): master_req = slave_req.related_master_req assert master_req is not None @@ -208,10 +205,11 @@ def _copy_master_req_to_slave_req(self, slave_req: InferReq): kv_len = master_req.cur_kv_len - self.model.req_manager.req_to_token_indexs[slave_req.req_idx][ - 0:kv_len - ] = self.model.req_manager.req_to_token_indexs[master_req.req_idx][0:kv_len] - # torch.cuda.current_stream().synchronize() + req_manager = self.model.req_manager + # Radix insertion may remap the master prefix; publish it before the GPU copy. + req_manager.copy_token_index_prefixes_to_gpu(req_prefix_lens) + req_prefix_lens.clear() + req_manager.copy_token_index_prefix(master_req.req_idx, slave_req.req_idx, kv_len) slave_req.shared_kv_node = master_req.shared_kv_node slave_req.cur_kv_len = kv_len # slave 在 prefill 阶段没有独立分配 KV，共享完成后需要同步其已持有的 KV 长度。 diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index 413a86b94e..645ad5e5ad 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -1,5 +1,4 @@ # 该文件用于提供在数据dp并行的推理模式下，共享kv cache trans相关的功能函数模块 -import time import numpy as np import dataclasses import torch @@ -9,7 +8,6 @@ from lightllm.utils.dist_utils import get_dp_rank_in_node from lightllm.server.core.objs.shm_array import ShmArray from ...infer_batch import InferReq -from lightllm.utils.dist_utils import get_current_device_id from lightllm.server.router.model_infer.infer_batch import g_infer_context import torch.distributed as dist @@ -71,15 +69,14 @@ def build_shared_kv_trans_tasks( # 计算需要传输的 kv 长度， 不能超过 req.get_cur_total_len() - 1 trans_size = min(max_req_radix_cache_len, req.get_cur_total_len() - 1) - req.cur_kv_len target_kv_len = req.cur_kv_len + trans_size - alloc_token_num = req._kv_cache_alloc_need(target_kv_len) if trans_size > 0 else 0 + alloc_token_num = ( + g_infer_context.req_manager.get_need_alloc_token_num(req, target_kv_len) if trans_size > 0 else 0 + ) if is_current_dp_handle and trans_size > 0 and alloc_token_num <= g_infer_context.get_can_alloc_token_num(): - assert req.hold_kv_len == req.cur_kv_len - mem_indexes = self.backend._alloc_req_kv_mem(req, alloc_token_num) - assert mem_indexes is not None - # mem_indexes 只描述需要复制的逻辑 KV；页尾预留槽位已经由 - # _alloc_req_kv_mem 写入请求表，但不参与本次跨 DP 传输。 - mem_indexes = mem_indexes[:trans_size] + mem_indexes = g_infer_context.req_manager.alloc_token_indexes( + req, target_kv_len, g_infer_context.radix_cache + ) max_kv_len_dp_rank = self.shared_req_infos.arr[req_index, :, self._KV_LEN_INDEX].argmax() max_kv_len_req_idx = int(self.shared_req_infos.arr[req_index, max_kv_len_dp_rank, self._REQ_IDX_INDEX]) max_kv_len_mem_manager_index = max_kv_len_dp_rank * self.backend.dp_world_size + self.backend.rank_in_dp diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index 7d3f34c4a7..34c876f624 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -1,14 +1,13 @@ import torch import numpy as np from typing import List, Tuple -from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context +from lightllm.server.router.model_infer.infer_batch import InferReq from lightllm.common.basemodel.batch_objs import ModelInput INT64_MAX = torch.iinfo(torch.int64).max def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> Tuple[ModelInput, List[InferReq]]: - run_reqs = [] total_token_num = 0 input_ids = [] b_req_idx = [] @@ -16,12 +15,11 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> b_q_seq_len = [] batch_multimodal_params = [] b_ready_cache_len = [] - b_mtp_index = [] + b_prefill_start_loc = [0] b_prefill_has_output = [] b_is_decode_req = [] for req in req_objs: - run_reqs.append(req) batch_multimodal_params.append(req.multimodal_params) b_req_idx.append(req.req_idx) @@ -30,7 +28,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> else: input_token_ids = req.get_input_token_ids() - b_prefill_has_output.append(False if len(input_token_ids) < req.get_cur_total_len() else True) + b_prefill_has_output.append(len(input_token_ids) >= req.get_cur_total_len()) seq_len = len(input_token_ids) input_token_len = seq_len - req.cur_kv_len @@ -39,10 +37,10 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> b_seq_len.append(seq_len) b_q_seq_len.append(input_token_len) + b_prefill_start_loc.append(b_prefill_start_loc[-1] + input_token_len) input_ids.append(input_id) total_token_num += seq_len b_ready_cache_len.append(req.cur_kv_len) - b_mtp_index.append(0) if hasattr(req, "is_decode_req_mixed_in_prefill"): b_is_decode_req.append(True) del req.is_decode_req_mixed_in_prefill @@ -55,15 +53,17 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> max_cache_len = max(b_ready_cache_len, default=0) max_q_seq_len = max(b_q_seq_len, default=0) - input_ids = np.concatenate(input_ids, dtype=np.int64) if input_ids else np.empty((0,), dtype=np.int64) - input_ids = torch.tensor(input_ids, dtype=torch.int64, device="cpu") - b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu") - b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu") - b_is_decode_req = torch.tensor(b_is_decode_req, dtype=torch.bool, device="cpu") - b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu") - b_ready_cache_len = torch.tensor(b_ready_cache_len, dtype=torch.int32, device="cpu") - b_q_seq_len = torch.tensor(b_q_seq_len, dtype=torch.int32, device="cpu") - b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len + token_ids = torch.empty(b_prefill_start_loc[-1], dtype=torch.int64, device="cpu", pin_memory=True) + if input_ids: + np.concatenate(input_ids, out=token_ids.numpy()) + input_ids = token_ids + b_req_idx, b_seq_len, b_mtp_index, b_ready_cache_len, b_prefill_start_loc = torch.from_numpy( + np.asarray( + (b_req_idx, b_seq_len, [0] * len(req_objs), b_ready_cache_len, b_prefill_start_loc[:-1]), + dtype=np.int32, + ) + ).unbind() + b_is_decode_req = torch.tensor(b_is_decode_req, dtype=torch.bool, device="cpu", pin_memory=True) model_input = ModelInput( batch_size=b_seq_len.shape[0], @@ -83,7 +83,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> multimodal_params=batch_multimodal_params, ) - return model_input, run_reqs + return model_input, list(req_objs) def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[InferReq]]: @@ -92,37 +92,25 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_req_idx = [] b_mtp_index = [] b_seq_len = [] - b_q_seq_len = [] multimodal_params = [] for req in req_objs: - run_reqs.append(req) - b_req_idx.append(req.req_idx) seq_len = req.get_cur_total_len() assert req.cur_kv_len == seq_len - 1, f"{req.cur_kv_len} {seq_len}" - b_seq_len.append(seq_len) - b_q_seq_len.append(1) - total_token_num += seq_len - b_mtp_index.append(0) - multimodal_params.append(req.multimodal_params) - # process the draft tokens. - for step in range(req.mtp_step): + for step in range(req.mtp_step + 1): run_reqs.append(req) b_req_idx.append(req.req_idx) - seq_len += 1 - b_seq_len.append(seq_len) - total_token_num += seq_len - b_mtp_index.append(step + 1) + b_seq_len.append(seq_len + step) + total_token_num += seq_len + step + b_mtp_index.append(step) multimodal_params.append(req.multimodal_params) - b_q_seq_len.append(1) # 空 DP rank 同样构建完整的 decode ModelInput；BaseModel 会在 token # gather 和 attention 初始化之前补入内部 dummy request。 max_kv_seq_len = max(b_seq_len, default=0) - max_q_seq_len = max(b_q_seq_len, default=1) - b_req_idx = torch.tensor(b_req_idx, dtype=torch.int32, device="cpu") - b_seq_len = torch.tensor(b_seq_len, dtype=torch.int32, device="cpu") - b_mtp_index = torch.tensor(b_mtp_index, dtype=torch.int32, device="cpu") + b_req_idx, b_seq_len, b_mtp_index = torch.from_numpy( + np.asarray((b_req_idx, b_seq_len, b_mtp_index), dtype=np.int32) + ).unbind() b_position_delta = build_b_position_delta(multimodal_params) b_shared_seq_len = torch.tensor( @@ -132,12 +120,13 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In [-1 if req.shared_kv_node is None else req.shared_kv_node.time_id % INT64_MAX for req in run_reqs], dtype=torch.int64, device="cpu", + pin_memory=True, ) model_input = ModelInput( batch_size=b_seq_len.shape[0], total_token_num=total_token_num, - max_q_seq_len=max_q_seq_len, + max_q_seq_len=1, max_kv_seq_len=max_kv_seq_len, input_ids=None, b_req_idx=b_req_idx, diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 736e52f439..b8a4d8636e 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -1,4 +1,3 @@ -import threading import torch.distributed as dist import torch import dataclasses @@ -8,11 +7,9 @@ from collections import deque from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient -from lightllm.utils.config_utils import is_hybrid_att_model from lightllm.utils.envs_utils import get_env_start_args from ..infer_batch import InferReq from lightllm.utils.dist_utils import create_new_group_for_current_dp -from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu, load_cpu_kv_to_gpu from lightllm.server.router.model_infer.infer_batch import g_infer_context from lightllm.utils.log_utils import init_logger from lightllm.utils.multimodal_chunk_utils import clamp_mm_cache_len @@ -105,7 +102,6 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): if need_token_num >= 128 and req.shm_req.input_len >= 256: assert req.cur_kv_len % self.args.page_size == 0 assert match_tokens % self.args.page_size == 0 - assert need_token_num % self.args.page_size == 0 assert req.hold_kv_len == req.cur_kv_len if need_token_num <= idle_token_num: # 计算需要加载的页面（只加载未匹配的部分） @@ -113,8 +109,9 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): assert ready_page_num <= matched_page_num need_pages = page_list[ready_page_num:matched_page_num] - mem_indexes = self.backend._alloc_req_kv_mem(req, need_token_num) - assert mem_indexes is not None + mem_indexes = g_infer_context.req_manager.alloc_token_indexes( + req, match_tokens, g_infer_context.radix_cache + ) if self.need_sync_compute_stream(): # TODO fa3 现在必须使用同步模式, 未来需要移除 @@ -124,8 +121,8 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): mem_manager = self.backend.model.mem_manager req_manager = self.backend.model.req_manager - mem_indexes_cuda = mem_indexes.cuda(non_blocking=True) - page_indexes_cuda = torch.tensor(need_pages, dtype=torch.int32, device="cpu").cuda( + mem_indexes_cuda = mem_indexes.pin_memory().cuda(non_blocking=True) + page_indexes_cuda = torch.tensor(need_pages, dtype=torch.int32, device="cpu", pin_memory=True).cuda( non_blocking=True ) # hybrid 页面加载必须按完整 page 处理，否则可能缺失恢复运行态所需的 checkpoint，所以 @@ -286,6 +283,7 @@ def _start_kv_cache_offload_task( move_token_num = page_len_list[item_size - 1] assert req.cur_kv_len >= move_token_num + self.backend.model.req_manager.wait_token_indexes() token_indexes = self.backend.model.req_manager.req_to_token_indexs[req.req_idx, 0:move_token_num] mem_manager = self.backend.model.mem_manager diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py index 786c60b30a..476eeda68f 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py @@ -1,4 +1,3 @@ -import random import torch.multiprocessing as mp from lightllm.server.pd_io_struct import PDChunckedTransTask, PDChunckedTransTaskGroup, PDAbortReq from lightllm.server.router.model_infer.mode_backend.chunked_prefill.impl import ChunkedPrefillBackend @@ -102,16 +101,12 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: # 提前释放有问题的 mem_index。cur_kv_len 只表示已经传输完成的 # 逻辑长度，hold_kv_len 还包含最后一个页面中尚未使用的预留槽位。 old_prefix_len = 0 if req_obj.shared_kv_node is None else req_obj.shared_kv_node.node_prefix_total_len - error_mem_len = req_obj.hold_kv_len - old_prefix_len - if error_mem_len > 0: - mem_indexes = ( - self.model.req_manager.req_to_token_indexs[ - req_obj.req_idx, old_prefix_len : req_obj.hold_kv_len - ] - .detach() - .cpu() + if req_obj.hold_kv_len > old_prefix_len: + self.model.req_manager.free_token_index_blocks( + self.model.req_manager.get_cpu_token_index_blocks( + req_obj.req_idx, old_prefix_len, req_obj.hold_kv_len + ) ) - self.model.mem_manager.free(mem_indexes) req_obj.cur_kv_len = old_prefix_len req_obj.hold_kv_len = old_prefix_len if self.is_master_in_dp: @@ -131,14 +126,7 @@ def _decode_node_gen_trans_tasks(self, req_obj: InferReq): trans_page_size = self.args.pd_kv_page_size assert trans_page_size % self.args.page_size == 0, "pd_kv_page_size must be divisible by page_size" req_obj.pd_trans_kv_start_index = req_obj.cur_kv_len - assert req_obj.hold_kv_len == req_obj.cur_kv_len - need_mem_size = req_obj._kv_cache_alloc_need(input_len) - - mem_indexes = self._alloc_req_kv_mem(req_obj, need_mem_size) - assert mem_indexes is not None - # 传输只覆盖真实 KV；最后一个模型页面中尚未使用的部分继续留在 - # req_to_token_indexs 中，供后续 decode 直接切片使用。 - mem_indexes = mem_indexes[: input_len - req_obj.cur_kv_len] + mem_indexes = self.model.req_manager.alloc_token_indexes(req_obj, input_len, self.radix_cache) while req_obj.pd_trans_kv_start_index < input_len: cur_page_size = min(trans_page_size, input_len - req_obj.pd_trans_kv_start_index) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py index 8641e3b69c..0d22f9f9a1 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py @@ -1,6 +1,5 @@ import torch.multiprocessing as mp -import random -from typing import List, Tuple +from typing import List from lightllm.server.router.model_infer.infer_batch import InferReq from lightllm.server.pd_io_struct import PDAbortReq, PDChunckedTransTask from lightllm.utils.log_utils import init_logger @@ -122,12 +121,9 @@ def _create_pd_trans_task( pd_decode_node_info = req_obj.sampling_param.pd_decode_node if page_kind == "kv": - mem_indexes = ( - self.model.req_manager.req_to_token_indexs[req_obj.req_idx, kv_start_index:kv_end_index] - .detach() - .cpu() - .tolist() - ) + mem_indexes = self.model.req_manager.get_cpu_token_indexes( + req_obj.req_idx, kv_start_index, kv_end_index + ).tolist() req_idx = None elif page_kind == "att_state": mem_indexes = [] diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/base.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/base.py index 9a82cf85da..e5cd7d54fd 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/base.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/base.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, Optional import torch diff --git a/lightllm/server/router/req_queue/base_queue.py b/lightllm/server/router/req_queue/base_queue.py index 8d4ac38085..139e9e74a6 100644 --- a/lightllm/server/router/req_queue/base_queue.py +++ b/lightllm/server/router/req_queue/base_queue.py @@ -1,6 +1,5 @@ import time -from typing import List, Dict -from lightllm.utils.infer_utils import calculate_time +from typing import List from ..batch import Batch, Req from lightllm.server.core.objs import FinishStatus from lightllm.utils.config_utils import get_fixed_kv_len, get_running_max_req_size_per_dp diff --git a/lightllm/server/router/req_queue/dp_base_queue.py b/lightllm/server/router/req_queue/dp_base_queue.py index 2d5c5b039b..f3dbfb7650 100644 --- a/lightllm/server/router/req_queue/dp_base_queue.py +++ b/lightllm/server/router/req_queue/dp_base_queue.py @@ -1,4 +1,3 @@ -import random from typing import List from ..batch import Batch, Req from lightllm.server.router.req_queue.base_queue import BaseQueue diff --git a/lightllm/utils/custom_kernel_utis.py b/lightllm/utils/custom_kernel_utis.py index 1a0b46ca0d..9ec3945bc9 100644 --- a/lightllm/utils/custom_kernel_utis.py +++ b/lightllm/utils/custom_kernel_utis.py @@ -4,28 +4,6 @@ from typing import List -def custom_cat(tensors): - """ - 直接调用 torch 的 cat操作，会造成多个流同步阻塞，用 custom_cat 进行替换。 - 注意返回的 tensor 为 cpu pin_memory 类型, 只会在一些特殊的场景使用。 - """ - if not isinstance(tensors, (list, tuple)): - raise ValueError("Input must be a list of tensors") - - assert tensors[0].is_cuda and len(tensors[0].shape) == 1 - sizes = [t.shape[0] for t in tensors] - dest_size = sum(sizes) - out_tensor = torch.empty((dest_size,), dtype=tensors[0].dtype, device="cpu", pin_memory=True) - - start_loc = 0 - for t, size in zip(tensors, sizes): - out_tensor[start_loc : (start_loc + size)].copy_(t, non_blocking=True) - start_loc += size - torch.cuda.current_stream().synchronize() - - return out_tensor - - def torch_cat_3(tensors: List[torch.Tensor], dim=0): if not tensors: raise ValueError("at least one tensor") diff --git a/skills/test_model/SKILL.md b/skills/test_model/SKILL.md index b77432b6ad..f5880ea268 100644 --- a/skills/test_model/SKILL.md +++ b/skills/test_model/SKILL.md @@ -57,3 +57,43 @@ lm_eval --model local-completions \ | 离线缓存模式 | 约 20s | 因此，除非有明确理由，`skills/test_model` 下的 `lm_eval` 测试都应默认启用离线缓存变量。 + +## Paged KV 全量回归 + +对 paged KV 的修改，先运行 `python -m pytest unit_tests -ra`，再执行本目录全部子 skill 的 +完整精度评测。GSM8K 应包含 1319 题，MMMU validation 应包含 900 题；保留 +`--log_samples` 结果，核对题数、空响应、服务端错误与精度，不能用 `--limit` 替代全量。 + +普通端到端场景增加 `--page_size 16`；CPU/disk cache 场景增加 `--page_size 128`， +并保留子 skill 要求的连续两轮精度评测。CPU cache 的页大小须为物理 page size 的整数倍。 + +额外验证 CPU cache 在 GPU 前缀被逐出后的真实回载。以下三组分别启动独立服务： + +| 模型 | KV 类型 | page_size | CPU cache 页大小 | 额外参数 | +|------|---------|-----------|------------------|----------| +| Qwen3-8B | 默认 | 16 | 128 | 无 | +| Qwen3-8B | int8kv | 128 | 128 | `--llm_kv_type int8kv` | +| Qwen3.5-0.8B | 默认 | 128 | 512 | `--linear_att_cache_size 10 --linear_att_hash_page_size 256 --linear_att_page_block_num 2` | + +每组使用 `--tp 2 --enable_cpu_cache --cpu_cache_storage_size 16 --enable_prompt_logprobs` +以及 `--max_total_token_num 32768 --max_req_total_len 16384 --chunked_prefill_size 257`； +设置表中 `--page_size` 和 `--cpu_cache_token_page_size`。257 特意不整除物理页大小， +用于覆盖 chunk 跨页边界。按子 skill 的方法检查服务日志、端口，并完成真实请求 warmup。 + +服务就绪后执行（变量对应本组实际启动参数）： + +```bash +python test/acc/paged_cpu_cache_probe.py \ + --url "http://127.0.0.1:${PORT}" --model-dir "${MODEL_DIR}" \ + --server-log "${LOG_DIR}/server.log" --output "${LOG_DIR}/cpu-cache-probe.json" \ + --page-size "${PAGE_SIZE}" --cpu-page-size "${CPU_PAGE_SIZE}" --kv-capacity 32768 +``` + +Qwen3.5 的探针额外传入 `--hash-page-size 256`，与服务端 hybrid hash page 保持一致。 + +探针要求回载日志包含实际 CPU cache 命中，并覆盖部分 GPU 前缀同时命中的情况； +GPU 常驻前缀参考与 CPU 回载使用相同的 prefill 起点，避免 chunk/query 形状改变产生的 +BF16 舍入差异干扰传输校验；32 个输出 token 必须相同，logprob 最大绝对差不超过 0.02。 +另外检查 prompt logprobs 禁用缓存复用后仍返回全部输入位置的有效结果。 +GPU↔CPU KV 逐值一致性、碎片页和异步引用生命周期由 +`unit_tests/server/router/model_infer/mode_backend/test_paged_cpu_cache.py` 覆盖。 diff --git a/test/advanced_config/fp8_calibration_per_head/test_kv_cache_calib_per_head_qwen2.5_14b.json b/test/advanced_config/fp8_calibration_per_head/test_kv_cache_calib_per_head_qwen2.5_14b.json index 4dec53ae4d..ea4bfd3f18 100644 --- a/test/advanced_config/fp8_calibration_per_head/test_kv_cache_calib_per_head_qwen2.5_14b.json +++ b/test/advanced_config/fp8_calibration_per_head/test_kv_cache_calib_per_head_qwen2.5_14b.json @@ -875,5 +875,502 @@ 0.0538504496216774, 0.0736607164144516 ] - ] + ], + "q_calibration": { + "version": "1.0", + "architectures": "Qwen2ForCausalLM", + "tensor": "q", + "calibration_stage": "prefill_and_decode_rope_pair_norm_bound", + "quant_type": "per_head", + "num_layers": 48, + "num_head": 8, + "qmin": -448.0, + "qmax": 448.0, + "scales_shape": [ + 48, + 8 + ], + "scales": [ + [ + 0.024359168484807014, + 0.014588438905775547, + 0.016693606972694397, + 0.019085492938756943, + 0.0331459566950798, + 0.031198419630527496, + 0.02046395093202591, + 0.02064039558172226 + ], + [ + 0.011792901903390884, + 0.02182481810450554, + 0.018223032355308533, + 0.0289231538772583, + 0.020388314500451088, + 0.013806928880512714, + 0.017620379105210304, + 0.025998810306191444 + ], + [ + 0.013750087469816208, + 0.016150565817952156, + 0.018145889043807983, + 0.00908266194164753, + 0.018548380583524704, + 0.03387211635708809, + 0.010761430487036705, + 0.01710752211511135 + ], + [ + 0.013337633572518826, + 0.01626240462064743, + 0.027586648240685463, + 0.012137503363192081, + 0.013598266057670116, + 0.01659339666366577, + 0.01721581071615219, + 0.01650177501142025 + ], + [ + 0.010620522312819958, + 0.011844871565699577, + 0.012841424904763699, + 0.017183411866426468, + 0.015775660052895546, + 0.012916413135826588, + 0.01092174556106329, + 0.015533560886979103 + ], + [ + 0.028964761644601822, + 0.02996930107474327, + 0.031513553112745285, + 0.02636159211397171, + 0.036943770945072174, + 0.037480611354112625, + 0.04077959060668945, + 0.03499114513397217 + ], + [ + 0.032286763191223145, + 0.03441798314452171, + 0.03761911392211914, + 0.028371211141347885, + 0.03811360523104668, + 0.03547193109989166, + 0.0291505865752697, + 0.04010683298110962 + ], + [ + 0.04146439582109451, + 0.032927148044109344, + 0.035208091139793396, + 0.036458130925893784, + 0.03027242422103882, + 0.03373913839459419, + 0.031825900077819824, + 0.030813030898571014 + ], + [ + 0.03204596787691116, + 0.04382704943418503, + 0.0275579746812582, + 0.038165293633937836, + 0.050316475331783295, + 0.03951522707939148, + 0.032703764736652374, + 0.027358578518033028 + ], + [ + 0.031024355441331863, + 0.04046265035867691, + 0.034903958439826965, + 0.0437946654856205, + 0.04088672623038292, + 0.04140536114573479, + 0.028904439881443977, + 0.038517311215400696 + ], + [ + 0.04395269602537155, + 0.03877522051334381, + 0.03959416225552559, + 0.03846414387226105, + 0.03584178164601326, + 0.038221195340156555, + 0.03670245409011841, + 0.03548015281558037 + ], + [ + 0.048821285367012024, + 0.03461268171668053, + 0.04205715283751488, + 0.04709198698401451, + 0.03540901839733124, + 0.031146535649895668, + 0.041428472846746445, + 0.03592243418097496 + ], + [ + 0.03548650071024895, + 0.029334072023630142, + 0.03227446600794792, + 0.03204076364636421, + 0.03773365914821625, + 0.04842876270413399, + 0.03427734971046448, + 0.039105020463466644 + ], + [ + 0.043562356382608414, + 0.04126812517642975, + 0.04234690964221954, + 0.03360659256577492, + 0.03938795626163483, + 0.04184635728597641, + 0.037781111896038055, + 0.030930759385228157 + ], + [ + 0.04419945925474167, + 0.030987508594989777, + 0.040564533323049545, + 0.037866491824388504, + 0.029172465205192566, + 0.03296279534697533, + 0.036964308470487595, + 0.028570324182510376 + ], + [ + 0.028622450307011604, + 0.03269753232598305, + 0.03126129135489464, + 0.03696795552968979, + 0.042203012853860855, + 0.038224563002586365, + 0.03194285184144974, + 0.04047691076993942 + ], + [ + 0.04120899736881256, + 0.038979705423116684, + 0.03419353440403938, + 0.03308095782995224, + 0.03498309478163719, + 0.03318129479885101, + 0.041171297430992126, + 0.025124337524175644 + ], + [ + 0.03802955895662308, + 0.03429017215967178, + 0.042303070425987244, + 0.04222817346453667, + 0.03234030306339264, + 0.047289956361055374, + 0.029485438019037247, + 0.03614581748843193 + ], + [ + 0.04146738350391388, + 0.0289231538772583, + 0.036455780267715454, + 0.03660660609602928, + 0.034280650317668915, + 0.03287815675139427, + 0.04227793961763382, + 0.03682270646095276 + ], + [ + 0.03308634087443352, + 0.040413301438093185, + 0.0339091531932354, + 0.051295943558216095, + 0.04447821527719498, + 0.03424642235040665, + 0.030958928167819977, + 0.03853381425142288 + ], + [ + 0.04004320129752159, + 0.03448706120252609, + 0.03118603676557541, + 0.03736995533108711, + 0.0357554666697979, + 0.04242594540119171, + 0.028710657730698586, + 0.035833101719617844 + ], + [ + 0.03122786432504654, + 0.026834741234779358, + 0.035893604159355164, + 0.028338786214590073, + 0.037639930844306946, + 0.0343923345208168, + 0.04374994710087776, + 0.034884508699178696 + ], + [ + 0.04359454661607742, + 0.029943428933620453, + 0.0341830849647522, + 0.03833219036459923, + 0.03369707241654396, + 0.03865843266248703, + 0.03293344005942345, + 0.0389111153781414 + ], + [ + 0.03847446292638779, + 0.03861941769719124, + 0.03062348999083042, + 0.033059436827898026, + 0.03301110863685608, + 0.03240063041448593, + 0.04058609902858734, + 0.02766369841992855 + ], + [ + 0.03376298025250435, + 0.03578426316380501, + 0.036965757608413696, + 0.0355362743139267, + 0.03761112689971924, + 0.03516975790262222, + 0.040464237332344055, + 0.03918806463479996 + ], + [ + 0.03457807004451752, + 0.03254459798336029, + 0.033271003514528275, + 0.04000820592045784, + 0.035757288336753845, + 0.0360182523727417, + 0.03488868102431297, + 0.03704201430082321 + ], + [ + 0.03541409224271774, + 0.03754438832402229, + 0.028201188892126083, + 0.029227353632450104, + 0.030108239501714706, + 0.03171919658780098, + 0.02996894344687462, + 0.02859300747513771 + ], + [ + 0.03389298915863037, + 0.038445938378572464, + 0.033777277916669846, + 0.0383286252617836, + 0.029926583170890808, + 0.031451381742954254, + 0.03698495030403137, + 0.03745793551206589 + ], + [ + 0.035882167518138885, + 0.034510985016822815, + 0.034302376210689545, + 0.03373913839459419, + 0.04205428436398506, + 0.03386495262384415, + 0.031501512974500656, + 0.030024299398064613 + ], + [ + 0.03443130850791931, + 0.027702871710062027, + 0.03374136611819267, + 0.034104444086551666, + 0.04278411343693733, + 0.041357651352882385, + 0.03488175943493843, + 0.03637121617794037 + ], + [ + 0.03374466672539711, + 0.0345914289355278, + 0.04181000217795372, + 0.03892875462770462, + 0.041846681386232376, + 0.033517103642225266, + 0.03965563699603081, + 0.03451665863394737 + ], + [ + 0.039268895983695984, + 0.03849337249994278, + 0.03940538689494133, + 0.03634502738714218, + 0.04064987599849701, + 0.041152529418468475, + 0.03435271233320236, + 0.030772149562835693 + ], + [ + 0.043006137013435364, + 0.034248851239681244, + 0.036484021693468094, + 0.033679064363241196, + 0.04005124047398567, + 0.03631957620382309, + 0.03748604655265808, + 0.03623544052243233 + ], + [ + 0.03875640034675598, + 0.03967435657978058, + 0.03847446292638779, + 0.032117847353219986, + 0.04112768918275833, + 0.03616621717810631, + 0.033236220479011536, + 0.037944477051496506 + ], + [ + 0.03583109378814697, + 0.03958548977971077, + 0.03162061423063278, + 0.04114099219441414, + 0.042592454701662064, + 0.033078882843256, + 0.03284982591867447, + 0.03279387578368187 + ], + [ + 0.030374785885214806, + 0.030821170657873154, + 0.035421665757894516, + 0.03689901903271675, + 0.034723613411188126, + 0.03361746668815613, + 0.03340458497405052, + 0.039419420063495636 + ], + [ + 0.034780725836753845, + 0.04221826046705246, + 0.0344911590218544, + 0.037793345749378204, + 0.03779105097055435, + 0.029417475685477257, + 0.04189971834421158, + 0.036193277686834335 + ], + [ + 0.032995421439409256, + 0.03746887296438217, + 0.03761148452758789, + 0.03769689053297043, + 0.038415949791669846, + 0.03746118023991585, + 0.035020794719457626, + 0.032971952110528946 + ], + [ + 0.040061336010694504, + 0.04114099219441414, + 0.03948236256837845, + 0.044569678604602814, + 0.038382332772016525, + 0.035295408219099045, + 0.030778804793953896, + 0.03820478543639183 + ], + [ + 0.03539469838142395, + 0.041420646011829376, + 0.03857550024986267, + 0.03856498748064041, + 0.037508368492126465, + 0.03838307410478592, + 0.03640927001833916, + 0.04143572598695755 + ], + [ + 0.040421463549137115, + 0.04514767602086067, + 0.03858954459428787, + 0.032494038343429565, + 0.04327918961644173, + 0.04091407358646393, + 0.04306666925549507, + 0.038458775728940964 + ], + [ + 0.03528358414769173, + 0.04331406578421593, + 0.037840064615011215, + 0.0432426743209362, + 0.038613639771938324, + 0.037182584404945374, + 0.029955970123410225, + 0.03499394655227661 + ], + [ + 0.03985505923628807, + 0.03641197830438614, + 0.04622425138950348, + 0.04795350134372711, + 0.037578094750642776, + 0.03200354799628258, + 0.03118741512298584, + 0.03464595600962639 + ], + [ + 0.03866660222411156, + 0.03503189980983734, + 0.03872193768620491, + 0.040000226348638535, + 0.035844773054122925, + 0.04460058733820915, + 0.03851550444960594, + 0.03525348752737045 + ], + [ + 0.039305880665779114, + 0.040484052151441574, + 0.0342857651412487, + 0.03715816140174866, + 0.03882361203432083, + 0.03441985324025154, + 0.037989936769008636, + 0.03467605635523796 + ], + [ + 0.028615113347768784, + 0.03879321739077568, + 0.03400520980358124, + 0.04875323921442032, + 0.03702318295836449, + 0.037619397044181824, + 0.039907313883304596, + 0.040792956948280334 + ], + [ + 0.033051058650016785, + 0.034135352820158005, + 0.0380941741168499, + 0.04974164813756943, + 0.033217087388038635, + 0.03705244138836861, + 0.032787661999464035, + 0.044617898762226105 + ], + [ + 0.02564295195043087, + 0.03051416203379631, + 0.029742099344730377, + 0.03528114780783653, + 0.02563738077878952, + 0.028006788343191147, + 0.030505722388625145, + 0.03242846950888634 + ] + ] + } } \ No newline at end of file diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 279840d5af..58d7080edd 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -568,18 +568,21 @@ def _ensure_req_kv_capacity(self, req_idx: torch.Tensor, target_kv_lens: torch.T req_id = int(req_id) max_target_len_by_req[req_id] = max(max_target_len_by_req.get(req_id, 0), int(target_len)) - mem_manager = self.model.req_manager.mem_manager - page_size = mem_manager.page_size + req_manager = self.model.req_manager + page_size = req_manager.mem_manager.page_size + allocations = [] for req_id, target_len in max_target_len_by_req.items(): new_hold_kv_len = align_up(target_len, page_size) old_hold_kv_len = self._hold_kv_len_by_req.get(req_id, 0) if new_hold_kv_len <= old_hold_kv_len: continue - alloc_token_num = new_hold_kv_len - old_hold_kv_len - mem_indexes = mem_manager.alloc(alloc_token_num) - self.model.req_manager.req_to_token_indexs[req_id, old_hold_kv_len:new_hold_kv_len].copy_(mem_indexes) - self._hold_kv_len_by_req[req_id] = new_hold_kv_len + req = SimpleNamespace(req_idx=req_id, hold_kv_len=old_hold_kv_len) + allocations.append((req, new_hold_kv_len - old_hold_kv_len)) + if allocations: + req_manager.alloc_req_token_indexes(allocations) + for req, _ in allocations: + self._hold_kv_len_by_req[req.req_idx] = req.hold_kv_len def _forward_prefill_input(self, model_input: ModelInput, allow_overlap: bool) -> ModelOutput: if allow_overlap and self.args.enable_prefill_microbatch_overlap and model_input.batch_size > 1: @@ -703,11 +706,6 @@ def concat_optional(value0: torch.Tensor | None, value1: torch.Tensor | None): output1.mtp_collector.confidence_logits, ), ), - prefill_mem_indexes_ready_event=( - output0.prefill_mem_indexes_ready_event - if output0.prefill_mem_indexes_ready_event is not None - else output1.prefill_mem_indexes_ready_event - ), prompt_logics=concat_optional(output0.prompt_logics, output1.prompt_logics), ) diff --git a/unit_tests/common/basemodel/attention/test_fp8_page_size.py b/unit_tests/common/basemodel/attention/test_fp8_page_size.py deleted file mode 100644 index 107728c26d..0000000000 --- a/unit_tests/common/basemodel/attention/test_fp8_page_size.py +++ /dev/null @@ -1,45 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import patch - -import pytest - -from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 -from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend -from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend - - -@pytest.mark.parametrize( - "backend_class", - [Fa3AttBackend, MlaFa3AttBackend, FlashInferAttBackend, MlaFlashInferAttBackend], -) -def test_default_infer_page_size_matches_model_page_size(backend_class): - backend = object.__new__(backend_class) - backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) - - backend._init_infer_page_size() - - assert backend.infer_page_size == 16 - - -@pytest.mark.parametrize( - ("backend_class", "backend_module"), - [ - (Fp8Fa3AttBackend, fa3_fp8), - (Fp8FlashInferAttBackend, flashinfer_fp8), - ], -) -def test_fp8_attention_backend_uses_single_token_infer_pages(backend_class, backend_module): - backend = object.__new__(backend_class) - backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) - - with patch.object(backend_module.logger, "warning") as warning: - backend._init_infer_page_size() - - assert backend.infer_page_size == 1 - warning.assert_called_once() - assert "infer_page_size=1" in warning.call_args.args[0] diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index cb0d9fef54..686f9af8f7 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -136,7 +136,7 @@ def pad_once(model_input, new_batch_size): graph_flags_at_att_init = [] executions = [] - def create_infer_state(model_input): + def create_infer_state(model_input, microbatch_index=0): infer_state = SimpleNamespace( b_req_idx=model_input.b_req_idx, b_seq_len=model_input.b_seq_len, diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py index 54f2ccacf1..847e2a6494 100644 --- a/unit_tests/common/basemodel/test_overlap_utils.py +++ b/unit_tests/common/basemodel/test_overlap_utils.py @@ -82,12 +82,6 @@ def record_to_cuda(self): monkeypatch.setattr(ModelInput, "to_cuda", record_to_cuda) - def fake_gather(**kwargs): - events.append("gather") - kwargs["input_ids"][-1] = 90 + events.count("gather") - - monkeypatch.setattr(basemodel, "gather_token_prefill_decode_mixed", fake_gather) - model = TpPartBaseModel.__new__(TpPartBaseModel) model.args = SimpleNamespace(enable_prefill_decode_mixed=True, page_size=1) model.req_manager = SimpleNamespace( @@ -107,14 +101,14 @@ def fake_overlap_forward(input0, input1): outputs = model.microbatch_overlap_prefill(model_input0, model_input1) - assert events == ["to_cuda", "gather", "to_cuda", "gather", "forward"] + assert events == ["to_cuda", "to_cuda", "forward"] assert captured_inputs == [model_input0, model_input1] - assert captured_inputs[0].input_ids.tolist() == [10, 91] - assert captured_inputs[1].input_ids.tolist() == [92] + assert captured_inputs[0].input_ids.tolist() == [10, 11] + assert captured_inputs[1].input_ids.tolist() == [12] assert len(outputs) == 2 -def test_base_model_decode_accepts_two_inputs_and_skips_empty_gather(monkeypatch): +def test_base_model_decode_defers_gather_until_padded_forward(monkeypatch): model_input0 = _make_decode_input() model_input0.input_ids = None model_input1 = _make_empty_decode_input() @@ -128,12 +122,6 @@ def record_to_cuda(self): monkeypatch.setattr(ModelInput, "to_cuda", record_to_cuda) - def fake_gather(**kwargs): - events.append("gather") - return torch.arange(40, 46, dtype=torch.int64, device="cuda") - - monkeypatch.setattr(basemodel, "gather_token", fake_gather) - model = TpPartBaseModel.__new__(TpPartBaseModel) model.args = SimpleNamespace(page_size=1) model.req_manager = SimpleNamespace( @@ -153,10 +141,9 @@ def fake_overlap_forward(input0, input1): model.microbatch_overlap_decode(model_input0, model_input1) - assert events == ["to_cuda", "gather", "to_cuda", "forward"] + assert events == ["to_cuda", "to_cuda", "forward"] assert captured_inputs == [model_input0, model_input1] - assert captured_inputs[0].input_ids.tolist() == [40, 41, 42, 43, 44, 45] - assert captured_inputs[1].input_ids.numel() == 0 + assert all(model_input.input_ids is None for model_input in captured_inputs) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py index 154af9acd3..b3068db5fb 100644 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py @@ -26,6 +26,7 @@ @pytest.fixture(autouse=True) def autotune_environment(monkeypatch): + torch.manual_seed(42) monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py index 2f8544cf75..fe1b968cf3 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py @@ -117,6 +117,8 @@ def reference(inputs): h += k[token, :, :, None] * delta[:, None, :] output[0, token] = (h * q[token, :, :, None]).sum(1) state[inputs["ssm_state_write_indices"][row, token - start]] = h.to(state.dtype) + # Each MTP step observes the same persistent-state rounding as decode. + h = h.to(state.dtype).float() return output, state diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py index 11035c83a3..14f9047289 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -1,12 +1,10 @@ import pytest import torch -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) +from unit_tests.common.basemodel.triton_kernel.linear_att.test_mtp_fused_recurrent_autotune import reference if not torch.cuda.is_available(): pytest.skip("CUDA required", allow_module_level=True) @@ -26,32 +24,13 @@ def _run_both( a_raw, b_raw, ): - """Run old (via autograd.Function) and new (direct kernel) side-by-side.""" - state_old = initial_state.clone() - state_new = initial_state.clone() - - o_old, fs_old = fused_recurrent_gated_delta_rule( - q=q, - k=k, - v=v, - initial_state=state_old, - inplace_final_state=True, - cu_seqlens=cu_seqlens, - ssm_state_indices=ssm_state_indices, - ssm_state_write_indices=ssm_state_write_indices, - num_accepted_tokens=num_accepted_tokens, - use_qk_l2norm_in_kernel=True, - A_log=A_log, - dt_bias=dt_bias, - a_raw=a_raw, - b_raw=b_raw, - ) - - o_new, fs_new = mtp_fused_recurrent_gated_delta_rule( + # The old FLA path retains FP32 state between tokens, unlike sequential decode + # and MTP since #1569. Compare with independent per-step rounding instead. + inputs = dict( q=q, k=k, v=v, - initial_state=state_new, + initial_state=initial_state.clone(), cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, ssm_state_write_indices=ssm_state_write_indices, @@ -61,8 +40,10 @@ def _run_both( a_raw=a_raw, b_raw=b_raw, ) - - return o_old, o_new, fs_old, fs_new + output_ref, state_ref = reference(inputs) + output, state = mtp_fused_recurrent_gated_delta_rule(**inputs) + torch.testing.assert_close(output.float(), output_ref, atol=2e-3, rtol=2e-2) + torch.testing.assert_close(state.float(), state_ref.float(), atol=2e-3, rtol=2e-2) @pytest.mark.parametrize("mtp_step", [1, 2, 3]) @@ -82,11 +63,11 @@ def test_mtp_verify_path(mtp_step): a_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) b_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) ssm_state = torch.randn(cache_slots, HV, K, V, device="cuda", dtype=torch.bfloat16) - ssm_idx = torch.randint(0, cache_slots, (batch, seqlen), device="cuda", dtype=torch.int32) + ssm_idx = torch.randperm(cache_slots, device="cuda", dtype=torch.int32)[:num_tokens].view(batch, seqlen) cu_seqlens = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * seqlen num_accepted = torch.full((batch,), seqlen, device="cuda", dtype=torch.int32) - o_old, o_new, fs_old, fs_new = _run_both( + _run_both( q, k, v, @@ -101,15 +82,6 @@ def test_mtp_verify_path(mtp_step): b_raw, ) - assert torch.equal( - o_old, o_new - ), f"output mismatch, max diff={torch.abs(o_old.float() - o_new.float()).max().item():.6f}" - if not torch.equal(fs_old, fs_new): - assert torch.allclose(fs_old.float(), fs_new.float(), rtol=1e-2, atol=5.0), ( - f"final_state mismatch at mtp_step={mtp_step}, " - f"max diff={torch.abs(fs_old.float() - fs_new.float()).max().item():.6f}" - ) - @pytest.mark.parametrize("seed", [1, 2, 3, 4]) def test_variable_length_cu_seqlens(seed): @@ -135,10 +107,10 @@ def test_variable_length_cu_seqlens(seed): b_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) ssm_state = torch.randn(cache_slots, HV, K, V, device="cuda", dtype=torch.bfloat16) # 2D indices: S+1 = max_len columns, unused columns for short seqs are irrelevant - ssm_idx = torch.randint(0, cache_slots, (batch, max_len), device="cuda", dtype=torch.int32) + ssm_idx = torch.randperm(cache_slots, device="cuda", dtype=torch.int32)[: batch * max_len].view(batch, max_len) num_accepted = lengths.to(device="cuda", dtype=torch.int32) - o_old, o_new, fs_old, fs_new = _run_both( + _run_both( q, k, v, @@ -153,15 +125,6 @@ def test_variable_length_cu_seqlens(seed): b_raw, ) - assert torch.equal( - o_old, o_new - ), f"seed={seed}: output mismatch, max diff={torch.abs(o_old.float() - o_new.float()).max().item():.6f}" - if not torch.equal(fs_old, fs_new): - assert torch.allclose(fs_old.float(), fs_new.float(), rtol=1e-2, atol=5.0), ( - f"seed={seed}: final_state mismatch, " - f"max diff={torch.abs(fs_old.float() - fs_new.float()).max().item():.6f}" - ) - if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py b/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py index 55e7d0724a..c1ef686f8e 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py +++ b/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py @@ -4,19 +4,7 @@ if not torch.cuda.is_available(): pytest.skip("requires CUDA", allow_module_level=True) -from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params, page_table_copy - - -def test_page_table_copy_uses_page_bases(): - req_to_token_indexs = torch.tensor([list(range(40, 52)), list(range(80, 92))], dtype=torch.int32, device="cuda") - page_table = torch.empty((2, 3), dtype=torch.int32, device="cuda") - page_table_copy( - page_table=page_table, - req_to_token_indexs=req_to_token_indexs, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - page_size=4, - ) - assert page_table.cpu().tolist() == [[20, 21, 22], [10, 11, 12]] +from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params def _reference_dynamic_spec_fa3_decode_params(b_req_idx, b_seq_len, b_mark_mtp_shared_group, hold_req_id): diff --git a/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py b/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py index 0e66266dea..15a4447466 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py +++ b/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py @@ -48,34 +48,3 @@ def repack_kv_ref(req_to_token_indexs, b_req_idx, b_seq_len, b_start_loc, output out_page_indices=output, ) assert torch.allclose(output.float(), ref.float()) - - -@pytest.mark.parametrize("page_size, page_count", [(1, 3), (3, 3), (4, 3), (16, 3), (4, 65), (16, 65)]) -def test_repack_kv_index_with_pages(page_size, page_count): - req_to_token_indexs = torch.tensor( - [ - list(range(10 * page_size, (10 + page_count) * page_size)), - list(range(20 * page_size, (20 + page_count) * page_size)), - list(range(30 * page_size, (30 + page_count) * page_size)), - ], - dtype=torch.int32, - device="cuda", - ) - req_indexes = torch.tensor([2, 0, 1], dtype=torch.int32, device="cuda") - # 使用不满的末页，并覆盖超过单个 block（64 页）的请求。 - max_seq_len = page_count * page_size - (page_size > 1) - seq_lens = torch.tensor([2 * page_size, 1, max_seq_len], dtype=torch.int32, device="cuda") - starts = torch.tensor([0, 2, 3], dtype=torch.int32, device="cuda") - output = torch.empty((3 + page_count,), dtype=torch.int32, device="cuda") - - repack_kv_index( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=req_indexes, - b_token_len=seq_lens, - b_page_start_loc=starts, - max_token_len=max_seq_len, - out_page_indices=output, - page_size=page_size, - ) - - assert output.cpu().tolist() == [30, 31, 10] + list(range(20, 20 + page_count)) diff --git a/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py b/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py index ec67af6077..80cc75ee33 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py +++ b/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py @@ -81,39 +81,6 @@ def test_scatter_routing_topk_respects_layer_index(): assert torch.equal(routing_buffer, expected) -def test_scatter_routing_topk_ignores_hold_indexes(): - _skip_without_cuda() - - kv_cache_size = 8 - num_moe_layers = 2 - topk = 2 - topk_ids = torch.arange(8, dtype=torch.int64, device="cuda").view(4, topk) - mem_indexes = torch.tensor([3, kv_cache_size, kv_cache_size + 1, -1], dtype=torch.int32, device="cuda") - routing_buffer = torch.zeros( - (kv_cache_size, num_moe_layers, topk), - dtype=torch.int16, - device="cpu", - pin_memory=True, - ) - routing_buffer_ptr = torch.tensor([routing_buffer.data_ptr()], dtype=torch.uint64, device="cuda") - - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=1, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=2, - ) - torch.cuda.synchronize() - - expected = torch.zeros_like(routing_buffer) - expected[3, 1, :] = topk_ids[0].cpu().to(torch.int16) - assert torch.equal(routing_buffer, expected) - - def test_scatter_routing_topk_is_cuda_graph_capturable(): _skip_without_cuda() diff --git a/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py b/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py deleted file mode 100644 index 0dab062934..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py +++ /dev/null @@ -1,75 +0,0 @@ -import pytest -import torch -from types import SimpleNamespace - -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( - select_kv_index_from_req, - select_kv_index_from_req_prefill, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def test_select_kv_index_from_req_for_decode(): - req_to_token_indexs = torch.tensor( - [ - [10, 11, 12, 13, 14, 15, 16, 17], - [20, 21, 22, 23, 24, 25, 26, 27], - ], - dtype=torch.int32, - device="cuda", - ) - - mem_indexes = select_kv_index_from_req( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 6], dtype=torch.int32, device="cuda"), - ) - - assert mem_indexes.cpu().tolist() == [22, 15] - - -def test_select_kv_index_from_req_for_prefill(): - req_to_token_indexs = torch.tensor( - [ - [10, 11, 12, 13, 14, 15, 16, 17], - [20, 21, 22, 23, 24, 25, 26, 27], - ], - dtype=torch.int32, - device="cuda", - ) - - mem_indexes = select_kv_index_from_req_prefill( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([4, 7], dtype=torch.int32, device="cuda"), - b_ready_cache_len=torch.tensor([2, 3], dtype=torch.int32, device="cuda"), - b_start_loc=torch.tensor([0, 2], dtype=torch.int32, device="cuda"), - max_q_seq_len=4, - token_num=6, - ) - - assert mem_indexes.cpu().tolist() == [22, 23, 13, 14, 15, 16] - - -def test_model_selects_reserved_indexes_when_execution_starts(): - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(page_size=4) - model.req_manager = SimpleNamespace( - req_to_token_indexs=torch.tensor( - [[10, 11, 12, 13], [20, 21, 22, 23]], - dtype=torch.int32, - device="cuda", - ) - ) - model_input = SimpleNamespace( - is_prefill=False, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 4], dtype=torch.int32, device="cuda"), - ) - - mem_indexes = model._select_mem_indexes(model_input) - - assert mem_indexes.cpu().tolist() == [22, 13] diff --git a/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py b/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py index 3e14d5baaa..772644c28f 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py +++ b/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py @@ -17,14 +17,19 @@ ], ) def test_sp_pad_copy(token_num, hidden_dim, sp_world_size): - in_tensor = torch.randn((token_num, hidden_dim), dtype=torch.float16, device="cuda") + if token_num % sp_world_size: + # TPSP pads the model input before splitting it between ranks. + with pytest.raises(AssertionError, match="in_token_num % sp_world_size != 0"): + sp_pad_copy(in_tensor=in_tensor, sp_rank_id=0, sp_world_size=sp_world_size) + padding = in_tensor[: sp_world_size - token_num % sp_world_size] + in_tensor = torch.cat([in_tensor, padding]) out_tensors = [ sp_pad_copy(in_tensor=in_tensor, sp_rank_id=rank_id, sp_world_size=sp_world_size) for rank_id in range(sp_world_size) ] out_tensor = torch.cat(out_tensors, dim=0) - assert torch.equal(in_tensor, out_tensor[0:token_num, :]) + assert torch.equal(in_tensor, out_tensor) if __name__ == "__main__": diff --git a/unit_tests/common/test_paged_mem_manager.py b/unit_tests/common/test_paged_mem_manager.py deleted file mode 100644 index 68bc514991..0000000000 --- a/unit_tests/common/test_paged_mem_manager.py +++ /dev/null @@ -1,32 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager -from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager - - -def test_kv_cache_allocation_requires_complete_pages(): - manager = MemoryManager.__new__(MemoryManager) - manager.page_size = 4 - manager.allocator = SimpleNamespace(alloc=lambda size: torch.arange(size, dtype=torch.int32)) - - with pytest.raises(AssertionError, match="must be a multiple of page_size 4"): - manager.alloc(5) - - assert manager.alloc(8).tolist() == list(range(8)) - - -@pytest.mark.parametrize( - "manager_class", - [MemoryManager, Deepseek2MemoryManager, PPLINT4KVMemoryManager, PPLINT8KVMemoryManager], -) -def test_kv_cache_buffer_size_requires_complete_pages(manager_class): - manager = manager_class.__new__(manager_class) - manager.page_size = 4 - - with pytest.raises(AssertionError, match="KV cache size 5 must be a multiple of page_size 4"): - manager._init_buffers(size=5, dtype=torch.float16, head_num=1, head_dim=8, layer_num=1) diff --git a/unit_tests/common/test_pinned_memory_non_blocking_copy.py b/unit_tests/common/test_pinned_memory_non_blocking_copy.py deleted file mode 100644 index 594157560b..0000000000 --- a/unit_tests/common/test_pinned_memory_non_blocking_copy.py +++ /dev/null @@ -1,39 +0,0 @@ -import pytest -import torch -from torch.profiler import ProfilerActivity, profile - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def _profile_request_table_copy(non_blocking: bool): - token_num = 1024 * 1024 - mem_indexes = torch.arange(token_num, dtype=torch.int32, pin_memory=True) - req_to_token_indexes = torch.empty((1, token_num), dtype=torch.int32, device="cuda") - - # 先完成分配等准备操作，避免其 CUDA runtime 调用混入待验证区间。 - torch.cuda.synchronize() - with profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - acc_events=True, - ) as prof: - req_to_token_indexes[0].copy_(mem_indexes, non_blocking=non_blocking) - if non_blocking: - # 异步路径需要在读取结果前显式等待；这里会记录为 cudaDeviceSynchronize， - # 与 copy_ 内部可能产生的 cudaStreamSynchronize 可以明确区分。 - torch.cuda.synchronize() - - runtime_calls = {event.key for event in prof.key_averages()} - assert torch.equal(req_to_token_indexes[0].cpu(), mem_indexes) - return runtime_calls - - -def test_pinned_cpu_to_gpu_non_blocking_copy_avoids_internal_stream_sync(): - blocking_calls = _profile_request_table_copy(non_blocking=False) - non_blocking_calls = _profile_request_table_copy(non_blocking=True) - - # 两条路径底层都提交异步 H2D memcpy；同步路径随后在 copy_ 内部等待当前 stream。 - assert "cudaMemcpyAsync" in blocking_calls - assert "cudaMemcpyAsync" in non_blocking_calls - assert "cudaStreamSynchronize" in blocking_calls - assert "cudaStreamSynchronize" not in non_blocking_calls diff --git a/unit_tests/common/test_quantized_paged_mem_manager.py b/unit_tests/common/test_quantized_paged_mem_manager.py deleted file mode 100644 index 84fe266613..0000000000 --- a/unit_tests/common/test_quantized_paged_mem_manager.py +++ /dev/null @@ -1,33 +0,0 @@ -from types import SimpleNamespace - -import torch - -from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager - - -def test_quantized_kv_buffers_include_the_full_hold_page(monkeypatch): - allocations = [] - - def fake_empty(shape, **kwargs): - allocations.append(shape) - return SimpleNamespace(shape=shape) - - monkeypatch.setattr(torch, "empty", fake_empty) - - int8_manager = PPLINT8KVMemoryManager.__new__(PPLINT8KVMemoryManager) - int8_manager.page_size = 4 - int8_manager.group_quant_size = 8 - int8_manager._init_buffers(size=32, dtype=torch.float16, head_num=2, head_dim=16, layer_num=3) - - int4_manager = PPLINT4KVMemoryManager.__new__(PPLINT4KVMemoryManager) - int4_manager.page_size = 4 - int4_manager.group_quant_size = 8 - int4_manager._init_buffers(size=32, dtype=torch.float16, head_num=2, head_dim=16, layer_num=3) - - assert allocations == [ - (3, 36, 4, 16), - (3, 36, 4, 2), - (3, 36, 4, 8), - (3, 36, 4, 2), - ] diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py deleted file mode 100644 index 224c308f28..0000000000 --- a/unit_tests/common/test_req_manager_page.py +++ /dev/null @@ -1,549 +0,0 @@ -from types import MethodType, SimpleNamespace - -import pytest -import torch - -from lightllm.common.req_manager import ReqManager -from lightllm.server.router.model_infer.mode_backend import generic_pre_process -from lightllm.server.router.model_infer import infer_batch -from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext -from lightllm.server.router.model_infer.mode_backend import base_backend - - -class _FakeMemManager: - page_size = 4 - - def __init__(self): - self.next_index = 0 - self.alloc_sizes = [] - - def alloc(self, size): - self.alloc_sizes.append(size) - result = torch.arange(self.next_index, self.next_index + size, dtype=torch.int32) - self.next_index += size - return result - - -def test_bind_mem_manager_initializes_hold_request_with_reserved_page(): - req_manager = ReqManager.__new__(ReqManager) - req_manager.HOLD_REQUEST_ID = 2 - req_manager.req_list = SimpleNamespace(is_all_free=lambda: True) - req_manager.req_to_token_indexs = torch.full((3, 8), -1, dtype=torch.int32) - mem_manager = SimpleNamespace(page_size=4, HOLD_TOKEN_MEMINDEXES=(100, 101, 102, 103)) - - req_manager.bind_mem_manager(mem_manager) - - assert req_manager.mem_manager is mem_manager - assert req_manager.req_to_token_indexs[2].tolist() == [100, 101, 102, 103] * 2 - assert req_manager.req_to_token_indexs[:2].eq(-1).all() - - -def test_hold_request_indexes_require_all_requests_released(): - req_manager = ReqManager.__new__(ReqManager) - req_manager.req_list = SimpleNamespace(is_all_free=lambda: False) - - with pytest.raises(AssertionError, match="all requests are released"): - req_manager.init_hold_request_indexs() - - -def _make_context(monkeypatch): - mem_manager = _FakeMemManager() - req_manager = SimpleNamespace( - mem_manager=mem_manager, - req_to_token_indexs=torch.full((2, 16), -1, dtype=torch.int32), - ) - context = SimpleNamespace( - args=SimpleNamespace(page_size=4), - req_manager=req_manager, - radix_cache=None, - ) - monkeypatch.setattr(generic_pre_process, "g_infer_context", context) - monkeypatch.setattr(base_backend, "g_infer_context", context) - backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) - backend.args = context.args - return context, backend - - -def _make_req(req_idx): - req = SimpleNamespace(req_idx=req_idx, cur_kv_len=0, hold_kv_len=0) - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - return req - - -def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypatch): - context, backend = _make_context(monkeypatch) - req = _make_req(0) - - mem_indexes = backend._alloc_req_kv_mem(req, alloc_token_num=4) - assert mem_indexes.tolist() == [0, 1, 2, 3] - assert req.hold_kv_len == 4 - assert context.req_manager.mem_manager.alloc_sizes == [4] - assert context.req_manager.req_to_token_indexs[0, :4].tolist() == [0, 1, 2, 3] - - req.cur_kv_len = 3 - assert backend._alloc_req_kv_mem(req, alloc_token_num=0) is None - assert context.req_manager.mem_manager.alloc_sizes == [4] - - req.cur_kv_len = 4 - backend._alloc_req_kv_mem(req, alloc_token_num=4) - assert req.hold_kv_len == 8 - assert context.req_manager.mem_manager.alloc_sizes == [4, 4] - assert context.req_manager.req_to_token_indexs[0, :8].tolist() == list(range(8)) - - -def test_reservation_fills_each_request_table_row(monkeypatch): - _, backend = _make_context(monkeypatch) - req0 = _make_req(0) - req1 = _make_req(1) - - backend._alloc_req_kv_mem(req0, alloc_token_num=4) - backend._alloc_req_kv_mem(req1, alloc_token_num=4) - - assert generic_pre_process.g_infer_context.req_manager.req_to_token_indexs[0, :4].tolist() == [0, 1, 2, 3] - assert generic_pre_process.g_infer_context.req_manager.req_to_token_indexs[1, :4].tolist() == [4, 5, 6, 7] - assert req0.hold_kv_len == req1.hold_kv_len == 4 - - -def test_alloc_req_kv_mem_forwards_non_blocking_copy_option(monkeypatch): - class _CopyTarget: - def __init__(self): - self.non_blocking_values = [] - - def copy_(self, source, non_blocking=False): - self.non_blocking_values.append(non_blocking) - - class _ReqToTokenIndexes: - def __init__(self, target): - self.target = target - - def __getitem__(self, key): - return self.target - - target = _CopyTarget() - context = SimpleNamespace( - req_manager=SimpleNamespace( - mem_manager=_FakeMemManager(), - req_to_token_indexs=_ReqToTokenIndexes(target), - ), - radix_cache=None, - ) - monkeypatch.setattr(base_backend, "g_infer_context", context) - backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) - backend.args = SimpleNamespace(page_size=4) - req = _make_req(0) - - backend._alloc_req_kv_mem(req, alloc_token_num=4) - backend._alloc_req_kv_mem(req, alloc_token_num=4, no_blcoking_copy=True) - - assert target.non_blocking_values == [False, True] - - -def test_need_token_num_distinguishes_compute_and_page_allocation(): - req = SimpleNamespace( - args=SimpleNamespace(page_size=4), - cur_kv_len=3, - hold_kv_len=4, - mtp_step=0, - get_chuncked_input_token_len=lambda: 6, - get_cur_total_len=lambda: 10, - ) - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - - assert InferReq.prefill_need_token_num(req, is_chuncked_prefill=True) == (3, 4) - assert InferReq.prefill_need_token_num(req, is_chuncked_prefill=False) == (7, 8) - assert InferReq.decode_need_token_num(req) == (1, 0) - - req.cur_kv_len = 4 - assert InferReq.decode_need_token_num(req) == (1, 4) - - req.hold_kv_len = 8 - assert InferReq.prefill_need_token_num(req, is_chuncked_prefill=True) == (2, 0) - assert InferReq.prefill_need_token_num(req, is_chuncked_prefill=False) == (6, 4) - - req.mtp_step = 2 - assert InferReq.decode_need_token_num(req) == (6, 4) - req.hold_kv_len = 12 - assert InferReq.decode_need_token_num(req) == (6, 0) - - -@pytest.mark.parametrize( - "hold_kv_len, can_alloc_token_num, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause", - [ - (4, 8, 3, [4], False), - (8, 0, 3, [], False), - (4, 4, 6, [4], True), - ], -) -def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( - monkeypatch, hold_kv_len, can_alloc_token_num, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause -): - context, backend = _make_context(monkeypatch) - backend.args.enable_cpu_cache = False - backend.args.enable_prefill_decode_mixed = False - backend.support_overlap = False - backend.disable_chunked_prefill = False - backend.batch_max_tokens = batch_max_tokens - backend.is_master_in_dp = True - backend._timer_merge_radix_tree = lambda: None - backend.prefill_queue_strategy = SimpleNamespace(reorder=lambda reqs, **kwargs: reqs) - context.get_can_alloc_token_num = lambda: can_alloc_token_num - context.cache_placement_controller = SimpleNamespace(set_req_cache_way=lambda reqs: None) - context.filter_reqs = lambda finished_reqs: None - context.pause_reqs = lambda reqs, is_master_in_dp: None - - reqs = [_make_req(0), _make_req(1)] - for req in reqs: - req.args = context.args - req.cur_kv_len = 3 - req.hold_kv_len = hold_kv_len - req.filter_mark = False - req.wait_pause = False - req.paused = False - req.infer_aborted = False - req.finish_status = infer_batch.FinishStatus() - req.is_slave_req = lambda: False - req.get_cur_total_len = lambda: 6 - req.get_chuncked_input_token_len = lambda: 6 - req.prefill_need_token_num = MethodType(InferReq.prefill_need_token_num, req) - start_index = req.req_idx * hold_kv_len - context.req_manager.req_to_token_indexs[req.req_idx, :hold_kv_len] = torch.arange( - start_index, start_index + hold_kv_len, dtype=torch.int32 - ) - context.req_manager.mem_manager.next_index = 2 * hold_kv_len - backend._filter_not_ready_reqs = lambda req_ids: reqs - copy_modes = [] - original_alloc_req_kv_mem = backend._alloc_req_kv_mem - - def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): - copy_modes.append(no_blcoking_copy) - return original_alloc_req_kv_mem( - req_obj, - alloc_token_num, - no_blcoking_copy=no_blcoking_copy, - ) - - backend._alloc_req_kv_mem = _record_alloc - - prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0, 1]) - - assert prefill_reqs == [reqs[0]] - assert decode_reqs == [] - assert prefill_tokens == 3 - assert context.req_manager.mem_manager.alloc_sizes == expected_alloc_sizes - assert reqs[0].hold_kv_len == 8 - assert not reqs[0].wait_pause - assert reqs[1].hold_kv_len == hold_kv_len - assert reqs[1].wait_pause is second_req_wait_pause - assert copy_modes == [True] - - -def test_decode_scheduler_uses_non_blocking_request_table_copy(monkeypatch): - context, backend = _make_context(monkeypatch) - backend.args.enable_cpu_cache = False - backend.args.enable_prefill_decode_mixed = False - backend.args.run_mode = "normal" - backend.support_overlap = False - backend.batch_max_tokens = 8 - backend.is_master_in_dp = True - backend._timer_merge_radix_tree = lambda: None - backend.prefill_queue_strategy = SimpleNamespace(reorder=lambda reqs, **kwargs: reqs) - context.get_can_alloc_token_num = lambda: 12 - context.cache_placement_controller = SimpleNamespace(set_req_cache_way=lambda reqs: None) - context.filter_reqs = lambda finished_reqs: None - context.pause_reqs = lambda reqs, is_master_in_dp: None - - req = _make_req(0) - req.args = context.args - req.cur_kv_len = 4 - req.hold_kv_len = 4 - req.mtp_step = 0 - req.filter_mark = False - req.wait_pause = False - req.paused = False - req.infer_aborted = False - req.finish_status = infer_batch.FinishStatus() - req.get_cur_total_len = lambda: 5 - req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) - context.req_manager.req_to_token_indexs[0, :4] = torch.arange(4, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 4 - backend._filter_not_ready_reqs = lambda req_ids: [req] - - copy_modes = [] - original_alloc_req_kv_mem = backend._alloc_req_kv_mem - - def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): - copy_modes.append(no_blcoking_copy) - return original_alloc_req_kv_mem( - req_obj, - alloc_token_num, - no_blcoking_copy=no_blcoking_copy, - ) - - backend._alloc_req_kv_mem = _record_alloc - - prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0]) - - assert prefill_reqs == [] - assert decode_reqs == [req] - assert prefill_tokens == 0 - assert req.hold_kv_len == 16 - assert context.req_manager.mem_manager.alloc_sizes == [12] - assert copy_modes == [True] - - -@pytest.mark.parametrize( - "page_size, expected_alloc_token_num", - [(1, 9), (3, 9), (4, 12), (7, 14), (8, 8)], -) -def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_steps( - monkeypatch, page_size, expected_alloc_token_num -): - context, backend = _make_context(monkeypatch) - context.args.page_size = page_size - context.req_manager.req_to_token_indexs = torch.full((2, 64), -1, dtype=torch.int32) - backend.args.enable_cpu_cache = False - backend.args.enable_prefill_decode_mixed = False - backend.args.run_mode = "normal" - backend.support_overlap = False - backend.batch_max_tokens = 8 - backend.is_master_in_dp = True - backend._timer_merge_radix_tree = lambda: None - backend.prefill_queue_strategy = SimpleNamespace(reorder=lambda reqs, **kwargs: reqs) - context.get_can_alloc_token_num = lambda: 32 - context.cache_placement_controller = SimpleNamespace(set_req_cache_way=lambda reqs: None) - context.filter_reqs = lambda finished_reqs: None - context.pause_reqs = lambda reqs, is_master_in_dp: None - - req = _make_req(0) - req.args = context.args - req.cur_kv_len = page_size - req.hold_kv_len = page_size - req.mtp_step = 0 - req.filter_mark = False - req.wait_pause = False - req.paused = False - req.infer_aborted = False - req.finish_status = infer_batch.FinishStatus() - req.get_cur_total_len = lambda: req.cur_kv_len + 1 - req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) - context.req_manager.req_to_token_indexs[0, :page_size] = torch.arange(page_size, dtype=torch.int32) - context.req_manager.mem_manager.next_index = page_size - backend._filter_not_ready_reqs = lambda req_ids: [req] - - for _ in range(expected_alloc_token_num): - _, decode_reqs, _ = backend._get_classed_reqs(req_ids=[0]) - assert decode_reqs == [req] - req.cur_kv_len += 1 - - assert req.hold_kv_len == page_size + expected_alloc_token_num - assert context.req_manager.mem_manager.alloc_sizes == [expected_alloc_token_num] - - -def test_decode_reserves_mtp_headroom(monkeypatch): - context, backend = _make_context(monkeypatch) - req = _make_req(0) - req.cur_kv_len = 3 - req.hold_kv_len = 4 - req.mtp_step = 2 - req.multimodal_params = {"images": [], "audios": []} - req.shared_kv_node = None - req.get_cur_total_len = lambda: 4 - req.get_radix_cache_shared_len = lambda: 0 - req.args = context.args - - context.req_manager.req_to_token_indexs[0, :4] = torch.arange(4, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 4 - token_num, alloc_token_num = InferReq.decode_need_token_num(req) - assert token_num == 6 - assert alloc_token_num == 8 - backend._alloc_req_kv_mem(req, alloc_token_num) - - model_input, run_reqs = generic_pre_process.prepare_decode_inputs([req]) - - assert run_reqs == [req, req, req] - assert model_input.b_seq_len.tolist() == [4, 5, 6] - assert req.hold_kv_len == 12 - assert context.req_manager.mem_manager.alloc_sizes == [8] - assert context.req_manager.req_to_token_indexs[0, :12].tolist() == list(range(12)) - - -def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): - context, backend = _make_context(monkeypatch) - context.args.page_size = 1 - req = _make_req(0) - req.cur_kv_len = 3 - req.hold_kv_len = 3 - req.mtp_step = 2 - req.multimodal_params = {"images": [], "audios": []} - req.shared_kv_node = None - req.get_cur_total_len = lambda: 4 - req.get_radix_cache_shared_len = lambda: 0 - req.args = context.args - context.req_manager.req_to_token_indexs[0, :3] = torch.arange(3, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 3 - token_num, alloc_token_num = InferReq.decode_need_token_num(req) - assert token_num == alloc_token_num == 6 - backend._alloc_req_kv_mem(req, alloc_token_num) - - model_input, _ = generic_pre_process.prepare_decode_inputs([req]) - - assert req.hold_kv_len == 9 - assert context.req_manager.mem_manager.alloc_sizes == [6] - assert context.req_manager.req_to_token_indexs[0, :9].tolist() == list(range(9)) - - -def test_page_size_one_frees_all_preallocated_indexes(): - infer_context = InferenceContext.__new__(InferenceContext) - infer_context.args = SimpleNamespace(page_size=1) - infer_context.radix_cache = None - infer_context.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(16, dtype=torch.int32)[None, :]) - req = SimpleNamespace( - req_idx=0, - cur_kv_len=3, - hold_kv_len=12, - shm_req=SimpleNamespace(shm_cur_kv_len=3), - ) - free_token_indexes = [] - - infer_context.free_a_req_mem(free_token_indexes, req) - - assert free_token_indexes[0].tolist() == list(range(12)) - assert req.cur_kv_len == req.hold_kv_len == 0 - - -def test_linear_attention_frees_reserved_page_tail(monkeypatch): - infer_context = InferenceContext.__new__(InferenceContext) - infer_context.args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) - infer_context.radix_cache = SimpleNamespace() - infer_context.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(16, dtype=torch.int32)[None, :]) - req = SimpleNamespace( - req_idx=0, - cur_kv_len=5, - hold_kv_len=8, - hybrid_cache_len=0, - tail_small_page_buffer_id=None, - hybrid_len_to_big_page_id={}, - shared_kv_node=None, - ) - free_token_indexes = [] - monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", True) - monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: infer_context.args) - - infer_context._hybrid_att_free_req(free_token_indexes, req) - - assert free_token_indexes[0].tolist() == list(range(8)) - - -def test_linear_attention_frees_preallocated_pages_before_first_forward(monkeypatch): - infer_context = InferenceContext.__new__(InferenceContext) - infer_context.args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) - infer_context.radix_cache = SimpleNamespace() - infer_context.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, dtype=torch.int32)[None, :]) - req = SimpleNamespace( - req_idx=0, - cur_kv_len=0, - hold_kv_len=8, - hybrid_cache_len=0, - tail_small_page_buffer_id=None, - hybrid_len_to_big_page_id={}, - shared_kv_node=None, - ) - free_token_indexes = [] - monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", True) - monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: infer_context.args) - - infer_context._hybrid_att_free_req(free_token_indexes, req) - - assert free_token_indexes[0].tolist() == list(range(8)) - - -def _make_paused_req(req_idx, target_kv_len, page_size, match_kv_len=0, chunk_kv_len=None): - if chunk_kv_len is None: - chunk_kv_len = target_kv_len - req = SimpleNamespace( - args=SimpleNamespace(page_size=page_size), - req_id=req_idx, - req_idx=req_idx, - cur_kv_len=0, - hold_kv_len=0, - paused=True, - match_call_count=0, - shared_kv_node=None, - shm_req=SimpleNamespace(is_paused=True, shm_cur_kv_len=0), - get_cur_total_len=lambda: target_kv_len, - get_chuncked_input_token_len=lambda: chunk_kv_len, - ) - - def match_radix_cache(): - req.match_call_count += 1 - req.cur_kv_len = match_kv_len - req.hold_kv_len = match_kv_len - req.shm_req.shm_cur_kv_len = match_kv_len - - req._match_radix_cache = match_radix_cache - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - req.prefill_need_token_num = MethodType(InferReq.prefill_need_token_num, req) - return req - - -def test_recover_paused_reqs_uses_page_allocation_need(monkeypatch): - infer_context = InferenceContext.__new__(InferenceContext) - infer_context.args = SimpleNamespace(page_size=4) - infer_context.backend = SimpleNamespace(disable_chunked_prefill=False) - infer_context.radix_cache = None - freed_indexes = [] - infer_context.req_manager = SimpleNamespace( - req_to_token_indexs=torch.arange(16, dtype=torch.int32).reshape(2, 8), - free_token=lambda indexes: freed_indexes.extend(indexes.tolist()), - ) - infer_context.get_can_alloc_token_num = lambda: 5 - large_req = _make_paused_req(req_idx=0, target_kv_len=5, page_size=4) - small_req = _make_paused_req(req_idx=1, target_kv_len=4, page_size=4) - monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", False) - monkeypatch.setattr(infer_batch, "custom_cat", lambda tensors: torch.cat(tensors)) - - infer_context.recover_paused_reqs([large_req, small_req], is_master_in_dp=True) - - assert large_req.paused is True - assert large_req.shm_req.is_paused is True - assert large_req.match_call_count == 0 - assert small_req.paused is True - assert small_req.shm_req.is_paused is True - assert small_req.match_call_count == 0 - assert freed_indexes == [] - - -def test_recover_paused_reqs_requires_capacity_for_the_full_sequence(monkeypatch): - infer_context = InferenceContext.__new__(InferenceContext) - infer_context.args = SimpleNamespace(page_size=4) - infer_context.radix_cache = None - infer_context.req_manager = SimpleNamespace( - req_to_token_indexs=torch.arange(8, dtype=torch.int32)[None, :], - free_token=lambda indexes: None, - ) - infer_context.get_can_alloc_token_num = lambda: 4 - req = _make_paused_req(req_idx=0, target_kv_len=8, page_size=4, chunk_kv_len=4) - monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", False) - monkeypatch.setattr(infer_batch, "custom_cat", lambda tensors: torch.cat(tensors)) - - infer_context.recover_paused_reqs([req], is_master_in_dp=True) - - assert req.paused is True - assert req.shm_req.is_paused is True - assert req.match_call_count == 0 - - -def test_recover_paused_reqs_accounts_for_rematched_prefix(monkeypatch): - infer_context = InferenceContext.__new__(InferenceContext) - infer_context.args = SimpleNamespace(page_size=4) - infer_context.backend = SimpleNamespace(disable_chunked_prefill=False) - infer_context.get_can_alloc_token_num = lambda: 12 - req = _make_paused_req(req_idx=0, target_kv_len=9, page_size=4, match_kv_len=8) - monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", False) - - infer_context.recover_paused_reqs([req], is_master_in_dp=True) - - assert req.cur_kv_len == req.hold_kv_len == 8 - assert req.paused is False - assert req.shm_req.is_paused is False - assert req.match_call_count == 1 diff --git a/unit_tests/models/test_qwen3_dspark_model_output.py b/unit_tests/models/test_qwen3_dspark_model_output.py index 8e2085fa00..27484e1b85 100644 --- a/unit_tests/models/test_qwen3_dspark_model_output.py +++ b/unit_tests/models/test_qwen3_dspark_model_output.py @@ -61,7 +61,7 @@ def test_dspark_without_block_size_uses_mtp_step(model_class, mtp_step): "model_class", [Qwen3DFlashModel, Qwen3_5DFlashModel, Qwen3DSparkModel, Qwen3_5DSparkModel], ) -def test_parallel_block_decode_commits_target_hiddens_directly(model_class): +def test_parallel_block_decode_commits_target_hiddens_directly(model_class, monkeypatch): model = model_class.__new__(model_class) model._cos_cached = torch.arange(24).view(6, 4) model._sin_cached = model._cos_cached + 100 @@ -70,7 +70,14 @@ def test_parallel_block_decode_commits_target_hiddens_directly(model_class): target_hiddens = torch.arange(6, dtype=torch.float32).view(2, 3) mem_indexes = torch.tensor([7, 11]) - model._select_mem_indexes = lambda _: mem_indexes + model.req_manager = SimpleNamespace(req_to_token_indexs=object()) + + def build_inputs(inputs, table): + assert inputs is model_input + assert table is model.req_manager.req_to_token_indexs + return mem_indexes, None + + monkeypatch.setattr("lightllm.models.qwen3_dflash.model.build_kv_indexes_and_input_ids", build_inputs) observed_states = [] class PreInfer: diff --git a/unit_tests/server/core/objs/test_req.py b/unit_tests/server/core/objs/test_req.py index 5966497ef1..042ef16bae 100644 --- a/unit_tests/server/core/objs/test_req.py +++ b/unit_tests/server/core/objs/test_req.py @@ -1,7 +1,6 @@ import pytest import easydict -from types import SimpleNamespace -from lightllm.server.core.objs.req import Req, ChunkedPrefillReq, SamplingParams +from lightllm.server.core.objs.req import Req from lightllm.server.core.objs.token_metadata import ReqFinalTokenMetadata from lightllm.utils import shm_utils from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args @@ -85,17 +84,6 @@ def test_final_token_metadata_read_returns_actual_prompt_tokens(req): ] -def test_chunked_req_get_tuple_tokens_adds_page_and_async_reserve(): - req = SimpleNamespace( - input_len=10, - shm_cur_output_len=0, - shm_cur_kv_len=0, - sample_params=SimpleNamespace(ignore_eos=True, max_new_tokens=5), - ) - - assert ChunkedPrefillReq.get_tuple_tokens(req, False, 10) == (11, 26) - - def test_finish_status(req): req.finish_status.set_status(req.finish_status.FINISHED_STOP) assert req.finish_status.is_finished() diff --git a/unit_tests/server/httpserver/test_running_request_lifecycle.py b/unit_tests/server/httpserver/test_running_request_lifecycle.py index 0487e45281..3f6eeefd9f 100644 --- a/unit_tests/server/httpserver/test_running_request_lifecycle.py +++ b/unit_tests/server/httpserver/test_running_request_lifecycle.py @@ -554,12 +554,3 @@ async def run(): assert manager.run_reqs_count_mark.get_value() == 0 asyncio.run(run()) - - -def test_real_supported_max_req_total_len_reserves_mtp_overlap_tokens(): - manager = HttpServerManager.__new__(HttpServerManager) - manager.shm_max_total_token_num = _ValueMark(1000) - manager.max_req_total_len = 2000 - manager.args = SimpleNamespace(mtp_step=2, page_size=4) - - assert manager.get_real_supported_max_req_total_len() == 951 diff --git a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py index f04341a983..d6ae6ccbcf 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -95,7 +95,6 @@ def test_case3(): def test_case4(): - tree = RadixCache(100, 2) ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64, device="cpu")) ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu")) @@ -258,6 +257,8 @@ def test_case10(): ) assert tree_node is not None assert size == 3 + tree.dec_node_ref_counter(tree_node) + tree.mem_manager = SimpleNamespace(allocator=SimpleNamespace(can_use_mem_size=95), free=lambda indexes: None) tree.flush_cache() tree_node, size, values = tree.match_prefix( torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu"), update_refs=True @@ -632,11 +633,16 @@ def test_hybrid_request_release_frees_or_transfers_image_page_slots( tree, freed = hybrid_tree context = InferenceContext() + context.is_hybrid_att_model = True context.args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: context.args) context.radix_cache = tree - context.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(100, 132).view(1, -1)) + table = torch.arange(100, 132).view(1, -1) + context.req_manager = SimpleNamespace( + get_cpu_token_indexes=lambda row, start, end: table[row, start:end].clone(), + get_cpu_token_index_blocks=lambda row, start, end: table[row, start:end].numpy(), + ) req = InferReq.__new__(InferReq) req.req_idx = req.cur_output_len = 0 req.cur_kv_len, req.hold_kv_len, req.hybrid_cache_len = computed, max(24, computed), cache_len @@ -653,7 +659,9 @@ def test_hybrid_request_release_frees_or_transfers_image_page_slots( ) released = [] context._hybrid_att_free_req(released, req) - assert torch.cat(released).tolist() == list(range(100 + retained, 100 + max(24, computed))) + assert [index for block in released for index in block.tolist()] == list( + range(100 + retained, 100 + max(24, computed)) + ) assert not req.hybrid_len_to_big_page_id assert req.tail_small_page_buffer_id is None assert tree.get_tree_total_tokens_num() == retained diff --git a/unit_tests/server/router/model_infer/mode_backend/test_diverse_backend.py b/unit_tests/server/router/model_infer/mode_backend/test_diverse_backend.py deleted file mode 100644 index ab02d5892c..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_diverse_backend.py +++ /dev/null @@ -1,74 +0,0 @@ -from types import SimpleNamespace - -import torch - -# 沿用生产代码的模块加载顺序，避免独立运行本测试文件时形成循环导入。 -from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 -from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.server.router.model_infer.mode_backend.diverse_backend.impl import DiversehBackend - - -class _FakeRadixCache: - def __init__(self): - self.dec_nodes = [] - self.add_nodes = [] - - def dec_node_ref_counter(self, node): - self.dec_nodes.append(node) - - def add_node_ref_counter(self, node): - self.add_nodes.append(node) - - -def test_diverse_copy_skips_slaves_waiting_for_overlap_cleanup(): - active_slave = SimpleNamespace(filter_mark=False) - filtered_slave = SimpleNamespace(filter_mark=True) - master_req = SimpleNamespace(slave_reqs=[active_slave, filtered_slave]) - backend = DiversehBackend.__new__(DiversehBackend) - - batch_idx, run_reqs = backend._diverse_copy( - master_reqs=[master_req], - b_prefill_has_out=[True], - ) - - assert batch_idx == [0, 0] - assert run_reqs == [master_req, active_slave] - - -def test_copy_master_kv_updates_slave_hold_kv_len(): - req_to_token_indexs = torch.full((2, 8), -1, dtype=torch.int32) - req_to_token_indexs[0, :4] = torch.tensor([10, 11, 12, 13], dtype=torch.int32) - old_shared_node = object() - new_shared_node = object() - master_req = SimpleNamespace( - req_idx=0, - cur_kv_len=4, - cur_output_len=1, - shared_kv_node=new_shared_node, - ) - slave_req = SimpleNamespace( - req_idx=1, - cur_kv_len=2, - hold_kv_len=2, - cur_output_len=0, - shared_kv_node=old_shared_node, - related_master_req=master_req, - shm_req=SimpleNamespace(input_len=4), - args=SimpleNamespace(page_size=1), - ) - backend = DiversehBackend.__new__(DiversehBackend) - backend.model = SimpleNamespace( - req_manager=SimpleNamespace(req_to_token_indexs=req_to_token_indexs), - ) - backend.radix_cache = _FakeRadixCache() - backend.is_master_in_dp = False - - backend._copy_master_req_to_slave_req(slave_req) - - assert req_to_token_indexs[1, :4].tolist() == [10, 11, 12, 13] - assert slave_req.shared_kv_node is new_shared_node - assert slave_req.cur_kv_len == 4 - assert slave_req.hold_kv_len == 4 - assert InferReq._kv_cache_alloc_need(slave_req, target_kv_len=5) == 1 - assert backend.radix_cache.dec_nodes == [old_shared_node] - assert backend.radix_cache.add_nodes == [new_shared_node] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py index 42519a7415..574f17a938 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py @@ -1,31 +1,8 @@ from types import SimpleNamespace -import torch from lightllm.server.router.model_infer.mode_backend import generic_pre_process -def _patch_empty_input_context(monkeypatch): - mem_manager = SimpleNamespace( - HOLD_TOKEN_MEMINDEXES=(-1,), - alloc=lambda size: torch.empty((size,), dtype=torch.int32), - ) - infer_context = SimpleNamespace( - args=SimpleNamespace(page_size=1), - req_manager=SimpleNamespace( - HOLD_REQUEST_ID=-1, - mem_manager=mem_manager, - req_to_token_indexs=torch.zeros((16, 32), dtype=torch.int32), - ), - radix_cache=None, - ) - monkeypatch.setattr(generic_pre_process, "g_infer_context", infer_context) - return infer_context - - -def _patch_overlap_input_context(monkeypatch): - return _patch_empty_input_context(monkeypatch) - - def _make_prefill_req(req_idx: int, token_num: int): input_token_ids = [req_idx] * token_num return SimpleNamespace( @@ -53,8 +30,6 @@ def _make_decode_req(req_idx: int): def test_prepare_prefill_inputs_allows_empty_batch(monkeypatch): - _patch_empty_input_context(monkeypatch) - model_input, run_reqs = generic_pre_process.prepare_prefill_inputs([], is_chuncked_mode=True) assert run_reqs == [] @@ -68,8 +43,6 @@ def test_prepare_prefill_inputs_allows_empty_batch(monkeypatch): def test_prepare_decode_inputs_allows_empty_batch(monkeypatch): - _patch_empty_input_context(monkeypatch) - model_input, run_reqs = generic_pre_process.prepare_decode_inputs([]) assert run_reqs == [] @@ -84,7 +57,6 @@ def test_prepare_decode_inputs_allows_empty_batch(monkeypatch): def test_overlap_prefill_balances_request_token_load_without_padding(monkeypatch): - _patch_overlap_input_context(monkeypatch) reqs = [ _make_prefill_req(req_idx=0, token_num=8), _make_prefill_req(req_idx=1, token_num=7), @@ -111,7 +83,6 @@ def test_overlap_prefill_balances_request_token_load_without_padding(monkeypatch def test_overlap_prefill_balances_single_token_request_normally(monkeypatch): - _patch_overlap_input_context(monkeypatch) req = _make_prefill_req(req_idx=7, token_num=1) ( @@ -132,7 +103,6 @@ def test_overlap_prefill_balances_single_token_request_normally(monkeypatch): def test_overlap_decode_builds_two_unpadded_inputs(monkeypatch): - _patch_overlap_input_context(monkeypatch) reqs = [_make_decode_req(req_idx=index) for index in range(3)] ( @@ -155,7 +125,6 @@ def test_overlap_decode_builds_two_unpadded_inputs(monkeypatch): def test_overlap_decode_preserves_empty_microbatch(monkeypatch): - _patch_overlap_input_context(monkeypatch) req = _make_decode_req(req_idx=7) ( diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index 7fcbb94827..df0de5567e 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -42,15 +42,17 @@ def test_non_gpu_cache_tiers_release_owned_tokens_without_radix_insert(): released_refs = [] context = InferenceContext() context.is_hybrid_att_model = False - context.req_manager = SimpleNamespace(req_to_token_indexs=torch.tensor([[10, 11, 12, 13, 14, 15, 16]])) + context.req_manager = SimpleNamespace( + get_cpu_token_index_blocks=lambda row, start, end: torch.arange(10 + start, 10 + end).numpy() + ) context.radix_cache = SimpleNamespace(dec_node_ref_counter=released_refs.append) shared_node = SimpleNamespace(node_prefix_total_len=2) req = SimpleNamespace(req_idx=0, cur_kv_len=5, hold_kv_len=7, shared_kv_node=shared_node) - free_token_indexes = [] + free_token_index_blocks = [] - context._free_req_mem_without_radix_insert(free_token_indexes, req) + context._free_req_mem_without_radix_insert(free_token_index_blocks, req) - assert free_token_indexes[0].tolist() == [12, 13, 14, 15, 16] + assert free_token_index_blocks[0].tolist() == [12, 13, 14, 15, 16] assert released_refs == [shared_node] assert req.shared_kv_node is None @@ -60,7 +62,7 @@ def test_legacy_cache_tiers_still_insert_gpu_radix_cache(): context.radix_cache = object() context.is_hybrid_att_model = False inserted_reqs = [] - context._full_att_free_req = lambda free_token_index, req: inserted_reqs.append(req) + context._full_att_free_req = lambda free_token_index_blocks, req: inserted_reqs.append(req) req = SimpleNamespace( cache_tiers=(CacheTier.GPU, CacheTier.CPU, CacheTier.DISK), cur_kv_len=3, @@ -136,7 +138,9 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): freed_big_pages = [] context = InferenceContext() context.is_hybrid_att_model = True - context.req_manager = SimpleNamespace(req_to_token_indexs=torch.tensor([[10, 11, 12]])) + context.req_manager = SimpleNamespace( + get_cpu_token_index_blocks=lambda row, start, end: torch.arange(10 + start, 10 + end).numpy() + ) context.radix_cache = SimpleNamespace( small_page_buffers=SimpleNamespace(free_state_cache=freed_small_pages.extend), big_page_buffers=SimpleNamespace(free_state_cache=freed_big_pages.extend), @@ -149,11 +153,11 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): tail_small_page_buffer_id=7, hybrid_len_to_big_page_id={128: 8, 256: 9}, ) - free_token_indexes = [] + free_token_index_blocks = [] - context._free_req_mem_without_radix_insert(free_token_indexes, req) + context._free_req_mem_without_radix_insert(free_token_index_blocks, req) - assert free_token_indexes[0].tolist() == [10, 11, 12] + assert free_token_index_blocks[0].tolist() == [10, 11, 12] assert freed_small_pages == [7] assert freed_big_pages == [8, 9] assert req.tail_small_page_buffer_id is None @@ -174,7 +178,8 @@ def alloc(need_size): next_index += need_size return mem_indexes - def alloc_req_kv_mem(req, alloc_token_num): + def alloc_token_indexes(req, end, radix_cache): + alloc_token_num = end - req.hold_kv_len mem_indexes = alloc(alloc_token_num) new_hold_kv_len = req.hold_kv_len + alloc_token_num table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes @@ -191,7 +196,6 @@ def alloc_req_kv_mem(req, alloc_token_num): mem_manager=SimpleNamespace(operator=operator), req_manager=SimpleNamespace(req_to_token_indexs=table), ), - _alloc_req_kv_mem=alloc_req_kv_mem, ) module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) module.args = SimpleNamespace(page_size=4, cpu_cache_token_page_size=132) @@ -200,7 +204,8 @@ def alloc_req_kv_mem(req, alloc_token_num): module.need_sync_compute_stream = lambda: False module.cpu_cache_client = SimpleNamespace() context = SimpleNamespace( - req_manager=SimpleNamespace(req_to_token_indexs=table, mem_manager=SimpleNamespace(alloc=alloc)), + req_manager=SimpleNamespace(req_to_token_indexs=table, alloc_token_indexes=alloc_token_indexes), + radix_cache=None, get_can_alloc_token_num=lambda: 200, ) req = SimpleNamespace( diff --git a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py deleted file mode 100644 index 89e15fc93d..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py +++ /dev/null @@ -1,145 +0,0 @@ -from types import SimpleNamespace - -import numpy as np -import pytest -import torch - -# 先加载通用预处理模块，沿用生产初始化顺序，避免 basemodel 与 req_manager -# 在独立运行本测试文件时形成模块级循环导入。 -from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 -from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans -from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl -from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_impl - - -def _bind_alloc_need(req, page_size): - req.args = SimpleNamespace(page_size=page_size) - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - return req - - -def _reserve_into_table(table, req, alloc_sizes, next_index, alloc_token_num): - if alloc_token_num == 0: - return None - alloc_sizes.append(alloc_token_num) - new_hold_kv_len = req.hold_kv_len + alloc_token_num - mem_indexes = torch.arange(next_index[0], next_index[0] + alloc_token_num, dtype=torch.int32) - table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes - next_index[0] += alloc_token_num - req.hold_kv_len = new_hold_kv_len - return mem_indexes - - -@pytest.mark.parametrize("is_hybrid_att_model", [False, True]) -def test_pd_decode_reserves_model_pages_but_transfers_only_logical_kv(monkeypatch, is_hybrid_att_model): - table = torch.full((1, 16), -1, dtype=torch.int32) - table[0, :4] = torch.arange(4, dtype=torch.int32) - alloc_sizes = [] - next_index = [4] - tasks = [] - backend = decode_impl.PDDecodeNode.__new__(decode_impl.PDDecodeNode) - backend.args = SimpleNamespace(pd_kv_page_size=4, page_size=4) - backend.model = SimpleNamespace(req_manager=SimpleNamespace(req_to_token_indexs=table)) - backend.is_master_in_dp = False - backend._alloc_req_kv_mem = lambda req, size: _reserve_into_table(table, req, alloc_sizes, next_index, size) - - def create_task(**kwargs): - tasks.append( - ( - kwargs["kv_start_index"], - kwargs["kv_end_index"], - kwargs["mem_indexes"], - kwargs.get("page_kind", "kv"), - ) - ) - kwargs["group"].task_list.append(SimpleNamespace()) - - backend._create_pd_trans_task = create_task - monkeypatch.setattr(decode_impl.g_infer_context, "is_hybrid_att_model", is_hybrid_att_model) - req = _bind_alloc_need( - SimpleNamespace( - req_idx=0, - cur_kv_len=4, - hold_kv_len=4, - pd_trans_kv_start_index=0, - shm_req=SimpleNamespace(input_len=10), - ), - page_size=4, - ) - - backend._decode_node_gen_trans_tasks(req) - - assert alloc_sizes == [8] - assert req.cur_kv_len == 10 - assert req.hold_kv_len == 12 - expected_tasks = [ - (4, 8, [4, 5, 6, 7], "kv"), - (8, 10, [8, 9], "kv"), - ] - if is_hybrid_att_model: - expected_tasks.append((10, 10, [], "att_state")) - assert tasks == expected_tasks - assert table[0, :12].tolist() == list(range(12)) - - -def test_pd_decode_rejects_unaligned_transfer_page_size(): - backend = decode_impl.PDDecodeNode.__new__(decode_impl.PDDecodeNode) - backend.args = SimpleNamespace(pd_kv_page_size=3, page_size=4) - req = SimpleNamespace(cur_kv_len=4, shm_req=SimpleNamespace(input_len=10)) - - with pytest.raises(AssertionError, match="pd_kv_page_size must be divisible by page_size"): - backend._decode_node_gen_trans_tasks(req) - - -def test_pd_prefill_rejects_unaligned_transfer_page_size(): - backend = prefill_impl.PDChunkedPrefillForPrefillNode.__new__(prefill_impl.PDChunkedPrefillForPrefillNode) - backend.args = SimpleNamespace(pd_kv_page_size=3, page_size=4) - req = SimpleNamespace(cur_kv_len=4, shm_req=SimpleNamespace(input_len=10)) - - with pytest.raises(AssertionError, match="pd_kv_page_size must be divisible by page_size"): - backend._prefill_chuncked_handle_func(req, next_token_id=1, next_token_prob=1.0, output_len=0) - - -def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypatch): - table = torch.full((1, 16), -1, dtype=torch.int32) - table[0, :4] = torch.arange(4, dtype=torch.int32) - alloc_sizes = [] - next_index = [4] - req = _bind_alloc_need( - SimpleNamespace( - req_idx=0, - cur_kv_len=4, - hold_kv_len=4, - get_cur_total_len=lambda: 20, - ), - page_size=4, - ) - source_table = torch.arange(32, dtype=torch.int32).reshape(2, 16) - backend = SimpleNamespace( - model=SimpleNamespace(mem_manager=SimpleNamespace()), - dp_world_size=1, - rank_in_dp=0, - node_nccl_group=None, - mem_managers=[SimpleNamespace(req_to_token_indexs=source_table)] * 2, - ) - backend._alloc_req_kv_mem = lambda cur_req, size: _reserve_into_table(table, cur_req, alloc_sizes, next_index, size) - module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) - module.backend = backend - module.dp_rank_in_node = 0 - module.shared_req_infos = SimpleNamespace(arr=np.array([[[4, 0], [10, 0]]], dtype=np.int64)) - context = SimpleNamespace( - req_manager=SimpleNamespace(req_to_token_indexs=table), - get_can_alloc_token_num=lambda: 100, - ) - monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) - monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) - - tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) - - assert alloc_sizes == [8] - assert req.cur_kv_len == 4 - assert req.hold_kv_len == 12 - assert len(tasks) == 1 - assert tasks[0].mem_indexes.tolist() == list(range(4, 10)) - assert len(tasks[0].max_kv_len_mem_indexes) == 6 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py index c498093fe2..95f9157b3b 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py @@ -3,6 +3,7 @@ import pytest +from lightllm.common.req_manager import ReqManager from lightllm.server.core.objs import FinishStatus from lightllm.server.core.objs import req as req_module from lightllm.server.router.model_infer.mode_backend import base_backend @@ -24,13 +25,14 @@ def _make_infer_req(cur_output_len: int, shm_output_len: int): finish_status=FinishStatus(), cpu_cache_task_status=SimpleNamespace(is_not_started=MagicMock(return_value=True)), cur_kv_len=10, + hold_kv_len=8, cur_output_len=cur_output_len, shm_req=SimpleNamespace(shm_cur_output_len=shm_output_len), sampling_param=SimpleNamespace( shm_param=SimpleNamespace(max_new_tokens=65535), ), get_cur_total_len=MagicMock(return_value=11), - decode_need_token_num=MagicMock(return_value=(1, 1)), + decode_need_token_num=MagicMock(return_value=1), ) @@ -51,7 +53,11 @@ def _classify_without_token_capacity(monkeypatch, req, support_overlap=True): backend._filter_not_ready_reqs = MagicMock(return_value=reqs) backend.prefill_queue_strategy = SimpleNamespace(reorder=MagicMock(side_effect=lambda reqs, **kwargs: reqs)) + backend.batch_max_tokens = 16 infer_context = base_backend.g_infer_context + manager = ReqManager.__new__(ReqManager) + manager.mem_manager = SimpleNamespace(page_size=4, allocator=SimpleNamespace(can_use_mem_size=0)) + monkeypatch.setattr(infer_context, "req_manager", manager) monkeypatch.setattr(infer_context, "get_can_alloc_token_num", MagicMock(return_value=0)) monkeypatch.setattr( infer_context, diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py b/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py deleted file mode 100644 index 9e44ed774b..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py +++ /dev/null @@ -1,52 +0,0 @@ -from types import SimpleNamespace - -import torch - -from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend - - -def test_preload_prompt_cache_keeps_only_complete_kv_pages(monkeypatch): - loaded = {} - - class FakeMemManager: - def alloc(self, size): - loaded["alloc_size"] = size - return torch.arange(size, dtype=torch.int32) - - def load_index_kv_buffer(self, indexes, buffers): - loaded["indexes"] = indexes.clone() - loaded["buffers"] = buffers - - class FakeRadixCache: - def __init__(self): - self.mem_manager = FakeMemManager() - - def insert(self, token_ids, indexes): - loaded["insert"] = (token_ids.clone(), indexes.clone()) - - def match_prefix(self, token_ids, update_refs): - loaded["match"] = (token_ids.clone(), update_refs) - - backend = ModeBackend.__new__(ModeBackend) - backend.logger = SimpleNamespace(info=lambda _: None) - backend.weight_dir = "/model" - backend.args = SimpleNamespace(page_size=4) - backend.radix_cache = FakeRadixCache() - kv_buffer = torch.arange(12).reshape(1, 6, 2) - monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) - monkeypatch.setattr(torch, "load", lambda *args, **kwargs: {"kv_buffer": kv_buffer}) - - backend.preload_prompt_cache_kv_buffer( - { - "prompt_cache_kv_buffer": {"rank_0": "prompt_cache.pt"}, - "prompt_cache_token_ids": [10, 11, 12, 13, 14, 15], - } - ) - - assert loaded["alloc_size"] == 4 - assert loaded["indexes"].tolist() == [0, 1, 2, 3] - assert loaded["buffers"]["kv_buffer"].shape == (1, 4, 2) - assert loaded["insert"][0].tolist() == [10, 11, 12, 13] - assert loaded["insert"][1].tolist() == [0, 1, 2, 3] - assert loaded["match"][0].tolist() == [10, 11, 12, 13] - assert loaded["match"][1] is True diff --git a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py b/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py deleted file mode 100644 index dff0e892a9..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py +++ /dev/null @@ -1,33 +0,0 @@ -from types import SimpleNamespace - -import pytest - -from lightllm.server.router.model_infer.mode_backend.rl_backend_ops import RlBackendOps -from lightllm.utils.torch_memory_saver_utils import MemoryTag - - -@pytest.mark.parametrize( - ("tags", "should_init_hold_indexes"), - [ - ([MemoryTag.WEIGHT], False), - ([MemoryTag.KV_CACHE], True), - (None, True), - ], -) -def test_resume_memory_initializes_hold_indexes_only_for_kv_cache(tags, should_init_hold_indexes): - events = [] - model = SimpleNamespace( - torch_memory_saver=SimpleNamespace(resume=lambda tag: events.append(("resume", tag))), - req_manager=SimpleNamespace(init_hold_request_indexs=lambda: events.append(("init_hold_indexes", None))), - ) - backend = SimpleNamespace(model=model, logger=None) - ops = RlBackendOps(backend) - ops._clear_cuda_cache = lambda: None - - ops._resume_memory_tags(tags) - - init_events = [event for event in events if event[0] == "init_hold_indexes"] - assert bool(init_events) is should_init_hold_indexes - if should_init_hold_indexes: - last_resume_index = max(index for index, event in enumerate(events) if event[0] == "resume") - assert events.index(("init_hold_indexes", None)) > last_resume_index diff --git a/unit_tests/server/router/model_infer/structured_output/test_state.py b/unit_tests/server/router/model_infer/structured_output/test_state.py index 4843d4dac3..22edee9bc0 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_state.py +++ b/unit_tests/server/router/model_infer/structured_output/test_state.py @@ -167,6 +167,7 @@ def read_artifact(req_id=req_id, owner_dp_rank=owner_dp_rank): context = infer_batch.InferenceContext( req_manager=SimpleNamespace( alloc=count().__next__, + copy_token_index_prefixes_to_gpu=lambda req_prefix_lens: None, req_sampling_params_manager=SimpleNamespace(init_req_sampling_params=lambda req: None), ), shm_req_manager=SimpleNamespace(get_req_obj_by_index=lambda index: shared_reqs[index]), diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index 189101c950..38d8dab766 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -105,6 +105,7 @@ def test_full_attention_gpu_restore_stays_outside_image( from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context + from lightllm.common.req_manager import ReqManager from lightllm.utils import shm_utils from lightllm.utils.kv_cache_utils import CpuKVCacheMeta, compute_token_list_hash @@ -119,11 +120,12 @@ def test_full_attention_gpu_restore_stays_outside_image( monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", False) monkeypatch.setattr(g_infer_context, "radix_cache", radix_cache) - monkeypatch.setattr( - g_infer_context, - "req_manager", - SimpleNamespace(req_to_token_indexs=torch.empty((1, len(query)), dtype=torch.int64)), - ) + req_manager = ReqManager.__new__(ReqManager) + req_manager.mem_manager = SimpleNamespace(page_size=page_size) + req_manager._cpu_token_index_blocks = torch.full( + (1, (len(query) + page_size - 1) // page_size), -1, dtype=torch.int32 + ).numpy() + monkeypatch.setattr(g_infer_context, "req_manager", req_manager) req = InferReq.__new__(InferReq) req.args = SimpleNamespace(page_size=page_size) req.sampling_param = SimpleNamespace(disable_prompt_cache=False) @@ -131,8 +133,11 @@ def test_full_attention_gpu_restore_stays_outside_image( req.shared_kv_node = None req.mm_input_spans = [(700, 1200)] req.shm_req = SimpleNamespace(input_len=len(query), shm_prompt_ids=SimpleNamespace(arr=query)) - req._match_radix_cache() + req_prefix_lens = {} + req._match_radix_cache(req_prefix_lens=req_prefix_lens) assert req.cur_kv_len == req.hold_kv_len == gpu_end + assert req_prefix_lens == ({0: gpu_end} if gpu_end else {}) + assert req_manager.get_cpu_token_index_blocks(0, 0, gpu_end).tolist() == query[:gpu_end:page_size] assert not any(start < gpu_end < end for start, end in req.mm_input_spans) assert radix_cache.get_refed_tokens_num() == gpu_end @@ -255,15 +260,14 @@ def test_cpu_cache_match_and_load_use_a_complete_restore_point( ) table = torch.arange(8192, dtype=torch.int32).view(1, -1) - def allocate(req, count): - req.hold_kv_len += count - return torch.arange(req.cur_kv_len, req.cur_kv_len + count, dtype=torch.int32) + def allocate(req, end, radix_cache): + req.hold_kv_len = end + return torch.arange(req.cur_kv_len, end, dtype=torch.int32) module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) module.args = SimpleNamespace(page_size=256) module.backend = SimpleNamespace( is_master_in_dp=True, - _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table), mem_manager=SimpleNamespace( @@ -276,6 +280,11 @@ def allocate(req, count): module.cpu_cache_client = client module.init_sync_group = None module.need_sync_compute_stream = lambda: False + monkeypatch.setattr( + common.g_infer_context, + "req_manager", + SimpleNamespace(alloc_token_indexes=allocate), + ) monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 8192) monkeypatch.setattr(common.g_infer_context, "is_hybrid_att_model", hybrid) monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) diff --git a/unit_tests/server/test_dp_request_capacity.py b/unit_tests/server/test_dp_request_capacity.py index 1e63550f45..a0ae03d1df 100644 --- a/unit_tests/server/test_dp_request_capacity.py +++ b/unit_tests/server/test_dp_request_capacity.py @@ -24,8 +24,8 @@ def test_hybrid_capacity_uses_local_dp_count(monkeypatch, global_capacity, dp, n assert ShmReqManager.get_max_req_num(None) == global_capacity -@pytest.mark.parametrize("mode", ["full_attention", "diverse", "cache_fetch"]) -def test_modes_requiring_existing_capacity_keep_global_default(monkeypatch, mode): +@pytest.mark.parametrize("mode,expected", [("full_attention", 4), ("diverse", 32), ("cache_fetch", 32)]) +def test_capacity_depends_on_mode(monkeypatch, mode, expected): monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: mode != "full_attention") args = StartArgs( running_max_req_size=32, @@ -33,7 +33,7 @@ def test_modes_requiring_existing_capacity_keep_global_default(monkeypatch, mode diverse_mode=mode == "diverse", enable_dp_prompt_cache_fetch=mode == "cache_fetch", ) - assert get_running_max_req_size_per_dp(args) == 32 + assert get_running_max_req_size_per_dp(args) == expected @pytest.mark.parametrize("capacity", [0, -1]) diff --git a/unit_tests/server/test_pd_start_args.py b/unit_tests/server/test_pd_start_args.py deleted file mode 100644 index def9608eb8..0000000000 --- a/unit_tests/server/test_pd_start_args.py +++ /dev/null @@ -1,75 +0,0 @@ -import pytest - -from lightllm.server.api_start import _launch_subprocesses -from lightllm.server.core.objs.start_args_type import StartArgs - - -def test_pd_kv_page_size_must_be_divisible_by_model_page_size(monkeypatch): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - args = StartArgs( - run_mode="decode", - page_size=4, - pd_kv_page_size=6, - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - ) - - with pytest.raises(AssertionError, match="--pd_kv_page_size must be divisible by --page_size"): - _launch_subprocesses(args) - - -def test_normal_mode_does_not_validate_pd_kv_page_size(monkeypatch): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda model_dir: False) - - def validation_finished(args): - raise RuntimeError("PD page-size validation passed") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) - args = StartArgs( - run_mode="normal", - model_dir="unused", - page_size=4, - pd_kv_page_size=6, - eos_id=[2], - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - ) - - with pytest.raises(RuntimeError, match="PD page-size validation passed"): - _launch_subprocesses(args) - - -@pytest.mark.parametrize("llm_kv_type", ["fp8kv_sph", "fp8kv_spt"]) -def test_fp8_kv_cache_allows_multi_token_model_pages(monkeypatch, llm_kv_type): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda model_dir: False) - - def validation_finished(args): - raise RuntimeError("FP8 page-size validation passed") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) - args = StartArgs( - llm_kv_type=llm_kv_type, - kv_quant_calibration_config_path="unused", - model_dir="unused", - page_size=16, - eos_id=[2], - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - ) - - with pytest.raises(RuntimeError, match="FP8 page-size validation passed"): - _launch_subprocesses(args) diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py deleted file mode 100644 index 6b57270125..0000000000 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ /dev/null @@ -1,153 +0,0 @@ -import sys -from pathlib import Path -from types import SimpleNamespace - -import numpy as np -import torch - - -REPO_ROOT = Path(__file__).resolve().parents[1] -STATIC_BENCHMARK_DIR = REPO_ROOT / "test" / "benchmark" / "static_inference" -for path in (REPO_ROOT, STATIC_BENCHMARK_DIR): - if str(path) not in sys.path: - sys.path.append(str(path)) - -from static_benchmark import ( - BenchmarkCase, - StaticBenchmarkExecutor, - decode_profile_batch_divisor, - resolve_batch_max_prefill_cases, -) - - -class FakeMemManager: - def __init__(self, page_size: int): - self.page_size = page_size - self.next_index = 0 - self.alloc_sizes = [] - - def alloc(self, size: int) -> torch.Tensor: - assert size % self.page_size == 0 - self.alloc_sizes.append(size) - indexes = torch.arange(self.next_index, self.next_index + size, dtype=torch.int32) - self.next_index += size - return indexes - - -def make_executor(page_size: int = 4): - mem_manager = FakeMemManager(page_size=page_size) - req_manager = SimpleNamespace( - mem_manager=mem_manager, - req_to_token_indexs=torch.full((4, 32), -1, dtype=torch.int32), - ) - model = SimpleNamespace(mem_manager=mem_manager, req_manager=req_manager) - args = SimpleNamespace(mtp_mode=None, mtp_step=0) - executor = StaticBenchmarkExecutor(args=args, model=model, draft_models=[], token_source=None) - return executor, mem_manager, req_manager - - -def assert_complete_physical_pages(indexes: torch.Tensor, page_size: int): - for start in range(0, indexes.numel(), page_size): - page = indexes[start : start + page_size] - assert torch.all(page // page_size == page[0] // page_size) - - -def test_prefill_allocates_and_reuses_complete_pages_per_request(): - executor, mem_manager, req_manager = make_executor(page_size=4) - req_idx = torch.tensor([0, 1], dtype=torch.int32) - - model_input = executor._make_prefill_input( - token_chunk=np.arange(6, dtype=np.int64).reshape(2, 3), - req_idx=req_idx, - ready_cache_len=0, - ) - - assert mem_manager.alloc_sizes == [4, 4] - assert model_input.b_is_decode_req.tolist() == [False, False] - assert executor._hold_kv_len_by_req == {0: 4, 1: 4} - assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :4], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :4], page_size=4) - - # 两个请求都还在已有页内继续 decode，不应产生新的分配。 - executor._make_decode_input( - batch_size=2, - req_idx=req_idx, - mtp_index=torch.zeros(2, dtype=torch.int32), - seq_len=torch.tensor([4, 4], dtype=torch.int32), - input_ids=torch.tensor([1, 2], dtype=torch.int64), - max_kv_seq_len=4, - ) - assert mem_manager.alloc_sizes == [4, 4] - - # 跨过页边界后，每个请求分别追加一整页。 - executor._make_decode_input( - batch_size=2, - req_idx=req_idx, - mtp_index=torch.zeros(2, dtype=torch.int32), - seq_len=torch.tensor([5, 5], dtype=torch.int32), - input_ids=torch.tensor([3, 4], dtype=torch.int64), - max_kv_seq_len=5, - ) - assert mem_manager.alloc_sizes == [4, 4, 4, 4] - assert executor._hold_kv_len_by_req == {0: 8, 1: 8} - assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :8], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :8], page_size=4) - - -def test_mtp_expanded_rows_allocate_once_per_request(): - executor, mem_manager, req_manager = make_executor(page_size=4) - executor.args.mtp_mode = "eagle_with_att" - executor.args.mtp_step = 3 - - model_input = executor._make_decode_input( - batch_size=4, - req_idx=torch.tensor([0, 0, 1, 1], dtype=torch.int32), - mtp_index=torch.tensor([0, 1, 0, 1], dtype=torch.int32), - seq_len=torch.tensor([3, 4, 3, 4], dtype=torch.int32), - input_ids=torch.tensor([1, 2, 3, 4], dtype=torch.int64), - max_kv_seq_len=4, - extra_kv_len=2, - ) - - assert model_input.batch_size == 4 - assert mem_manager.alloc_sizes == [8, 8] - assert executor._hold_kv_len_by_req == {0: 8, 1: 8} - assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :8], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :8], page_size=4) - - -def test_profile_capacity_uses_page_aligned_per_request_cost(): - args = SimpleNamespace( - mtp_mode=None, - mtp_step=0, - page_size=4, - max_batch_size=0, - batch_max_tokens=64, - ) - prefill_case = BenchmarkCase( - name="prefill", - stage="prefill", - batch_size=8, - context_len=5, - output_len=0, - prefill_uncached_len=5, - prefill_batch_size_by_batch_max_tokens=8, - ) - - resolved = resolve_batch_max_prefill_cases(args, [prefill_case], profiled_max_total_token_num=24) - - assert resolved[0].batch_size == 3 - assert resolved[0].profiled_batch_divisor == 8 - - decode_case = BenchmarkCase( - name="decode", - stage="decode", - batch_size=1, - context_len=6, - output_len=2, - ) - assert decode_profile_batch_divisor(args, decode_case) == 20 - - args.mtp_mode = "eagle_with_att" - args.mtp_step = 3 - assert decode_profile_batch_divisor(args, decode_case) == 24 diff --git a/unit_tests/utils/test_config_utils_fixed_kv.py b/unit_tests/utils/test_config_utils_fixed_kv.py deleted file mode 100644 index 116e299e24..0000000000 --- a/unit_tests/utils/test_config_utils_fixed_kv.py +++ /dev/null @@ -1,33 +0,0 @@ -from types import SimpleNamespace - -import pytest - -from lightllm.utils import config_utils - - -@pytest.fixture(autouse=True) -def clear_fixed_kv_len_cache(): - config_utils.get_fixed_kv_len.cache_clear() - yield - config_utils.get_fixed_kv_len.cache_clear() - - -def test_fixed_kv_len_supports_page_size_one(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=1, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {"prompt_cache_token_ids": [1, 2, 3]}) - - assert config_utils.get_fixed_kv_len() == 3 - - -def test_fixed_kv_len_uses_only_complete_pages(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=2, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {"prompt_cache_token_ids": [1, 2, 3]}) - - assert config_utils.get_fixed_kv_len() == 2 - - -def test_no_fixed_kv_supports_paged_kv(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=4, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {}) - - assert config_utils.get_fixed_kv_len() == 0 diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py index a06e779392..a4843bd2fc 100644 --- a/unit_tests/utils/test_envs_utils.py +++ b/unit_tests/utils/test_envs_utils.py @@ -23,11 +23,11 @@ def test_stop_in_reasoning_reads_environment(monkeypatch): is_stop_matching_in_reasoning_enabled.cache_clear() -def test_pd_cache_high_priority_max_age_defaults_to_36_seconds(monkeypatch): +def test_pd_cache_high_priority_max_age_defaults_to_180_seconds(monkeypatch): monkeypatch.delenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MAX_AGE_SECONDS", raising=False) get_pd_cache_high_priority_max_age_seconds.cache_clear() - assert get_pd_cache_high_priority_max_age_seconds() == 36 + assert get_pd_cache_high_priority_max_age_seconds() == 180 get_pd_cache_high_priority_max_age_seconds.cache_clear() @@ -41,11 +41,11 @@ def test_pd_cache_high_priority_max_age_reads_environment_variable(monkeypatch): get_pd_cache_high_priority_max_age_seconds.cache_clear() -def test_pd_cache_high_priority_min_prompt_tokens_defaults_to_4096(monkeypatch): +def test_pd_cache_high_priority_min_prompt_tokens_defaults_to_2048(monkeypatch): monkeypatch.delenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MIN_PROMPT_TOKENS", raising=False) get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - assert get_pd_cache_high_priority_min_prompt_tokens() == 4096 + assert get_pd_cache_high_priority_min_prompt_tokens() == 2048 get_pd_cache_high_priority_min_prompt_tokens.cache_clear() @@ -59,11 +59,11 @@ def test_pd_cache_high_priority_min_prompt_tokens_reads_environment_variable(mon get_pd_cache_high_priority_min_prompt_tokens.cache_clear() -def test_pd_node_resource_wait_timeout_defaults_to_10_seconds(monkeypatch): +def test_pd_node_resource_wait_timeout_defaults_to_20_seconds(monkeypatch): monkeypatch.delenv("LIGHTLLM_PD_NODE_RESOURCE_WAIT_TIMEOUT_SECONDS", raising=False) get_pd_node_resource_wait_timeout_seconds.cache_clear() - assert get_pd_node_resource_wait_timeout_seconds() == 10 + assert get_pd_node_resource_wait_timeout_seconds() == 20 get_pd_node_resource_wait_timeout_seconds.cache_clear() diff --git a/unit_tests/utils/test_sgl_utils.py b/unit_tests/utils/test_sgl_utils.py index b8af0a91c5..85bb041e68 100644 --- a/unit_tests/utils/test_sgl_utils.py +++ b/unit_tests/utils/test_sgl_utils.py @@ -117,7 +117,6 @@ def test_fa3_decode_preserves_actual_length_before_graph_capture(monkeypatch, pa state = fa3_module.Fa3DecodeAttState( backend=SimpleNamespace( model=model, - page_size=page_size, infer_page_size=page_size, uses_causal_attention=lambda: True, uses_dynamic_spec_verify_layout=lambda: False, diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py index aa8182cde8..ef2f139f14 100644 --- a/unit_tests/utils/test_speculative_utils.py +++ b/unit_tests/utils/test_speculative_utils.py @@ -126,7 +126,7 @@ def test_fa3_prefill_state_owns_causality(monkeypatch, state_class): infer_state.input_ids = torch.empty(2, dtype=torch.int64) infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, dtype=torch.int32).reshape(2, 4)) state = state_class( - backend=SimpleNamespace(uses_causal_attention=lambda: False, page_size=1, infer_page_size=1), + backend=SimpleNamespace(uses_causal_attention=lambda: False, infer_page_size=1), _infer_state_weak_ref=weakref.ref(infer_state), ) From a1bd40b6a1878ca9ec42862b752a1d55577b03e1 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Wed, 14 Oct 2026 18:02:39 +0800 Subject: [PATCH 2/2] fix: keep KV allocator local in IPC exports and restore page regressions --- .../kv_cache_mem_manager/mem_manager.py | 8 + .../model_infer/mode_backend/base_backend.py | 2 - skills/test_model/SKILL.md | 31 ++- .../test_copy_kv_index_to_req.py | 92 ++++++++ unit_tests/common/test_mem_manager_ipc.py | 93 ++++++++ unit_tests/common/test_req_manager_page.py | 214 ++++++++++++++++++ .../test_hybrid_cpu_state_cache.py | 8 +- 7 files changed, 422 insertions(+), 26 deletions(-) create mode 100644 unit_tests/common/basemodel/triton_kernel/test_copy_kv_index_to_req.py create mode 100644 unit_tests/common/test_mem_manager_ipc.py create mode 100644 unit_tests/common/test_req_manager_page.py diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 8dfd7f1721..59a23d866f 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -59,6 +59,14 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False # 构建对外的操作类接口 self.operator: BaseMemManagerOperator = self.operator_class(self) + def __getstate__(self): + # IPC consumers access KV buffers and request indexes, while allocation + # remains owned by the inference process. NumPy pools would otherwise + # be embedded in every descriptor written by write_to_shm. + state = self.__dict__.copy() + state["allocator"] = None + return state + def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: k = self.kv_buffer[layer_index][:, : self.head_num, :] v = self.kv_buffer[layer_index][:, self.head_num :, :] 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 1ff50a82ac..aee51e8971 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -614,8 +614,6 @@ def _timer_merge_radix_tree(self): ) return - - def _classify_reqs_and_alloc_kv(self, candidates, available): """按优先级顺序接纳请求，再一次性分配所有必需的 token 容量。 diff --git a/skills/test_model/SKILL.md b/skills/test_model/SKILL.md index f5880ea268..579defb879 100644 --- a/skills/test_model/SKILL.md +++ b/skills/test_model/SKILL.md @@ -80,20 +80,17 @@ lm_eval --model local-completions \ 设置表中 `--page_size` 和 `--cpu_cache_token_page_size`。257 特意不整除物理页大小， 用于覆盖 chunk 跨页边界。按子 skill 的方法检查服务日志、端口，并完成真实请求 warmup。 -服务就绪后执行（变量对应本组实际启动参数）： - -```bash -python test/acc/paged_cpu_cache_probe.py \ - --url "http://127.0.0.1:${PORT}" --model-dir "${MODEL_DIR}" \ - --server-log "${LOG_DIR}/server.log" --output "${LOG_DIR}/cpu-cache-probe.json" \ - --page-size "${PAGE_SIZE}" --cpu-page-size "${CPU_PAGE_SIZE}" --kv-capacity 32768 -``` - -Qwen3.5 的探针额外传入 `--hash-page-size 256`，与服务端 hybrid hash page 保持一致。 - -探针要求回载日志包含实际 CPU cache 命中，并覆盖部分 GPU 前缀同时命中的情况； -GPU 常驻前缀参考与 CPU 回载使用相同的 prefill 起点，避免 chunk/query 形状改变产生的 -BF16 舍入差异干扰传输校验；32 个输出 token 必须相同，logprob 最大绝对差不超过 0.02。 -另外检查 prompt logprobs 禁用缓存复用后仍返回全部输入位置的有效结果。 -GPU↔CPU KV 逐值一致性、碎片页和异步引用生命周期由 -`unit_tests/server/router/model_infer/mode_backend/test_paged_cpu_cache.py` 覆盖。 +仓库目前没有自动执行 GPU 前缀逐出后真实 CPU cache 回载的探针。服务就绪后， +用以下人工步骤记录请求参数、响应与服务日志，不能把连续两轮精度评测当作已覆盖逐出回载： + +1. 对固定提示词使用 greedy 采样，记录 GPU 常驻前缀参考的 32 个输出 token 和 logprob。 +2. 用其他提示词填充 KV cache，确认目标 GPU 前缀已被逐出且其 CPU cache 仍可用。 +3. 再次请求目标提示词，确认日志中有实际 CPU cache 命中；另测仍有部分 GPU 前缀命中的情况。 + 参考与回载请求须使用相同的 prefill 起点，避免 chunk/query 形状改变造成的 BF16 舍入差异。 + 32 个输出 token 必须相同，logprob 最大绝对差不超过 0.02。 +4. 请求 prompt logprobs，检查禁用缓存复用后仍返回全部输入位置的有效结果。 + +Qwen3.5 服务的 `--linear_att_hash_page_size` 保持为 256。现有 +`unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py` +覆盖后端缓存任务状态和引用释放；它使用模拟传输，不能证明真实 GPU↔CPU KV 逐值一致性。 +未完成上述人工验证或逐值传输校验时，在结果中明确列为验证缺口。 diff --git a/unit_tests/common/basemodel/triton_kernel/test_copy_kv_index_to_req.py b/unit_tests/common/basemodel/triton_kernel/test_copy_kv_index_to_req.py new file mode 100644 index 0000000000..193e45efc7 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_copy_kv_index_to_req.py @@ -0,0 +1,92 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import build_kv_indexes_and_input_ids + + +@pytest.mark.parametrize("batch", [0, 1, 7, 257]) +@pytest.mark.parametrize("gather", [False, True]) +def test_fused_decode_matches_independent_indexing(batch, gather): + table = torch.arange(12 * 2048, device="cuda", dtype=torch.int32).reshape(12, 2048)[:, ::2] + next_ids = torch.arange(12 * 4, device="cuda", dtype=torch.int64).reshape(12, 4) + 100 + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) % 12 + ends = torch.arange(batch, device="cuda", dtype=torch.int32) % 1023 + 1 + steps = torch.arange(batch, device="cuda", dtype=torch.int32) % 4 + ids = None if gather else torch.full((batch,), 55, dtype=torch.int64, device="cuda") + model_input = SimpleNamespace( + is_prefill=False, + batch_size=batch, + b_req_idx=reqs, + input_ids=ids, + max_q_seq_len=1, + b_seq_len=ends, + b_mtp_index=steps, + b_is_decode_req=None, + ) + indexes, result_ids = build_kv_indexes_and_input_ids(model_input, table, next_ids if gather else None) + assert torch.equal(indexes, table[reqs.long(), ends.long() - 1]) + assert torch.equal(result_ids, next_ids[reqs.long(), steps.long()] if gather else ids) + + +@pytest.mark.parametrize("max_q", [1, 255, 513, 4097]) +def test_fused_mixed_prefill_matches_independent_indexing(max_q): + lengths = [max_q, 1, min(7, max_q)] + starts = [0, lengths[0], lengths[0] + lengths[1]] + ready = [2, 3, 4] + table = torch.arange(8 * 10000, device="cuda", dtype=torch.int32).reshape(8, 10000)[:, ::2] + next_ids = torch.arange(8 * 4, device="cuda", dtype=torch.int64).reshape(8, 4) + 900 + ids = torch.arange(sum(lengths), device="cuda", dtype=torch.int64) + expected_ids = ids.clone() + expected_ids[starts[1]] = next_ids[3, 2] + model_input = SimpleNamespace( + is_prefill=True, + batch_size=3, + b_req_idx=torch.tensor([1, 3, 5], device="cuda", dtype=torch.int32), + input_ids=ids, + max_q_seq_len=max_q, + b_seq_len=torch.tensor([r + q for r, q in zip(ready, lengths)], device="cuda", dtype=torch.int32), + b_ready_cache_len=torch.tensor(ready, device="cuda", dtype=torch.int32), + b_prefill_start_loc=torch.tensor(starts, device="cuda", dtype=torch.int32), + b_mtp_index=torch.tensor([0, 2, 0], device="cuda", dtype=torch.int32), + b_is_decode_req=torch.tensor([False, True, False], device="cuda"), + ) + expected_indexes = torch.cat([table[req, r : r + q] for req, r, q in zip([1, 3, 5], ready, lengths)]) + indexes, result_ids = build_kv_indexes_and_input_ids(model_input, table, next_ids) + assert torch.equal(indexes, expected_indexes) + assert torch.equal(result_ids, expected_ids) + + +def test_fused_decode_graph_replay_reads_current_tables(): + table = torch.arange(8 * 32, dtype=torch.int32, device="cuda").reshape(8, 32) + next_ids = torch.arange(8 * 4, dtype=torch.int64, device="cuda").reshape(8, 4) + reqs = torch.tensor([0, 1, 2, 7], dtype=torch.int32, device="cuda") + ends = torch.tensor([4, 5, 6, 2], dtype=torch.int32, device="cuda") + steps = torch.tensor([0, 1, 2, 0], dtype=torch.int32, device="cuda") + model_input = SimpleNamespace( + is_prefill=False, + batch_size=4, + b_req_idx=reqs, + input_ids=None, + max_q_seq_len=1, + b_seq_len=ends, + b_mtp_index=steps, + b_is_decode_req=None, + ) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + build_kv_indexes_and_input_ids(model_input, table, next_ids) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + indexes, ids = build_kv_indexes_and_input_ids(model_input, table, next_ids) + for iteration in range(3): + table.add_(1000) + next_ids.add_(100) + ends.add_(1) + graph.replay() + assert torch.equal(indexes, table[reqs.long(), ends.long() - 1]) + assert torch.equal(ids, next_ids[reqs.long(), steps.long()]) diff --git a/unit_tests/common/test_mem_manager_ipc.py b/unit_tests/common/test_mem_manager_ipc.py new file mode 100644 index 0000000000..9aff9bde31 --- /dev/null +++ b/unit_tests/common/test_mem_manager_ipc.py @@ -0,0 +1,93 @@ +from multiprocessing.reduction import ForkingPickler + +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.common.kv_cache_mem_manager import allocator as allocator_module +from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager +from lightllm.common.kv_cache_mem_manager.operator import NormalMemOperator + + +class _Counter: + def __init__(self, *args): + pass + + def set_value(self, value): + self.value = value + + +def _make_manager(monkeypatch, size, device="cpu"): + monkeypatch.setattr(allocator_module, "SharedInt", _Counter) + monkeypatch.setattr(allocator_module, "get_current_rank_in_node", lambda: 0) + manager = MemoryManager.__new__(MemoryManager) + manager.size = size + manager.page_size = 1 + manager.allocator = allocator_module.KvCacheAllocator(size) + manager.kv_buffer = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 4, 2, 4) + manager.req_to_token_indexs = torch.arange(4, dtype=torch.int32, device=device).reshape(1, 4) + manager.operator = NormalMemOperator(manager) + return manager + + +def test_ipc_descriptor_size_does_not_scale_with_local_allocator(monkeypatch): + small = _make_manager(monkeypatch, 256) + large = _make_manager(monkeypatch, 131072) + small_payload = ForkingPickler.dumps(small) + large_payload = ForkingPickler.dumps(large) + print(f"IPC pickle sizes: small={len(small_payload)} large={len(large_payload)}") + assert len(large_payload) < 4096 + assert abs(len(large_payload) - len(small_payload)) < 128 + + +def test_ipc_roundtrip_retains_shared_buffers_and_local_owner(monkeypatch): + source = _make_manager(monkeypatch, 256) + allocator = source.allocator + exported = ForkingPickler.loads(ForkingPickler.dumps(source)) + assert source.allocator is allocator + assert allocator.alloc(1).tolist() == [0] + assert exported.allocator is None + assert exported.operator.mem_manager is exported + exported.kv_buffer.fill_(17) + exported.req_to_token_indexs.fill_(23) + assert torch.all(source.kv_buffer == 17) + assert torch.all(source.req_to_token_indexs == 23) + + +def _update_exported_buffers(payload, connection): + try: + exported = ForkingPickler.loads(payload) + assert exported.allocator is None + assert exported.operator.mem_manager is exported + exported.kv_buffer.fill_(17) + exported.req_to_token_indexs.fill_(23) + if exported.kv_buffer.is_cuda: + torch.cuda.synchronize() + connection.send("complete") + finally: + connection.close() + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_ipc_buffers_remain_shared_between_processes(monkeypatch, device): + source = _make_manager(monkeypatch, 256, device) + allocator = source.allocator + context = mp.get_context("spawn") + receiver, sender = context.Pipe(duplex=False) + process = context.Process(target=_update_exported_buffers, args=(ForkingPickler.dumps(source).tobytes(), sender)) + try: + process.start() + sender.close() + assert receiver.poll(60), "IPC consumer did not complete" + assert receiver.recv() == "complete" + process.join(10) + assert process.exitcode == 0 + assert source.allocator is allocator + assert torch.all(source.kv_buffer == 17) + assert torch.all(source.req_to_token_indexs == 23) + finally: + if process.is_alive(): + process.terminate() + process.join(10) + receiver.close() + sender.close() diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py new file mode 100644 index 0000000000..95b7a6eb30 --- /dev/null +++ b/unit_tests/common/test_req_manager_page.py @@ -0,0 +1,214 @@ +from types import MethodType, SimpleNamespace + +import numpy as np +import pytest +import torch + +from lightllm.common.kv_cache_mem_manager import allocator as allocator_module +from lightllm.common.req_manager import base as req_module +from lightllm.common.req_manager import req_sampling_params +from lightllm.server.router.model_infer.infer_batch import InferReq +from lightllm.server.router.model_infer.mode_backend import base_backend + + +class _Counter: + def __init__(self, *args): + pass + + def set_value(self, value): + self.value = value + + +def _make_manager(monkeypatch, page_size, count=12): + monkeypatch.setattr(allocator_module, "SharedInt", _Counter) + monkeypatch.setattr(allocator_module, "get_current_rank_in_node", lambda: 0) + monkeypatch.setattr(req_module, "get_env_start_args", lambda: SimpleNamespace(page_size=page_size)) + monkeypatch.setattr(req_sampling_params, "ReqSamplingParamsManager", lambda size: None) + alloc = allocator_module.KvCacheAllocator(page_size * 200, page_size) + memory = SimpleNamespace( + page_size=page_size, allocator=alloc, HOLD_TOKEN_MEMINDEXES=tuple(range(alloc.size, alloc.size + page_size)) + ) + return req_module.ReqManager(count, page_size * 10, memory) + + +@pytest.mark.parametrize("page_size", [1, 2, 4, 8, 64]) +@pytest.mark.parametrize("count", [0, 1, 4, 5, 9]) +def test_alloc_publish_reuse_and_owned_snapshot(monkeypatch, page_size, count): + manager = _make_manager(monkeypatch, page_size) + requests = [SimpleNamespace(req_idx=i, hold_kv_len=0, cur_kv_len=0) for i in range(count)] + allocations = [(req, page_size * (1 + i % 3)) for i, req in enumerate(requests)] + manager.alloc_req_token_indexes(allocations) + torch.cuda.synchronize() + expected_offset = 0 + snapshots = [] + for req, size in allocations: + expected = torch.arange(expected_offset, expected_offset + size, dtype=torch.int32) + assert torch.equal(manager.req_to_token_indexs[req.req_idx, :size].cpu(), expected) + snapshot = manager.get_cpu_token_indexes(req.req_idx, 0, size) + assert torch.equal(snapshot, expected) + snapshots.append((snapshot, expected)) + assert req.hold_kv_len == size + expected_offset += size + blocks = [manager.get_cpu_token_index_blocks(req.req_idx, 0, size).copy() for req, size in allocations] + if blocks: + manager.free_token_index_blocks(np.concatenate(blocks)) + for req, _ in allocations: + req.hold_kv_len = 0 + manager.alloc_req_token_indexes(list(reversed(allocations))) + torch.cuda.synchronize() + for snapshot, expected in snapshots: + assert torch.equal(snapshot, expected) + for req, size in allocations: + start = min(1, size) + assert torch.equal( + manager.get_cpu_token_indexes(req.req_idx, start, size), + manager.req_to_token_indexs[req.req_idx, start:size].cpu(), + ) + + +@pytest.mark.parametrize("page_size", [1, 4, 64]) +def test_prefix_cross_stream_publish(monkeypatch, page_size): + manager = _make_manager(monkeypatch, page_size) + writer, reader = torch.cuda.Stream(), torch.cuda.Stream() + source = torch.arange(page_size * 3, dtype=torch.int32) + page_size * 20 + with torch.cuda.stream(writer): + prefixes = {} + manager.write_token_index_prefix(0, 0, source, prefixes) + manager.copy_token_index_prefixes_to_gpu(prefixes) + with torch.cuda.stream(reader): + manager.copy_token_index_prefix(0, 1, len(source)) + observed = manager.req_to_token_indexs[1, : len(source)].clone() + reader.synchronize() + assert torch.equal(observed.cpu(), source) + assert torch.equal(manager.get_cpu_token_indexes(1, 0, len(source)), source) + + +@pytest.mark.parametrize("page_size", [1, 4, 64]) +def test_alloc_token_indexes_returns_logical_span_and_reuses_page_tail(monkeypatch, page_size): + manager = _make_manager(monkeypatch, page_size) + req = SimpleNamespace(req_idx=0, hold_kv_len=0, cur_kv_len=0) + indexes = manager.alloc_token_indexes(req, page_size + 1) + assert indexes.tolist() == list(range(page_size + 1)) + expected_hold = (page_size + 1 + page_size - 1) // page_size * page_size + assert req.hold_kv_len == expected_hold + capacity = manager.mem_manager.allocator.can_use_mem_size + req.cur_kv_len = page_size + tail = manager.alloc_token_indexes(req, expected_hold) + assert tail.tolist() == list(range(page_size, expected_hold)) + assert manager.mem_manager.allocator.can_use_mem_size == capacity + + +def _make_scheduler(monkeypatch, page_size, available, batch_max_tokens=8): + manager = _make_manager(monkeypatch, page_size) + context = SimpleNamespace( + req_manager=manager, + radix_cache=None, + get_can_alloc_token_num=lambda: available, + cache_placement_controller=SimpleNamespace(set_req_cache_way=lambda reqs: None), + filter_reqs=lambda finished_reqs: None, + pause_reqs=lambda reqs, is_master_in_dp: None, + ) + monkeypatch.setattr(base_backend, "g_infer_context", context) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + backend.args = SimpleNamespace(enable_cpu_cache=False, enable_prefill_decode_mixed=False, run_mode="normal") + backend.support_overlap = False + backend.disable_chunked_prefill = False + backend.batch_max_tokens = batch_max_tokens + backend.is_master_in_dp = True + backend._timer_merge_radix_tree = lambda: None + backend.prefill_queue_strategy = SimpleNamespace(reorder=lambda reqs, **kwargs: reqs) + allocations = [] + original_alloc = manager.alloc_req_token_indexes + + def record_alloc(items, radix_cache=None): + allocations.append(sum(size for _, size in items)) + return original_alloc(items, radix_cache) + + monkeypatch.setattr(manager, "alloc_req_token_indexes", record_alloc) + return manager, backend, allocations + + +def _make_scheduler_req(req_idx, current, total): + req = SimpleNamespace( + req_idx=req_idx, + cur_kv_len=current, + hold_kv_len=0, + mtp_step=0, + filter_mark=False, + wait_pause=False, + paused=False, + infer_aborted=False, + finish_status=SimpleNamespace(is_finished=lambda: False), + is_slave_req=lambda: False, + get_cur_total_len=total, + get_chuncked_input_token_len=total, + ) + req.prefill_need_token_num = MethodType(InferReq.prefill_need_token_num, req) + req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) + return req + + +@pytest.mark.parametrize( + "hold_kv_len, available, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause", + [(4, 8, 3, [4], False), (8, 0, 3, [], False), (4, 4, 6, [4], True)], +) +def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( + monkeypatch, hold_kv_len, available, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause +): + manager, backend, allocations = _make_scheduler(monkeypatch, 4, available, batch_max_tokens) + reqs = [_make_scheduler_req(i, 3, lambda: 6) for i in range(2)] + manager.alloc_req_token_indexes([(req, hold_kv_len) for req in reqs]) + allocations.clear() + backend._filter_not_ready_reqs = lambda req_ids: reqs + + prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0, 1]) + + assert prefill_reqs == [reqs[0]] + assert decode_reqs == [] + assert prefill_tokens == 3 + assert allocations == expected_alloc_sizes + assert reqs[0].hold_kv_len == 8 + assert not reqs[0].wait_pause + assert reqs[1].hold_kv_len == hold_kv_len + assert reqs[1].wait_pause is second_req_wait_pause + manager.wait_token_indexes().synchronize() + assert torch.equal(manager.req_to_token_indexs[0, :8].cpu(), manager.get_cpu_token_indexes(0, 0, 8)) + + +def test_decode_scheduler_publishes_batched_request_table(monkeypatch): + manager, backend, allocations = _make_scheduler(monkeypatch, 4, 12) + req = _make_scheduler_req(0, 4, lambda: 5) + manager.alloc_req_token_indexes([(req, 4)]) + allocations.clear() + backend._filter_not_ready_reqs = lambda req_ids: [req] + + prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0]) + + assert prefill_reqs == [] + assert decode_reqs == [req] + assert prefill_tokens == 0 + assert req.hold_kv_len == 16 + assert allocations == [12] + manager.wait_token_indexes().synchronize() + assert manager.req_to_token_indexs[0, :16].tolist() == list(range(16)) + + +@pytest.mark.parametrize("page_size, expected_alloc_token_num", [(1, 9), (3, 9), (4, 12), (7, 14), (8, 8)]) +def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_steps( + monkeypatch, page_size, expected_alloc_token_num +): + manager, backend, allocations = _make_scheduler(monkeypatch, page_size, 32) + req = _make_scheduler_req(0, page_size, lambda: req.cur_kv_len + 1) + manager.alloc_req_token_indexes([(req, page_size)]) + allocations.clear() + backend._filter_not_ready_reqs = lambda req_ids: [req] + + for _ in range(expected_alloc_token_num): + prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0]) + assert prefill_reqs == [] + assert decode_reqs == [req] + assert prefill_tokens == 0 + req.cur_kv_len += 1 + + assert req.hold_kv_len == page_size + expected_alloc_token_num + assert allocations == [expected_alloc_token_num] diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 7b60c15c4a..ddec21da0c 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -67,6 +67,7 @@ def linear_cache(monkeypatch, tmp_path, request): SimpleNamespace( get_big_page_ids_by_node=lambda node: [], free_one_small_page_buffer=lambda: None, + free_radix_cache_to_get_enough_token=lambda need: None, small_page_buffers=requests.small_page_buffers, ), ) @@ -144,16 +145,9 @@ def test_cpu_roundtrip_keeps_kv_continuous_and_restores_only_saved_states( big.conv_state_cache.buffer.zero_() big.ssm_state_cache.buffer.zero_() - def allocate(req, count): - allocated = mem.alloc(count) - requests.req_to_token_indexs[req.req_idx, :count] = allocated.cuda() - req.hold_kv_len = count - return allocated - module.args = get_env_start_args() module.backend = SimpleNamespace( is_master_in_dp=True, - _alloc_req_kv_mem=allocate, model=SimpleNamespace(mem_manager=mem, req_manager=requests), ) module.init_sync_group = None