From 15f6e1972e2e1dea59aa1fb64c6d4e79e7b93e7b Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 22 Sep 2026 13:17:41 +0800 Subject: [PATCH 01/22] 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 | 199 ++----- 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 | 274 +++------ .../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 | 142 ++--- .../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 | 241 ++++++-- .../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 | 351 ++++------- .../model_infer/mode_backend/base_backend.py | 166 ++---- .../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 | 22 +- .../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 | 3 +- .../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 | 550 ------------------ .../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 | 51 +- .../mode_backend/test_diverse_backend.py | 74 --- .../mode_backend/test_generic_pre_process.py | 31 - .../mode_backend/test_multi_level_kv_cache.py | 90 +-- .../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 -- .../req_queue/test_pd_queue_selection.py | 109 ---- 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 | 2 +- unit_tests/utils/test_speculative_utils.py | 2 +- 74 files changed, 1298 insertions(+), 3288 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/router/req_queue/test_pd_queue_selection.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 9c9b811954..5ee71e629e 100644 --- a/lightllm/common/basemodel/attention/fa3/fp.py +++ b/lightllm/common/basemodel/attention/fa3/fp.py @@ -2,7 +2,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 e124562f02..7cf6a5d524 100644 --- a/lightllm/common/basemodel/attention/fa3/mla.py +++ b/lightllm/common/basemodel/attention/fa3/mla.py @@ -2,7 +2,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 5a2d5b77d8..06a053d200 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,29 @@ 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 + ): + model_input = self._create_padded_decode_model_input(model_input, batch_size) + infer_state = self._create_inferstate(model_input, microbatch_index) + # Attention backends need the capture flag before initializing their state. + infer_state.is_cuda_graph = 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 +422,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 +479,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 +541,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 +555,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 +563,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,13 +585,7 @@ 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 会根据该标记准备 CUDA Graph capture 专用状态， - # 因此必须在 init_att_state 之前完成赋值。 - infer_state.is_cuda_graph = 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) if use_cuda_graph: if need_capture: @@ -630,7 +599,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" @@ -728,15 +696,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): @@ -754,24 +713,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( @@ -787,8 +730,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() @@ -797,19 +738,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): @@ -820,21 +748,16 @@ 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 = 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 = 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) + infer_state1 = self._create_decode_inferstate(model_input1, infer_batch_size, need_capture, 1) + + if use_cuda_graph: if need_capture: model_output0, model_output1 = self.graph.capture_decode( self._overlap_tpsp_token_forward, @@ -847,23 +770,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..9f4182d54f 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,102 @@ 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: 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) + offset = tl.load(Packet + 4 * row + 3) + pos = block * BLOCK + tl.arange(0, BLOCK) + value = tl.load(Packet + offset + pos // PAGE, pos < size, other=0) + pos % PAGE + 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..f92305a4b9 100644 --- a/lightllm/common/kv_cache_mem_manager/allocator.py +++ b/lightllm/common/kv_cache_mem_manager/allocator.py @@ -1,103 +1,71 @@ +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 - + """A stack of physical page bases; returned token indexes own their storage.""" + + 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_pages(self, need_size: int): + """Borrow page bases until the next allocator mutation; callers snapshot before DMA.""" + assert need_size % self.page_size == 0 + assert 0 <= need_size <= self.can_use_mem_size, "error alloc state" + start = (self.size - self.can_use_mem_size) // self.page_size self.can_use_mem_size -= need_size self.shared_can_use_token_num.set_value(self.can_use_mem_size) - - # 利用缓冲区返回，避免异步情况下的内存竞争 - if self._return_start + need_size > self._mem_state_return.shape[0]: - 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 + return self.free_pages[start : start + need_size // self.page_size] + + def alloc(self, need_size: int) -> torch.Tensor: + pages = self._alloc_pages(need_size) + indexes = torch.empty(need_size, dtype=torch.int32, device="cpu", pin_memory=True) + self.expand_pages(pages, indexes.numpy()) + return indexes + + def expand_pages(self, pages, out=None): + if out is None: + out = np.empty(len(pages) * self.page_size, dtype=np.int32) + if self.page_size == 1: + out[:] = pages else: - # 从 gpu 到 cpu 的拷贝操作是流内阻塞操作 - self.mem_state[start:end] = free_index - - self.mark_start -= len(free_index) - - self.can_use_mem_size += len(free_index) + # Short pages expand faster along the page-base axis than in short token runs. + np.add( + pages, + 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.release_pages(values[:: self.page_size]) + + def release_pages(self, pages): + size = len(pages) * 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_pages[end - len(pages) : end] = pages + 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_pages[:] = 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_pages = 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..b6781b8f95 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,15 @@ 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" ) + self._cpu_pages = 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 +78,192 @@ 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_pages[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_page_bases(self, req_idx, start, end): + """Return a CPU view of page bases for the token range [start, end).""" + page_size = self.mem_manager.page_size + assert start % page_size == end % page_size == 0 + return self._cpu_pages[req_idx, start // page_size : end // page_size] + + def get_cpu_token_indexes(self, req_idx, start, end): + """Materialize owned token indexes only for cache insertion and PD metadata.""" + page_size = self.mem_manager.page_size + first, offset = divmod(start, page_size) + bases = self._cpu_pages[req_idx, first : (end + page_size - 1) // page_size] + indexes = self.mem_manager.allocator.expand_pages(bases) + return torch.from_numpy(indexes[offset : offset + end - start]) + + 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_pages[req_idx, start // page_size : (start + 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.mem_manager.page_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_page_bases(dest_req_idx, 0, size)[:] = self.get_cpu_page_bases(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_pages(self, allocations, radix_cache=None): + """Reserve pages once; CPU and GPU tables consume the same page-base assignment.""" + 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 + pages = self.mem_manager.allocator._alloc_pages(need) + updates, offset = [], 0 + for req, size in allocations: + start = req.hold_kv_len + if size == page_size: + self._cpu_pages[req.req_idx, start // page_size] = pages[offset] + else: + self._cpu_pages[req.req_idx, start // page_size : (start + size) // page_size] = pages[ + offset : offset + size // page_size + ] + updates.append((req.req_idx, start, size, offset)) + req.hold_kv_len += size + offset += size // page_size + self._copy_token_indexes_to_gpu(updates, pages) + + 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_pages([(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, pages=None): + """Snapshot compact pages, using direct DMA for small batches and GPU expansion otherwise.""" + from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import update_req_token_indexes + + page_size = self.mem_manager.page_size + if pages is None: + parts = [ + self._cpu_pages[row, start // page_size : (start + size) // page_size] + for row, start, size, _ in updates + ] + pages = parts[0] if len(parts) == 1 else np.concatenate(parts) + if not len(pages): + return + packed = len(updates) > 4 + header = 4 * len(updates) if packed else 0 + packet = torch.empty( + header + len(pages) * (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 + meta[:, 3] += header + data[header:] = pages + else: + self.mem_manager.allocator.expand_pages(pages, data) + 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 * page_size : offset * page_size + 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 + size = max(0, (target + page_size - 1) // page_size * page_size - req.hold_kv_len) + assert size % page_size == 0 + return size + + def classify_reqs_and_alloc_kv(self, candidates, available, batch_max_tokens, radix_cache=None): + """按优先级顺序接纳请求，再一次性预留所有必需的页。 + + prefill 和 decode 两组请求都会跨调度轮次保留已分配的页。 + decode 的可选额外预留仅使用满足所有必要分配后剩余的空闲容量。 + """ + prefill, decode, rejected, allocations, headroom = [], [], [], [], [] + prefill_tokens = required = 0 + page_size = self.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 = self.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, self.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: + self.alloc_req_pages(allocations, radix_cache) + return prefill, decode, rejected, prefill_tokens + + 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_page_index): for req_index in free_req_indexes: self.req_list.free(req_index) if self.req_list.is_all_free(): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") - self.mem_manager.free(free_token_index) + self.mem_manager.allocator.release_pages(free_page_index) def free_req(self, free_req_index: int): self.req_list.free(free_req_index) @@ -112,10 +271,10 @@ 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) + def free_pages(self, free_page_index): + self.mem_manager.allocator.release_pages(free_page_index) return 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 87271b6145..b9ed18648b 100644 --- a/lightllm/common/req_manager/req_sampling_params.py +++ b/lightllm/common/req_manager/req_sampling_params.py @@ -35,6 +35,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. self.req_to_next_token_scores = ( torch.zeros_like(self.req_to_next_token_ids, dtype=torch.float32) if get_env_start_args().mtp_dynamic_verify @@ -81,15 +82,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 58929b1cc1..c66090bfe9 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 c7fc111a64..d2fab91e8d 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -29,7 +29,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 e9d462a256..0a70aea5cc 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 d72cb3edd2..b627ed22ef 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 ebc2035877..afc2560a73 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.req_id_generator 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.server.pd_io_struct import PDDecodeNodeInfo from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient @@ -90,6 +87,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, _ = r assert r_id not in self.requests_mapping.keys() @@ -100,11 +98,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 下，建立一组请求间的主从关系 @@ -127,25 +127,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_page_index: 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_page_index.append(self.req_manager.get_cpu_page_bases(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_page_index=free_page_index, req=req) else: if not self.is_hybrid_att_model: - self._full_att_free_req(free_token_index=free_token_index, req=req) + self._full_att_free_req(free_page_index=free_page_index, req=req) else: - self._hybrid_att_free_req(free_token_index=free_token_index, req=req) + self._hybrid_att_free_req(free_page_index=free_page_index, 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_page_index: 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_page_index.append(self.req_manager.get_cpu_page_bases(req.req_idx, shared_kv_len, req.hold_kv_len)) if self.is_hybrid_att_model: # 释放请求尾部尚未移交给 radix cache 的 hybrid attention 小页状态。 @@ -163,27 +163,27 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer 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_page_index: 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_page_index, req, cache_kv_len) + + def _cache_req_prefix(self, free_page_index: 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_page_index.append(self.req_manager.get_cpu_page_bases(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_page_index.append(self.req_manager.get_cpu_page_bases(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_page_index: 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 @@ -203,83 +203,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_page_index, 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_page_index, + 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): """ @@ -303,7 +252,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_page_index = [] for request_id in finished_request_ids: req: InferReq = self.requests_mapping.pop(request_id) if self.args.diverse_mode: @@ -312,7 +261,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_page_index, req) free_req_index.append(req.req_idx) # logger.info(f"infer release req id {req.shm_req.request_id}") @@ -322,9 +271,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_page_index: + free_page_index = np.concatenate(free_page_index) + self.req_manager.free(free_req_index, free_page_index) 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] @@ -348,13 +297,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_page_index = [] 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_page_index, req) assert req.wait_pause is True req.wait_pause = False req.paused = True @@ -362,27 +310,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 len(free_page_index) != 0: + free_page_index = np.concatenate(free_page_index) + self.req_manager.free_pages(free_page_index) 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 @@ -390,6 +339,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): @@ -457,7 +407,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 @@ -533,6 +483,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 @@ -602,9 +554,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): @@ -639,7 +591,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..." @@ -650,19 +609,12 @@ def _match_radix_cache(self): key = key[0 : len(key) - 1] # 最后一个不需要，因为需要一个额外的token，让其在prefill的时候输出下一个token的值 share_node, kv_len, value_tensor = g_infer_context.radix_cache.match_prefix(key, update_refs=True) 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..." @@ -684,94 +636,41 @@ 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 ) + 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) + 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) - 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 @@ -937,39 +836,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 + return target_kv_len - self.cur_kv_len - def decode_need_token_num(self) -> Tuple[int, int]: - """返回 (decode 需要的 token 数, 需要额外分配的 KV token 容量)。 - - 第一个值沿用分页分配前的需求估算：普通 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 22f3cbe5dc..9c14559076 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 @@ -411,11 +411,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: @@ -555,9 +555,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, @@ -643,30 +641,6 @@ def remaining_prefill_tokens(req: InferReq) -> int: ready_reqs.insert(0, target_req) return ready_reqs - # 一些可以复用的通用功能函数 - 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 _get_classed_reqs( self, req_ids: List[int] = None, @@ -691,6 +665,8 @@ def _get_classed_reqs( 3. finished_reqs 需要释放的请求, 包含正常结束和aborted退出的请求。 4. prefill_reqs 需要进行prefill操作的请求 5. decode_reqs 需要进行decode操作的请求 + + 为通过预算检查的两组请求预留 KV，保持尚未执行请求的跨轮容量。 """ # 定期对 radix cache 进行 merge，防止查询插入的操作效率下降 self._timer_merge_radix_tree() @@ -712,100 +688,70 @@ 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: - - 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.paused: - paused_reqs.append(req_obj) - continue + def candidates(): + for req_obj in ready_reqs: + if req_obj.filter_mark: + finished_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 + if req_obj.wait_pause: + wait_pause_reqs.append(req_obj) continue - else: - finished_reqs.append(req_obj) + + if req_obj.paused: + paused_reqs.append(req_obj) continue - if no_decode: - is_decode = False - else: - is_decode = req_obj.cur_kv_len + 1 == req_obj.get_cur_total_len() - if is_decode and strict_prefill and req_obj.cur_kv_len + 1 == req_obj.shm_req.input_len: + 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 + + if no_decode: is_decode = False + else: + is_decode = req_obj.cur_kv_len + 1 == req_obj.get_cur_total_len() + if is_decode and strict_prefill and req_obj.cur_kv_len + 1 == req_obj.shm_req.input_len: + is_decode = False - if is_decode: - # 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 + if is_decode: + token_num = req_obj.decode_need_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 + # 在 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 - ) - 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 + token_num = req_obj.prefill_need_token_num(is_chuncked_prefill=not self.disable_chunked_prefill) + yield req_obj, is_decode, token_num + + prefill_reqs, decode_reqs, rejected, prefill_tokens = g_infer_context.req_manager.classify_reqs_and_alloc_kv( + candidates(), can_alloc_token_num, self.batch_max_tokens, g_infer_context.radix_cache + ) + 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 else: - if wait_pause_count < pause_max_req_num: - req_obj.wait_pause = True - wait_pause_count += 1 + finished_reqs.append(req_obj) + self.logger.info( + f"force early finish for PD decode req_id={req_obj.req_id} because token capacity is insufficient" + ) + else: + req_obj.wait_pause = True # 先由控制器确定请求需要写入的缓存层级，再按是否包含 CPU cache 决定是否发起 offload。 cache_controller = g_infer_context.cache_placement_controller @@ -832,10 +778,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， @@ -957,7 +900,6 @@ def _sample_and_scatter_token( b_prefill_has_output_cpu: torch.Tensor = None, mask_func: Optional[Callable] = None, ): - if mask_func is not None: assert len(run_reqs) == logits.shape[0] mask_func(run_reqs, logits) 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 dedd3bb4a6..aad675cff4 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 12a2e9a048..f4b2a8f76d 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 @@ -95,7 +92,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: # 计算需要加载的页面（只加载未匹配的部分） @@ -103,8 +99,9 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): assert ready_page_num <= len(page_list) need_pages = page_list[ready_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 现在必须使用同步模式, 未来需要移除 @@ -114,8 +111,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，所以 @@ -278,6 +275,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..1ae59f8698 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,10 @@ 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_pages( + self.model.req_manager.get_cpu_page_bases(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 +124,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 0c25b8949a..2fbd8329df 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 f7f7173ebb..07b8ef32f5 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..1cc7329b9d 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_pages(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 3f6664d031..ca035fdaf2 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -3,7 +3,6 @@ import torch import pytest -from lightllm.common.basemodel import basemodel from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput @@ -131,7 +130,7 @@ def pad_once(model_input, new_batch_size): graph_flags_at_att_init = [] - 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 55ef9fcf37..46b9610c0d 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 @@ -24,6 +24,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 f45ac26ceb..0000000000 --- a/unit_tests/common/test_req_manager_page.py +++ /dev/null @@ -1,550 +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._reorder_pd_high_priority_reqs = lambda reqs: reqs - backend._reorder_long_prefill_reqs = lambda reqs: 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 = backend._get_classed_reqs(req_ids=[0, 1]) - - assert prefill_reqs == [reqs[0]] - assert decode_reqs == [] - 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._reorder_pd_high_priority_reqs = lambda reqs: reqs - backend._reorder_long_prefill_reqs = lambda reqs: 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 = backend._get_classed_reqs(req_ids=[0]) - - assert prefill_reqs == [] - assert decode_reqs == [req] - 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._reorder_pd_high_priority_reqs = lambda reqs: reqs - backend._reorder_long_prefill_reqs = lambda reqs: 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 3cae97ac51..823ca9d4fd 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 @@ -80,17 +79,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 a75af88847..fee9fb621b 100644 --- a/unit_tests/server/httpserver/test_running_request_lifecycle.py +++ b/unit_tests/server/httpserver/test_running_request_lifecycle.py @@ -418,12 +418,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 80265ac40d..8a7f016400 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -92,7 +92,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")) @@ -255,6 +254,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 @@ -269,53 +270,5 @@ def test_case10(): assert tree.root_node.ref_counter == 1 -def test_page_aligned_insert_and_match(): - tree = RadixCache(100, 99, page_size=4) - values = torch.arange(100, 110, dtype=torch.int64) - - prefix_len, _ = tree.insert(torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), values) - assert prefix_len == 0 - assert tree.get_tree_total_tokens_num() == 8 - - # A mismatch inside a page cannot produce a partial-page cache hit. - node, matched_len, matched_values = tree.match_prefix(torch.tensor([1, 2, 0, 4, 5, 6, 7, 8])) - assert node is None - assert matched_len == 0 - assert matched_values is None - - node, matched_len, matched_values = tree.match_prefix(torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 11])) - assert node is not None - assert matched_len == 8 - assert matched_values.tolist() == list(range(100, 108)) - - # The second sequence shares exactly one page and then branches by its - # complete second-page key. - prefix_len, _ = tree.insert( - torch.tensor([1, 2, 3, 4, 5, 6, 0, 8]), - torch.arange(200, 208, dtype=torch.int64), - ) - assert prefix_len == 4 - assert tree.get_tree_total_tokens_num() == 12 - - -def test_page_size_must_match_mem_manager(): - mem_manager = SimpleNamespace(page_size=8) - - with pytest.raises(ValueError, match="must match mem_manager page_size 8"): - RadixCache(100, 100, mem_manager=mem_manager, page_size=4) - - -def test_page_key_bytes_does_not_share_tensor_memory(): - token_ids = torch.tensor([1, 2, 3, 4], dtype=torch.int64) - expected_key = token_ids.clone().numpy().tobytes() - key = TreeNode(page_size=4).get_child_key(token_ids) - - token_ids[0] = 100 - - assert isinstance(key, bytes) - assert key == expected_key - assert key != TreeNode(page_size=4).get_child_key(token_ids) - - if __name__ == "__main__": pytest.main() 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 260dbe1065..e16964c0db 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_page_bases=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_page_indexes = [] - context._free_req_mem_without_radix_insert(free_token_indexes, req) + context._free_req_mem_without_radix_insert(free_page_indexes, req) - assert free_token_indexes[0].tolist() == [12, 13, 14, 15, 16] + assert free_page_indexes[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_page_index, 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_page_bases=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,80 +153,12 @@ 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_page_indexes = [] - context._free_req_mem_without_radix_insert(free_token_indexes, req) + context._free_req_mem_without_radix_insert(free_page_indexes, req) - assert free_token_indexes[0].tolist() == [10, 11, 12] + assert free_page_indexes[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 assert req.hybrid_len_to_big_page_id == {} - - -def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch): - table = torch.full((1, 136), -1, dtype=torch.int32) - table[0, :4] = torch.arange(4, dtype=torch.int32) - next_index = 4 - alloc_sizes = [] - loaded_indexes = [] - - def alloc(need_size): - nonlocal next_index - alloc_sizes.append(need_size) - mem_indexes = torch.arange(next_index, next_index + need_size, dtype=torch.int32) - next_index += need_size - return mem_indexes - - def alloc_req_kv_mem(req, alloc_token_num): - 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 - req.hold_kv_len = new_hold_kv_len - return mem_indexes - - operator = SimpleNamespace( - load_cpu_cache_to_gpu=lambda mem_indexes, **kwargs: loaded_indexes.extend(mem_indexes.tolist()) - ) - backend = SimpleNamespace( - is_master_in_dp=False, - radix_cache=None, - model=SimpleNamespace( - 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) - module.backend = backend - module.init_sync_group = object() - 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)), - get_can_alloc_token_num=lambda: 200, - ) - req = SimpleNamespace( - req_idx=0, - cur_kv_len=4, - hold_kv_len=4, - shm_req=SimpleNamespace( - input_len=256, - disk_prompt_cache_len=0, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: [7]), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: [132]), - ), - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - ) - monkeypatch.setattr(multi_level_kv_cache_impl, "g_infer_context", context) - monkeypatch.setattr(multi_level_kv_cache_impl.dist, "barrier", lambda group: None) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) - - module.load_cpu_cache_to_reqs([req]) - - assert alloc_sizes == [128] - assert req.cur_kv_len == 132 - assert req.hold_kv_len == 132 - assert loaded_indexes == list(range(132)) - assert table[0, :132].tolist() == list(range(132)) 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 73e133cc18..24ca0c9ecd 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), ) @@ -52,7 +54,11 @@ def _classify_without_token_capacity(monkeypatch, req, support_overlap=True): backend._reorder_pd_high_priority_reqs = MagicMock(side_effect=lambda reqs: reqs) backend._reorder_long_prefill_reqs = MagicMock(side_effect=lambda reqs: 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/req_queue/test_pd_queue_selection.py b/unit_tests/server/router/req_queue/test_pd_queue_selection.py deleted file mode 100644 index 7c8cd5cbf0..0000000000 --- a/unit_tests/server/router/req_queue/test_pd_queue_selection.py +++ /dev/null @@ -1,109 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import MagicMock - -from lightllm.server.router.req_queue import _get_req_queue_class -from lightllm.server.router.batch import Batch -from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_decode import PDDecodeQueue -from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_prefill import PDPrefillQueue - - -def _make_args(run_mode: str): - return SimpleNamespace( - run_mode=run_mode, - diverse_mode=False, - output_constraint_mode="none", - first_token_constraint_mode=False, - disable_chunked_prefill=False, - ) - - -def test_pd_prefill_uses_prefill_queue(): - args = _make_args("prefill") - - assert _get_req_queue_class(args, router=None, dp_size_in_node=1) is PDPrefillQueue - - -def test_pd_decode_uses_decode_queue(): - args = _make_args("decode") - - assert _get_req_queue_class(args, router=None, dp_size_in_node=1) is PDDecodeQueue - - -def test_pd_prefill_peak_tokens_sum_request_tuples(): - queue = PDPrefillQueue.__new__(PDPrefillQueue) - queue.dp_index = 0 - queue.args = SimpleNamespace(page_size=16) - queue.is_busy = lambda: False - queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) - reqs = [ - SimpleNamespace( - request_id="req-0", - input_len=10, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), - get_tuple_tokens=MagicMock(return_value=(11, 34)), - ), - SimpleNamespace( - request_id="req-1", - input_len=20, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), - get_tuple_tokens=MagicMock(return_value=(21, 34)), - ), - SimpleNamespace( - request_id="req-2", - input_len=100, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=1), - get_tuple_tokens=MagicMock(return_value=(101, 34)), - ), - ] - batch = Batch(batch_id=1, reqs=reqs, dp_size_in_node=2) - - assert queue._caclu_batch_estimated_peak_token_num(batch) == 100 - reqs[0].get_tuple_tokens.assert_called_once_with(False, 10) - reqs[1].get_tuple_tokens.assert_called_once_with(False, 10) - reqs[2].get_tuple_tokens.assert_not_called() - - -def test_pd_decode_sums_request_tuples_for_non_decode_requests(): - queue = PDDecodeQueue.__new__(PDDecodeQueue) - queue.args = SimpleNamespace(page_size=16) - queue.dp_index = 0 - queue.is_busy = lambda: False - queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) - reqs = [ - SimpleNamespace( - request_id="req-0", - input_len=20, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), - is_infer_decode=lambda: False, - get_tuple_tokens=MagicMock(return_value=(21, 34)), - ), - SimpleNamespace( - request_id="req-1", - input_len=100, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=1), - is_infer_decode=lambda: False, - get_tuple_tokens=MagicMock(return_value=(101, 34)), - ), - ] - batch = Batch(batch_id=1, reqs=reqs, dp_size_in_node=2) - - assert queue._caclu_batch_estimated_peak_token_num(batch) == 55 - reqs[0].get_tuple_tokens.assert_called_once_with(False, 10) - reqs[1].get_tuple_tokens.assert_not_called() - - -def test_pd_decode_uses_tuple_estimation(): - queue = PDDecodeQueue.__new__(PDDecodeQueue) - queue.args = SimpleNamespace(page_size=16) - queue.dp_index = 0 - queue.is_busy = lambda: False - queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) - req = SimpleNamespace( - request_id="req-0", - sample_params=SimpleNamespace(suggested_dp_index=0), - is_infer_decode=lambda: True, - get_tuple_tokens=lambda is_busy, ema_req_out_len: (16, 32), - ) - batch = Batch(batch_id=1, reqs=[req], dp_size_in_node=1) - - assert queue._caclu_batch_estimated_peak_token_num(batch) == 48 diff --git a/unit_tests/server/test_dp_request_capacity.py b/unit_tests/server/test_dp_request_capacity.py index f16001f8c7..aea1372403 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 f7c690bbe3..8a7e82fe83 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 7530c0c4d2..85153f43b3 100644 --- a/unit_tests/utils/test_sgl_utils.py +++ b/unit_tests/utils/test_sgl_utils.py @@ -107,7 +107,7 @@ 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, get_page_table_view=lambda att_batch_size, max_kv_len, microbatch_index: torch.zeros( diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py index f4f85d4b04..e3898aa801 100644 --- a/unit_tests/utils/test_speculative_utils.py +++ b/unit_tests/utils/test_speculative_utils.py @@ -125,7 +125,7 @@ def test_fa3_prefill_state_owns_causality(monkeypatch, state_class): 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=infer_state, ) From d33837231315b3e53f944fa175574a704020e651 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 13 Oct 2026 13:48:26 +0800 Subject: [PATCH 02/22] perf: prioritize token index management with paged KV support --- .../triton_kernel/copy_kv_index_to_req.py | 7 +- .../common/kv_cache_mem_manager/allocator.py | 45 +++++---- lightllm/common/req_manager/base.py | 91 +++++++++++-------- .../server/router/model_infer/infer_batch.py | 56 +++++++----- .../pd/decode_node_impl/decode_impl.py | 6 +- .../static_inference/static_benchmark.py | 2 +- .../mode_backend/test_multi_level_kv_cache.py | 18 ++-- 7 files changed, 126 insertions(+), 99 deletions(-) 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 9f4182d54f..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 @@ -85,14 +85,15 @@ def build_kv_indexes_and_input_ids(model_input: ModelInput, req_to_token_indexs, @triton.jit -def _update_req_token_indexes(Table, Packet, STRIDE: tl.constexpr, BLOCK: tl.constexpr, PAGE: tl.constexpr): +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) - offset = tl.load(Packet + 4 * row + 3) + # 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, pos < size, other=0) + pos % PAGE + 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) diff --git a/lightllm/common/kv_cache_mem_manager/allocator.py b/lightllm/common/kv_cache_mem_manager/allocator.py index f92305a4b9..079123237c 100644 --- a/lightllm/common/kv_cache_mem_manager/allocator.py +++ b/lightllm/common/kv_cache_mem_manager/allocator.py @@ -5,7 +5,7 @@ class KvCacheAllocator: - """A stack of physical page bases; returned token indexes own their storage.""" + """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 @@ -13,30 +13,35 @@ def __init__(self, size: int, page_size: int = 1) -> None: self.shared_can_use_token_num = SharedInt(f"mem_manger_can_use_token_num_{rank}") self.resize(size) - def _alloc_pages(self, need_size: int): - """Borrow page bases until the next allocator mutation; callers snapshot before DMA.""" - assert need_size % self.page_size == 0 + 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 = (self.size - self.can_use_mem_size) // self.page_size + 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_pages[start : start + need_size // self.page_size] + return self._free_token_index_blocks[start : start + size] def alloc(self, need_size: int) -> torch.Tensor: - pages = self._alloc_pages(need_size) - indexes = torch.empty(need_size, dtype=torch.int32, device="cpu", pin_memory=True) - self.expand_pages(pages, indexes.numpy()) - return indexes + 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: + out.numpy()[:] = index_blocks + else: + self.expand_token_index_blocks(index_blocks, out.numpy()) + return out - def expand_pages(self, pages, out=None): + def expand_token_index_blocks(self, index_blocks, out=None): if out is None: - out = np.empty(len(pages) * self.page_size, dtype=np.int32) + out = np.empty(len(index_blocks) * self.page_size, dtype=np.int32) if self.page_size == 1: - out[:] = pages + out[:] = index_blocks else: # Short pages expand faster along the page-base axis than in short token runs. np.add( - pages, + index_blocks, self.page_offsets[:, None], out=out.reshape(-1, self.page_size).T, order="C" if self.page_size < 8 else "F", @@ -47,25 +52,25 @@ 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.release_pages(values[:: self.page_size]) + self.free_token_index_blocks(values[:: self.page_size]) - def release_pages(self, pages): - size = len(pages) * 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_pages[end - len(pages) : end] = pages + 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) def free_all(self): - self.free_pages[:] = np.arange(0, self.size, self.page_size, dtype=np.int32) + 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) def resize(self, new_size: int) -> None: assert new_size % self.page_size == 0 self.size = new_size - self.free_pages = np.arange(0, new_size, self.page_size, dtype=np.int32) + 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/req_manager/base.py b/lightllm/common/req_manager/base.py index b6781b8f95..137cef90a9 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -54,7 +54,8 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana self.req_to_token_indexs = torch.zeros( (max_request_num + 1, max_sequence_length), dtype=torch.int32, device="cuda" ) - self._cpu_pages = np.empty((max_request_num + 1, max_sequence_length // page_size), dtype=np.int32) + # 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) @@ -78,23 +79,27 @@ def init_hold_request_indexs(self): # HOLD_REQUEST_ID 对应的请求行供 DP padding、overlap microbatch 等占位请求使用。将该行 # 按 page_size 划分后，每一页都映射到 mem_manager 额外保留的同一个物理页；这样占位请求 # 无论访问哪一个逻辑位置，都会落到合法且不会参与正常分配的 KV cache 地址上。 - self._cpu_pages[self.HOLD_REQUEST_ID].fill(self.mem_manager.HOLD_TOKEN_MEMINDEXES[0]) + 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]) ) - def get_cpu_page_bases(self, req_idx, start, end): - """Return a CPU view of page bases for the token range [start, end).""" + 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: + return self._cpu_token_index_blocks[req_idx, start:end] assert start % page_size == end % page_size == 0 - return self._cpu_pages[req_idx, start // page_size : end // page_size] + return self._cpu_token_index_blocks[req_idx, start // page_size : end // page_size] def get_cpu_token_indexes(self, req_idx, start, end): - """Materialize owned token indexes only for cache insertion and PD metadata.""" + """Return owned token indexes for cache insertion and asynchronous PD metadata.""" page_size = self.mem_manager.page_size + if page_size == 1: + return torch.from_numpy(self._cpu_token_index_blocks[req_idx, start:end].copy()) first, offset = divmod(start, page_size) - bases = self._cpu_pages[req_idx, first : (end + page_size - 1) // page_size] - indexes = self.mem_manager.allocator.expand_pages(bases) + 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]) def write_token_index_prefix(self, req_idx, start, indexes, req_prefix_lens): @@ -103,7 +108,9 @@ def write_token_index_prefix(self, req_idx, start, indexes, req_prefix_lens): assert start == 0 size = indexes.numel() assert size % page_size == 0 - self._cpu_pages[req_idx, start // page_size : (start + size) // page_size] = indexes.numpy()[::page_size] + self._cpu_token_index_blocks[req_idx, start // page_size : (start + size) // page_size] = indexes.numpy()[ + ::page_size + ] if size: req_prefix_lens[req_idx] = max(size, req_prefix_lens.get(req_idx, 0)) @@ -113,72 +120,76 @@ def copy_token_index_prefixes_to_gpu(self, req_prefix_lens): updates, offset = [], 0 for row, size in req_prefix_lens.items(): updates.append((row, 0, size, offset)) - offset += size // self.mem_manager.page_size + 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_page_bases(dest_req_idx, 0, size)[:] = self.get_cpu_page_bases(source_req_idx, 0, size) + 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_pages(self, allocations, radix_cache=None): - """Reserve pages once; CPU and GPU tables consume the same page-base assignment.""" + 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 - pages = self.mem_manager.allocator._alloc_pages(need) + 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_pages[req.req_idx, start // page_size] = pages[offset] + self._cpu_token_index_blocks[req.req_idx, start // page_size] = index_blocks[block_offset] else: - self._cpu_pages[req.req_idx, start // page_size : (start + size) // page_size] = pages[ - offset : offset + size // page_size - ] + 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 // page_size - self._copy_token_indexes_to_gpu(updates, pages) + 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_pages([(req, size)], radix_cache) + 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, pages=None): - """Snapshot compact pages, using direct DMA for small batches and GPU expansion otherwise.""" + 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 pages is None: + if index_blocks is None: parts = [ - self._cpu_pages[row, start // page_size : (start + size) // page_size] + self._cpu_token_index_blocks[row, start // page_size : (start + size) // page_size] for row, start, size, _ in updates ] - pages = parts[0] if len(parts) == 1 else np.concatenate(parts) - if not len(pages): + 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(pages) * (1 if packed else page_size), dtype=torch.int32, device="cpu", pin_memory=True + 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 - meta[:, 3] += header - data[header:] = pages + data[header:] = index_blocks + elif page_size == 1: + data[:] = index_blocks else: - self.mem_manager.allocator.expand_pages(pages, data) + self.mem_manager.allocator.expand_token_index_blocks(index_blocks, data) stream = self.wait_token_indexes() if packed: update_req_token_indexes( @@ -194,7 +205,7 @@ def _copy_token_indexes_to_gpu(self, updates, pages=None): 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 * page_size : offset * page_size + size], + packet if len(updates) == 1 else packet[offset : offset + size], non_blocking=True, ) self.token_indexes_ready.record(stream) @@ -203,14 +214,16 @@ def _copy_token_indexes_to_gpu(self, updates, pages=None): 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: + return max(0, target - req.hold_kv_len) size = max(0, (target + page_size - 1) // page_size * page_size - req.hold_kv_len) assert size % page_size == 0 return size def classify_reqs_and_alloc_kv(self, candidates, available, batch_max_tokens, radix_cache=None): - """按优先级顺序接纳请求，再一次性预留所有必需的页。 + """按优先级顺序接纳请求，再一次性分配所有必需的 token 容量。 - prefill 和 decode 两组请求都会跨调度轮次保留已分配的页。 + prefill 和 decode 两组请求都会跨调度轮次保留已分配的 token 容量。 decode 的可选额外预留仅使用满足所有必要分配后剩余的空闲容量。 """ prefill, decode, rejected, allocations, headroom = [], [], [], [], [] @@ -243,7 +256,7 @@ def classify_reqs_and_alloc_kv(self, candidates, available, batch_max_tokens, ra allocation[1] += extra extra_budget -= extra if allocations: - self.alloc_req_pages(allocations, radix_cache) + self.alloc_req_token_indexes(allocations, radix_cache) return prefill, decode, rejected, prefill_tokens def wait_token_indexes(self): @@ -257,13 +270,13 @@ def wait_token_indexes(self): def alloc(self): return self.req_list.alloc() - def free(self, free_req_indexes: List[int], free_page_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.allocator.release_pages(free_page_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) @@ -271,8 +284,8 @@ def free_req(self, free_req_index: int): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") return - def free_pages(self, free_page_index): - self.mem_manager.allocator.release_pages(free_page_index) + def free_token_index_blocks(self, free_token_index_blocks): + self.mem_manager.allocator.free_token_index_blocks(free_token_index_blocks) return def free_all(self): diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index afc2560a73..48cb9584d9 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -127,25 +127,27 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: return req_objs - def free_a_req_mem(self, free_page_index: List, req: "InferReq"): + def free_a_req_mem(self, free_token_index_blocks: List, req: "InferReq"): if self.radix_cache is None: - free_page_index.append(self.req_manager.get_cpu_page_bases(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_page_index=free_page_index, req=req) + self._free_req_mem_without_radix_insert(free_token_index_blocks=free_token_index_blocks, req=req) else: if not self.is_hybrid_att_model: - self._full_att_free_req(free_page_index=free_page_index, req=req) + self._full_att_free_req(free_token_index_blocks=free_token_index_blocks, req=req) else: - self._hybrid_att_free_req(free_page_index=free_page_index, req=req) + 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_page_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_page_index.append(self.req_manager.get_cpu_page_bases(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 小页状态。 @@ -163,27 +165,31 @@ def _free_req_mem_without_radix_insert(self, free_page_index: List, req: "InferR req.shared_kv_node = None return - def _full_att_free_req(self, free_page_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 - self._cache_req_prefix(free_page_index, req, cache_kv_len) + self._cache_req_prefix(free_token_index_blocks, req, cache_kv_len) - def _cache_req_prefix(self, free_page_index: List, req: "InferReq", cache_kv_len: int, **cache_metadata): + 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_page_index.append(self.req_manager.get_cpu_page_bases(req.req_idx, cache_kv_len, req.hold_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_page_index.append(self.req_manager.get_cpu_page_bases(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_page_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 @@ -211,7 +217,7 @@ def _hybrid_att_free_req(self, free_page_index: List, req: "InferReq"): elif shared_kv_len <= req.cur_kv_len: assert req.tail_small_page_buffer_id is None req.cur_kv_len = shared_kv_len - self._free_req_mem_without_radix_insert(free_page_index, req) + self._free_req_mem_without_radix_insert(free_token_index_blocks, req) return else: assert False, f"error state: cur_kv_len: {req.cur_kv_len}" @@ -222,7 +228,7 @@ def _hybrid_att_free_req(self, free_page_index: List, req: "InferReq"): state_idxs[-1] = req.tail_small_page_buffer_id req.tail_small_page_buffer_id = None self._cache_req_prefix( - free_page_index, + free_token_index_blocks, req, cache_kv_len, block_hashs=req.shm_req.hybrid_token_hash_list.get_all()[:page_num], @@ -252,7 +258,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_page_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: @@ -261,7 +267,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_page_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}") @@ -271,9 +277,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_page_index: - free_page_index = np.concatenate(free_page_index) - self.req_manager.free(free_req_index, free_page_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] @@ -297,12 +303,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_page_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_page_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 @@ -310,9 +316,9 @@ 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_page_index) != 0: - free_page_index = np.concatenate(free_page_index) - self.req_manager.free_pages(free_page_index) + if len(free_token_index_blocks) != 0: + 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): 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 1ae59f8698..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 @@ -102,8 +102,10 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: # 逻辑长度，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 if req_obj.hold_kv_len > old_prefix_len: - self.model.req_manager.free_pages( - self.model.req_manager.get_cpu_page_bases(req_obj.req_idx, old_prefix_len, req_obj.hold_kv_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 + ) ) req_obj.cur_kv_len = old_prefix_len req_obj.hold_kv_len = old_prefix_len diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 1cc7329b9d..58d7080edd 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -580,7 +580,7 @@ def _ensure_req_kv_capacity(self, req_idx: torch.Tensor, target_kv_lens: torch.T 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_pages(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 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 e16964c0db..a546f6fcbf 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 @@ -43,16 +43,16 @@ def test_non_gpu_cache_tiers_release_owned_tokens_without_radix_insert(): context = InferenceContext() context.is_hybrid_att_model = False context.req_manager = SimpleNamespace( - get_cpu_page_bases=lambda row, start, end: torch.arange(10 + start, 10 + end).numpy() + 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_page_indexes = [] + free_token_index_blocks = [] - context._free_req_mem_without_radix_insert(free_page_indexes, req) + context._free_req_mem_without_radix_insert(free_token_index_blocks, req) - assert free_page_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 @@ -62,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_page_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, @@ -139,7 +139,7 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): context = InferenceContext() context.is_hybrid_att_model = True context.req_manager = SimpleNamespace( - get_cpu_page_bases=lambda row, start, end: torch.arange(10 + start, 10 + end).numpy() + 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), @@ -153,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_page_indexes = [] + free_token_index_blocks = [] - context._free_req_mem_without_radix_insert(free_page_indexes, req) + context._free_req_mem_without_radix_insert(free_token_index_blocks, req) - assert free_page_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 From 80ca24c90dc8661f6217f539291c03f6363a0bb3 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 14 Oct 2026 07:00:25 +0000 Subject: [PATCH 03/22] refactor: mirror request token indexes on CPU and fuse input padding Maintain request token mappings in pinned CPU memory and upload pending ranges before inference. Fuse prefill and decode padding with Triton, reduce CPU input construction overhead, and mark synchronization sites for further analysis. Validation: 128 tests passed across padding kernels, model inputs and outputs, MTP batch alignment, CUDA Graph replay, and overlap paths. --- docs/kv_cache_page_size.md | 59 ++ lightllm/common/basemodel/attention/fa3/fp.py | 1 + .../common/basemodel/attention/fa3/mla.py | 2 +- lightllm/common/basemodel/basemodel.py | 277 +++++--- 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 +++++++++ .../triton_kernel/pad_decode_inputs.py | 127 ++++ .../triton_kernel/pad_prefill_inputs.py | 113 +++ lightllm/common/cuda_wrapper.py | 42 ++ lightllm/common/infer_utils.py | 15 + .../common/kv_cache_mem_manager/allocator.py | 145 ++-- .../deepseek2_mem_manager.py | 4 +- .../kv_cache_mem_manager/mem_manager.py | 36 +- .../ppl_int4kv_mem_manager.py | 4 + .../ppl_int8kv_mem_manager.py | 4 + .../qwen3next_mem_manager.py | 9 + lightllm/common/req_manager/base.py | 259 ++----- .../common/req_manager/req_sampling_params.py | 14 +- 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 | 360 ++++++---- .../router/model_infer/infer_req_ext.py | 2 +- .../model_infer/mode_backend/base_backend.py | 188 +++-- .../mode_backend/diverse_backend/impl.py | 38 +- .../dp_backend/dp_shared_kv_trans.py | 19 +- .../mode_backend/generic_pre_process.py | 58 +- .../mode_backend/multi_level_kv_cache.py | 24 +- .../pd/decode_node_impl/decode_impl.py | 23 +- .../pd/prefill_node_impl/prefill_impl.py | 9 +- .../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 | 23 + skills/test_model/SKILL.md | 40 -- ...t_kv_cache_calib_per_head_qwen2.5_14b.json | 499 +------------- .../static_inference/static_benchmark.py | 24 +- .../basemodel/attention/test_fp8_page_size.py | 45 ++ .../common/basemodel/test_model_input.py | 6 + .../common/basemodel/test_model_output.py | 10 +- .../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_pad_decode_inputs.py | 92 +++ .../triton_kernel/test_pad_prefill_inputs.py | 86 +++ .../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_cpu.py | 645 ++++++++++++++++++ unit_tests/common/test_req_manager_page.py | 514 ++++++++++++++ .../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 | 48 ++ .../mode_backend/test_diverse_backend.py | 78 +++ .../mode_backend/test_generic_pre_process.py | 31 + .../mode_backend/test_multi_level_kv_cache.py | 102 ++- .../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 + .../req_queue/test_pd_queue_selection.py | 109 +++ 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 | 158 +++++ .../utils/test_config_utils_fixed_kv.py | 33 + unit_tests/utils/test_envs_utils.py | 12 +- unit_tests/utils/test_sgl_utils.py | 2 +- unit_tests/utils/test_speculative_utils.py | 2 +- 82 files changed, 4503 insertions(+), 1361 deletions(-) create mode 100644 docs/kv_cache_page_size.md create mode 100644 lightllm/common/basemodel/triton_kernel/pad_decode_inputs.py create mode 100644 lightllm/common/basemodel/triton_kernel/pad_prefill_inputs.py create mode 100644 lightllm/common/infer_utils.py create mode 100644 unit_tests/common/basemodel/attention/test_fp8_page_size.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_pad_decode_inputs.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_pad_prefill_inputs.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py create mode 100644 unit_tests/common/test_paged_mem_manager.py create mode 100644 unit_tests/common/test_pinned_memory_non_blocking_copy.py create mode 100644 unit_tests/common/test_quantized_paged_mem_manager.py create mode 100644 unit_tests/common/test_req_manager_cpu.py create mode 100644 unit_tests/common/test_req_manager_page.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_diverse_backend.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py create mode 100644 unit_tests/server/router/req_queue/test_pd_queue_selection.py create mode 100644 unit_tests/server/test_pd_start_args.py create mode 100644 unit_tests/test_static_benchmark_paged_kv.py create 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 new file mode 100644 index 0000000000..8d45aeb1cf --- /dev/null +++ b/docs/kv_cache_page_size.md @@ -0,0 +1,59 @@ +# 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 接口。 + +`ReqManager.req_to_token_indexs_cpu` 是与 GPU 请求表同形状、同 dtype 的 CPU 镜像。 +调度分配、前缀匹配和 diverse 共享直接拷贝到镜像；缓存插入、释放和 PD 元数据直接读取镜像。 +缓存持有的索引必须复制，避免请求槽位复用后改变缓存内容。 +每个 `InferReq` 用 `move_to_gpu_index_len` 记录已经提交到推理流的 GPU 索引长度。 +`_get_classed_reqs` 返回前通过 backend 的 `_copy_req_token_indexes_to_gpu` 统一收集 +prefill/decode 请求的 +`[move_to_gpu_index_len, hold_kv_len)` 区间，通过一次原生 `cudaMemcpyBatchAsync` 提交， +成功后将长度更新到 `hold_kv_len`。`DuringApiCall` 源访问属性保证返回前读取完 CPU 数据， +后续计算依靠同流顺序读取 GPU 表；长度表示异步拷贝已排队，不表示设备已经完成拷贝。 +暂停/释放时重置长度，前缀重映射时回退长度以重新上传。 +DP 通过共享 CPU 镜像读取其他 rank 的分配元数据，无需等待下一次推理。 +请求筛选、地址计算、批量上传及长度更新均在这个 backend 函数内完成；`ReqManager` 只持有映射表。 +独立 benchmark 和模型初始化检查共用 `cuda_memcpy_batch_async` 底层调用。 + +## 生命周期 + +- 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 5ee71e629e..9c9b811954 100644 --- a/lightllm/common/basemodel/attention/fa3/fp.py +++ b/lightllm/common/basemodel/attention/fa3/fp.py @@ -2,6 +2,7 @@ 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 7cf6a5d524..e124562f02 100644 --- a/lightllm/common/basemodel/attention/fa3/mla.py +++ b/lightllm/common/basemodel/attention/fa3/mla.py @@ -2,7 +2,7 @@ import torch import triton from ..base_att import BasePrefillAttState, BaseDecodeAttState, AttControl -from typing import Tuple +from typing import Optional, TYPE_CHECKING, 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 06a053d200..b539be887c 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -6,7 +6,6 @@ import math import json import torch -import torch.nn.functional as F import triton from typing import final, List from tqdm import tqdm @@ -17,11 +16,17 @@ 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 build_kv_indexes_and_input_ids +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.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.common.basemodel.triton_kernel.pad_decode_inputs import pad_decode_inputs +from lightllm.common.basemodel.triton_kernel.pad_prefill_inputs import pad_prefill_inputs 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 @@ -325,17 +330,24 @@ def forward(self, model_input: ModelInput): else: return self._decode(model_input) - 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, + 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): infer_state = self.infer_state_class() infer_state.hidden_collector = self.hidden_collector_prototype.new_instance() infer_state.input_ids = model_input.input_ids @@ -366,7 +378,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 = model_input.mem_indexes + infer_state.mem_index = self._select_mem_indexes(model_input) infer_state.microbatch_index = microbatch_index infer_state.dist_group = dist_group_manager.get_group(microbatch_index) @@ -388,29 +400,6 @@ 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 - ): - model_input = self._create_padded_decode_model_input(model_input, batch_size) - infer_state = self._create_inferstate(model_input, microbatch_index) - # Attention backends need the capture flag before initializing their state. - infer_state.is_cuda_graph = 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 @@ -422,28 +411,24 @@ 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) - 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 - ) - new_model_input.b_mtp_index = F.pad( - new_model_input.b_mtp_index, (0, padded_batch_size), mode="constant", value=0 - ) - new_model_input.b_seq_len = F.pad(new_model_input.b_seq_len, (0, padded_batch_size), mode="constant", value=2) - if new_model_input.b_position_delta is not None: - new_model_input.b_position_delta = F.pad( - new_model_input.b_position_delta, (0, padded_batch_size), mode="constant", value=0 - ) - new_model_input.multimodal_params = new_model_input.multimodal_params + [ - {"images": [], "audios": []} for _ in range(padded_batch_size) - ] - - new_model_input.b_shared_seq_len = F.pad( - new_model_input.b_shared_seq_len, (0, padded_batch_size), mode="constant", value=0 - ) - new_model_input.b_shared_radix_node_id = F.pad( - new_model_input.b_shared_radix_node_id, (0, padded_batch_size), mode="constant", value=-1 + ( + new_model_input.input_ids, + new_model_input.b_req_idx, + new_model_input.b_mtp_index, + new_model_input.b_seq_len, + new_model_input.b_position_delta, + new_model_input.b_shared_seq_len, + new_model_input.b_shared_radix_node_id, + ) = pad_decode_inputs( + input_ids=model_input.input_ids, + b_req_idx=model_input.b_req_idx, + b_mtp_index=model_input.b_mtp_index, + b_seq_len=model_input.b_seq_len, + b_position_delta=model_input.b_position_delta, + b_shared_seq_len=model_input.b_shared_seq_len, + b_shared_radix_node_id=model_input.b_shared_radix_node_id, + new_batch_size=new_batch_size, + hold_request_id=self.req_manager.HOLD_REQUEST_ID, ) # 特殊模型，特殊模式的特殊变量的特殊 padding @@ -453,6 +438,10 @@ def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_s new_batch_size=new_batch_size, ) + new_model_input.multimodal_params = new_model_input.multimodal_params + [ + {"images": [], "audios": []} for _ in range(padded_batch_size) + ] + new_model_input.check_input() return new_model_input @@ -471,20 +460,31 @@ def _create_padded_prefill_model_input(self, model_input: ModelInput, new_handle new_model_input.max_q_seq_len = max(padded_token_num, model_input.max_q_seq_len) new_model_input.max_kv_seq_len = max(padded_token_num, model_input.max_kv_seq_len) new_model_input.max_cache_len = max(0, model_input.max_cache_len) - new_model_input.input_ids = F.pad(new_model_input.input_ids, (0, padded_token_num), mode="constant", value=1) - new_model_input.b_req_idx = F.pad( - new_model_input.b_req_idx, (0, 1), mode="constant", value=self.req_manager.HOLD_REQUEST_ID - ) - new_model_input.b_mtp_index = F.pad(new_model_input.b_mtp_index, (0, 1), mode="constant", value=0) - 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) - new_model_input.b_prefill_start_loc = F.pad( - model_input.b_prefill_start_loc.to(torch.int32), (0, 1), value=handle_token_num + ( + new_model_input.input_ids, + new_model_input.b_req_idx, + new_model_input.b_mtp_index, + new_model_input.b_seq_len, + new_model_input.b_ready_cache_len, + new_model_input.b_is_decode_req, + ) = pad_prefill_inputs( + input_ids=model_input.input_ids, + b_req_idx=model_input.b_req_idx, + b_mtp_index=model_input.b_mtp_index, + b_seq_len=model_input.b_seq_len, + b_ready_cache_len=model_input.b_ready_cache_len, + b_is_decode_req=model_input.b_is_decode_req, + new_handle_token_num=new_handle_token_num, + hold_request_id=self.req_manager.HOLD_REQUEST_ID, ) + 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 # 构建新的list, 使用 append 可能会让外面使用的数组引用发生变化，导致错误。 - 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": []}] + 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": []} + ] # 特殊模型，特殊模式的特殊变量的特殊 padding if new_model_input.mtp_draft_input_hiddens is not None: @@ -541,6 +541,16 @@ 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 @@ -555,7 +565,16 @@ def _prefill( handle_token_num=infer_handle_token_num ) - infer_state = self._create_prefill_inferstate(model_input, 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() model_output = self._context_forward(infer_state=infer_state) model_output = self._create_unpad_prefill_model_output( @@ -563,12 +582,29 @@ 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)) @@ -585,7 +621,13 @@ 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) - infer_state = self._create_decode_inferstate(model_input, infer_batch_size, need_capture) + 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 会根据该标记准备 CUDA Graph capture 专用状态， + # 因此必须在 init_att_state 之前完成赋值。 + infer_state.is_cuda_graph = need_capture + infer_state.init_some_extra_state(self) + infer_state.init_att_state() if use_cuda_graph: if need_capture: @@ -696,6 +738,15 @@ 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): @@ -713,8 +764,24 @@ 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 - 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_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() + model_output0, model_output1 = self._overlap_tpsp_context_forward(infer_state0, infer_state1=infer_state1) model_output0 = self._create_unpad_prefill_model_output( @@ -730,6 +797,8 @@ 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() @@ -738,6 +807,19 @@ 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): @@ -748,16 +830,21 @@ 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)) - 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: + if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch): 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 = 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 = 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) - infer_state1 = self._create_decode_inferstate(model_input1, infer_batch_size, need_capture, 1) - - if use_cuda_graph: if need_capture: model_output0, model_output1 = self.graph.capture_decode( self._overlap_tpsp_token_forward, @@ -770,11 +857,23 @@ 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 @@ -898,8 +997,10 @@ def _check_max_len_infer(self): b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda") page_size = self.mem_manager.page_size alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size - mem_indexes = self.mem_manager.alloc(alloc_token_num).cuda() - self.req_manager.req_to_token_indexs[b_req_idx[0], : len(mem_indexes)] = mem_indexes + mem_indexes = self.mem_manager.alloc(alloc_token_num) + req_idx = int(b_req_idx[0]) + self.req_manager.req_to_token_indexs_cpu[req_idx, : len(mem_indexes)].copy_(mem_indexes) + self.req_manager.req_to_token_indexs.copy_(self.req_manager.req_to_token_indexs_cpu, non_blocking=True) b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = self.batch_max_tokens b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda") @@ -979,8 +1080,10 @@ def _autotune_warmup(self): b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda") page_size = self.mem_manager.page_size alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size - mem_indexes = self.mem_manager.alloc(alloc_token_num).cuda() - self.req_manager.req_to_token_indexs[b_req_idx[0], : len(mem_indexes)] = mem_indexes + mem_indexes = self.mem_manager.alloc(alloc_token_num) + req_idx = int(b_req_idx[0]) + self.req_manager.req_to_token_indexs_cpu[req_idx, : len(mem_indexes)].copy_(mem_indexes) + self.req_manager.req_to_token_indexs.copy_(self.req_manager.req_to_token_indexs_cpu, non_blocking=True) b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = input_len b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda") diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index 447916fafc..d8f2329b24 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -32,8 +32,6 @@ 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 @@ -191,6 +189,9 @@ 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 8e28ea03b0..ee888c1a00 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -1,3 +1,4 @@ +import os import torch import torch.distributed as dist import copy @@ -7,10 +8,13 @@ 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 +from lightllm.utils.torch_memory_saver_utils import ( + TorchMemorySaverWrapper, + MemoryTag, +) from .infer_struct import InferStateInfo @@ -282,12 +286,6 @@ 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 @@ -346,11 +344,6 @@ 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 1eb7e3000b..28d44fefbf 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -1,4 +1,6 @@ +import os import torch +import copy import bisect import triton from typing import List, Tuple @@ -6,6 +8,7 @@ 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 f992e25b02..efe563d764 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,103 +1,204 @@ import torch + import triton import triton.language as tl -from lightllm.common.basemodel.batch_objs import ModelInput @triton.jit -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, +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, BLOCK: tl.constexpr, ): - 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: - 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_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 ) - 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]( + 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 + 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]( req_to_token_indexs, b_req_idx, - 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, + 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, ) - return out, input_ids @triton.jit -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 +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, ) + 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 6dd8890a2e..ad995b2b5d 100644 --- a/lightllm/common/basemodel/triton_kernel/dynamic_mtp_utils.py +++ b/lightllm/common/basemodel/triton_kernel/dynamic_mtp_utils.py @@ -209,40 +209,86 @@ 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) - 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_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 ) - 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, + 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, 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=triton.next_power_of_2(old_batch_size), + BLOCK_SIZE=BLOCK_SIZE, num_warps=8, num_stages=1, ) - for name in fields: - setattr(model_input, name, tensors["out_" + name]) + + 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 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 b9081e7776..23d8b0f818 100644 --- a/lightllm/common/basemodel/triton_kernel/fa3_utils.py +++ b/lightllm/common/basemodel/triton_kernel/fa3_utils.py @@ -275,3 +275,36 @@ 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 c1794ec00e..a5dfc6d226 100644 --- a/lightllm/common/basemodel/triton_kernel/gather_token_id.py +++ b/lightllm/common/basemodel/triton_kernel/gather_token_id.py @@ -86,3 +86,325 @@ 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/basemodel/triton_kernel/pad_decode_inputs.py b/lightllm/common/basemodel/triton_kernel/pad_decode_inputs.py new file mode 100644 index 0000000000..69ddef61f8 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/pad_decode_inputs.py @@ -0,0 +1,127 @@ +from typing import Optional + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _pad_decode_inputs_kernel( + input_ids, + b_req_idx, + b_mtp_index, + b_seq_len, + b_position_delta, + b_shared_seq_len, + b_shared_radix_node_id, + padded_input_ids, + padded_b_req_idx, + padded_b_mtp_index, + padded_b_seq_len, + padded_b_position_delta, + padded_b_shared_seq_len, + padded_b_shared_radix_node_id, + batch_size, + new_batch_size, + hold_request_id, + HAS_B_POSITION_DELTA: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + source_mask = offsets < batch_size + output_mask = offsets < new_batch_size + + # 保留已有请求，新增 HOLD 请求使用 token id 1、seq_len 2。 + token_ids = tl.load(input_ids + offsets, mask=source_mask, other=1) + tl.store(padded_input_ids + offsets, token_ids, mask=output_mask) + + req_idx = tl.load(b_req_idx + offsets, mask=source_mask, other=hold_request_id) + tl.store(padded_b_req_idx + offsets, req_idx, mask=output_mask) + + mtp_index = tl.load(b_mtp_index + offsets, mask=source_mask, other=0) + tl.store(padded_b_mtp_index + offsets, mtp_index, mask=output_mask) + + seq_len = tl.load(b_seq_len + offsets, mask=source_mask, other=2) + tl.store(padded_b_seq_len + offsets, seq_len, mask=output_mask) + + # 位置偏移按需补齐；HOLD 请求不携带位置偏移和共享缓存。 + if HAS_B_POSITION_DELTA: + position_delta = tl.load(b_position_delta + offsets, mask=source_mask, other=0) + tl.store(padded_b_position_delta + offsets, position_delta, mask=output_mask) + + shared_seq_len = tl.load(b_shared_seq_len + offsets, mask=source_mask, other=0) + tl.store(padded_b_shared_seq_len + offsets, shared_seq_len, mask=output_mask) + + shared_radix_node_id = tl.load(b_shared_radix_node_id + offsets, mask=source_mask, other=-1) + tl.store(padded_b_shared_radix_node_id + offsets, shared_radix_node_id, mask=output_mask) + + +@torch.no_grad() +def pad_decode_inputs( + input_ids: torch.Tensor, + b_req_idx: torch.Tensor, + b_mtp_index: torch.Tensor, + b_seq_len: torch.Tensor, + b_position_delta: Optional[torch.Tensor], + b_shared_seq_len: torch.Tensor, + b_shared_radix_node_id: torch.Tensor, + new_batch_size: int, + hold_request_id: int, +): + """一次 kernel 调用将 decode Token 和请求元数据补齐到 new_batch_size。 + + 调用方保证输入为同一 CUDA 设备上的连续一维 Tensor，各字段按请求对齐， + 且 new_batch_size 大于现有请求数。b_position_delta 为 None 时保持 None。 + 返回顺序与七个 Tensor 输入一致。 + """ + batch_size = b_req_idx.shape[0] + + padded_input_ids = torch.empty((new_batch_size,), dtype=input_ids.dtype, device=input_ids.device) + padded_b_req_idx = torch.empty((new_batch_size,), dtype=b_req_idx.dtype, device=b_req_idx.device) + padded_b_mtp_index = torch.empty((new_batch_size,), dtype=b_mtp_index.dtype, device=b_mtp_index.device) + padded_b_seq_len = torch.empty((new_batch_size,), dtype=b_seq_len.dtype, device=b_seq_len.device) + padded_b_position_delta = None + if b_position_delta is not None: + padded_b_position_delta = torch.empty( + (new_batch_size,), dtype=b_position_delta.dtype, device=b_position_delta.device + ) + padded_b_shared_seq_len = torch.empty( + (new_batch_size,), dtype=b_shared_seq_len.dtype, device=b_shared_seq_len.device + ) + padded_b_shared_radix_node_id = torch.empty( + (new_batch_size,), dtype=b_shared_radix_node_id.dtype, device=b_shared_radix_node_id.device + ) + + block_size = 256 + grid = (triton.cdiv(new_batch_size, block_size),) + _pad_decode_inputs_kernel[grid]( + input_ids=input_ids, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + b_seq_len=b_seq_len, + b_position_delta=b_position_delta, + b_shared_seq_len=b_shared_seq_len, + b_shared_radix_node_id=b_shared_radix_node_id, + padded_input_ids=padded_input_ids, + padded_b_req_idx=padded_b_req_idx, + padded_b_mtp_index=padded_b_mtp_index, + padded_b_seq_len=padded_b_seq_len, + padded_b_position_delta=padded_b_position_delta, + padded_b_shared_seq_len=padded_b_shared_seq_len, + padded_b_shared_radix_node_id=padded_b_shared_radix_node_id, + batch_size=batch_size, + new_batch_size=new_batch_size, + hold_request_id=hold_request_id, + HAS_B_POSITION_DELTA=b_position_delta is not None, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return ( + padded_input_ids, + padded_b_req_idx, + padded_b_mtp_index, + padded_b_seq_len, + padded_b_position_delta, + padded_b_shared_seq_len, + padded_b_shared_radix_node_id, + ) diff --git a/lightllm/common/basemodel/triton_kernel/pad_prefill_inputs.py b/lightllm/common/basemodel/triton_kernel/pad_prefill_inputs.py new file mode 100644 index 0000000000..c2828e251d --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/pad_prefill_inputs.py @@ -0,0 +1,113 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _pad_prefill_inputs_kernel( + input_ids, + b_req_idx, + b_mtp_index, + b_seq_len, + b_ready_cache_len, + b_is_decode_req, + padded_input_ids, + padded_b_req_idx, + padded_b_mtp_index, + padded_b_seq_len, + padded_b_ready_cache_len, + padded_b_is_decode_req, + handle_token_num, + padded_token_num, + batch_size, + hold_request_id, + BLOCK_SIZE: tl.constexpr, +): + block_start = tl.program_id(0) * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + # Token 按实际补齐数量扩展，新增位置填入 token id 1。 + token_ids = tl.load(input_ids + offsets, mask=offsets < handle_token_num, other=1) + tl.store(padded_input_ids + offsets, token_ids, mask=offsets < handle_token_num + padded_token_num) + + # 请求元数据只追加一个 HOLD 请求；后续 Token block 不再处理元数据。 + if block_start < batch_size + 1: + source_mask = offsets < batch_size + output_mask = offsets < batch_size + 1 + + req_idx = tl.load(b_req_idx + offsets, mask=source_mask, other=hold_request_id) + tl.store(padded_b_req_idx + offsets, req_idx, mask=output_mask) + + mtp_index = tl.load(b_mtp_index + offsets, mask=source_mask, other=0) + tl.store(padded_b_mtp_index + offsets, mtp_index, mask=output_mask) + + seq_len = tl.load(b_seq_len + offsets, mask=source_mask, other=padded_token_num) + tl.store(padded_b_seq_len + offsets, seq_len, mask=output_mask) + + ready_cache_len = tl.load(b_ready_cache_len + offsets, mask=source_mask, other=0) + tl.store(padded_b_ready_cache_len + offsets, ready_cache_len, mask=output_mask) + + is_decode_req = tl.load(b_is_decode_req + offsets, mask=source_mask, other=False) + tl.store(padded_b_is_decode_req + offsets, is_decode_req, mask=output_mask) + + +@torch.no_grad() +def pad_prefill_inputs( + input_ids: torch.Tensor, + b_req_idx: torch.Tensor, + b_mtp_index: torch.Tensor, + b_seq_len: torch.Tensor, + b_ready_cache_len: torch.Tensor, + b_is_decode_req: torch.Tensor, + new_handle_token_num: int, + hold_request_id: int, +): + """一次 kernel 调用补齐 prefill Token，并追加一个 HOLD 请求。 + + 调用方保证输入为同一 CUDA 设备上的连续一维 Tensor，元数据按请求对齐， + 且 new_handle_token_num 大于现有 Token 数。返回顺序与六个 Tensor 输入一致。 + """ + handle_token_num = input_ids.shape[0] + padded_token_num = new_handle_token_num - handle_token_num + batch_size = b_req_idx.shape[0] + new_batch_size = batch_size + 1 + + padded_input_ids = torch.empty((new_handle_token_num,), dtype=input_ids.dtype, device=input_ids.device) + padded_b_req_idx = torch.empty((new_batch_size,), dtype=b_req_idx.dtype, device=b_req_idx.device) + padded_b_mtp_index = torch.empty((new_batch_size,), dtype=b_mtp_index.dtype, device=b_mtp_index.device) + padded_b_seq_len = torch.empty((new_batch_size,), dtype=b_seq_len.dtype, device=b_seq_len.device) + padded_b_ready_cache_len = torch.empty( + (new_batch_size,), dtype=b_ready_cache_len.dtype, device=b_ready_cache_len.device + ) + padded_b_is_decode_req = torch.empty((new_batch_size,), dtype=b_is_decode_req.dtype, device=b_is_decode_req.device) + + block_size = 256 + grid = (triton.cdiv(max(new_handle_token_num, new_batch_size), block_size),) + _pad_prefill_inputs_kernel[grid]( + input_ids=input_ids, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + b_seq_len=b_seq_len, + b_ready_cache_len=b_ready_cache_len, + b_is_decode_req=b_is_decode_req, + padded_input_ids=padded_input_ids, + padded_b_req_idx=padded_b_req_idx, + padded_b_mtp_index=padded_b_mtp_index, + padded_b_seq_len=padded_b_seq_len, + padded_b_ready_cache_len=padded_b_ready_cache_len, + padded_b_is_decode_req=padded_b_is_decode_req, + handle_token_num=handle_token_num, + padded_token_num=padded_token_num, + batch_size=batch_size, + hold_request_id=hold_request_id, + BLOCK_SIZE=block_size, + num_warps=4, + ) + return ( + padded_input_ids, + padded_b_req_idx, + padded_b_mtp_index, + padded_b_seq_len, + padded_b_ready_cache_len, + padded_b_is_decode_req, + ) diff --git a/lightllm/common/cuda_wrapper.py b/lightllm/common/cuda_wrapper.py index 32823f1766..b089227004 100644 --- a/lightllm/common/cuda_wrapper.py +++ b/lightllm/common/cuda_wrapper.py @@ -38,6 +38,48 @@ cudaError_t = ctypes.c_int cudaMemcpyKind = ctypes.c_int +_cuda_memcpy_batch_available = True + + +def cuda_memcpy_pinned_h2d_batch_async(dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor]): + """在当前非默认流上，将 pinned CPU Tensor 批量异步拷贝到 GPU Tensor，仅支持 H2D。 + + 使用 Stream 模式；调用方保证源为 pinned 内存、目标在当前 CUDA 设备上， + 各对 Tensor 连续且形状和 dtype 一致，并保持源数据有效且不变直到传输完成。 + """ + global _cuda_memcpy_batch_available + + if not dst_gpu_tensors: + return + stream = torch.cuda.current_stream() + # CUDA 批量拷贝接口明确禁止使用 legacy 默认流（句柄为 0）。 + # 在异常回退前检查，避免默认流调用错误被吞掉并关闭全局 batch 开关。 + assert stream.cuda_stream != 0, "cuda_memcpy_pinned_h2d_batch_async 必须在非默认 CUDA 流上调用" + + if _cuda_memcpy_batch_available: + try: + from cuda.bindings import runtime as cuda_runtime + + attrs = cuda_runtime.cudaMemcpyAttributes() + # 使用 Stream 模式，让源数据读取遵循当前流顺序。 + # 调用方需保持源数据有效且不变，直到传输完成。 + attrs.srcAccessOrder = cuda_runtime.cudaMemcpySrcAccessOrder.cudaMemcpySrcAccessOrderStream + dst_ptrs = [tensor.data_ptr() for tensor in dst_gpu_tensors] + src_ptrs = [tensor.data_ptr() for tensor in src_pinned_tensors] + byte_sizes = [tensor.numel() * tensor.element_size() for tensor in src_pinned_tensors] + (error,) = cuda_runtime.cudaMemcpyBatchAsync( + dst_ptrs, src_ptrs, byte_sizes, len(byte_sizes), [attrs], [0], 1, stream.cuda_stream + ) + if error != cuda_runtime.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaMemcpyBatchAsync failed: {error.name}") + except Exception as error: + _cuda_memcpy_batch_available = False + logger.warning(f"Disabling CUDA batch copies; falling back to individual copies: {error}") + else: + return + + for dst, src in zip(dst_gpu_tensors, src_pinned_tensors): + dst.copy_(src, non_blocking=True) class cudaIpcMemHandle_t(ctypes.Structure): diff --git a/lightllm/common/infer_utils.py b/lightllm/common/infer_utils.py new file mode 100644 index 0000000000..e1b9cc3830 --- /dev/null +++ b/lightllm/common/infer_utils.py @@ -0,0 +1,15 @@ +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 079123237c..1331943b1a 100644 --- a/lightllm/common/kv_cache_mem_manager/allocator.py +++ b/lightllm/common/kv_cache_mem_manager/allocator.py @@ -1,76 +1,103 @@ -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: - """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 + 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 + 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] - 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: - out.numpy()[:] = index_blocks - else: - self.expand_token_index_blocks(index_blocks, out.numpy()) - return out - - def expand_token_index_blocks(self, index_blocks, out=None): - if out is None: - out = np.empty(len(index_blocks) * self.page_size, dtype=np.int32) - if self.page_size == 1: - out[:] = index_blocks + # 利用缓冲区返回，避免异步情况下的内存竞争 + 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 else: - # 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 + # 从 gpu 到 cpu 的拷贝操作是流内阻塞操作 + self.mem_state[start:end] = free_index + + self.mark_start -= len(free_index) + + self.can_use_mem_size += len(free_index) self.shared_can_use_token_num.set_value(self.can_use_mem_size) + if self.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._free_token_index_blocks[:] = np.arange(0, self.size, self.page_size, dtype=np.int32) - self.can_use_mem_size = self.size + 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.shared_can_use_token_num.set_value(self.can_use_mem_size) + return def resize(self, new_size: int) -> None: - assert new_size % self.page_size == 0 + """ + just for test code + """ self.size = new_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) + 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) 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 96ed9bc668..4be820b753 100644 --- a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py @@ -1,6 +1,8 @@ import torch +import os +import torch.distributed as dist from .mem_manager import MemoryManager -from typing import List, Any +from typing import List, Union, 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 8dfd7f1721..6c812ff21f 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -1,4 +1,7 @@ import math +import os +import re +import weakref import torch import torch.distributed as dist import torch.multiprocessing as mp @@ -26,7 +29,6 @@ class MemoryManager: - operator_class = NormalMemOperator def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): @@ -45,7 +47,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.page_size) + self.allocator = KvCacheAllocator(self.size) self._init_buffers( self.size, @@ -192,6 +194,9 @@ 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 @@ -204,6 +209,23 @@ 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]} @@ -226,6 +248,16 @@ def write_to_shm(self, req_manager): # 这个地方是一个不太优雅的设计，但是暂时这么做，可以让dp shared kv swap模块直接访问 req_manager 中的 req_to_token_indexs # 避免过多无用的数据复制和传输开销。 self.req_to_token_indexs: torch.Tensor = req_manager.req_to_token_indexs + # DP transfers need allocation metadata even before the next model forward publishes it. + self.req_to_token_indexs_cpu: torch.Tensor = req_manager.req_to_token_indexs_cpu.share_memory_() + # share_memory_ 会更换存储并丢失 pinned 属性，注册共享存储以继续直接异步上传。 + if not self.req_to_token_indexs_cpu.is_pinned(): + runtime = torch.cuda.cudart() + cpu_table = self.req_to_token_indexs_cpu + error = runtime.cudaHostRegister(cpu_table.data_ptr(), cpu_table.numel() * cpu_table.element_size(), 0) + if error != runtime.cudaError.success: + raise RuntimeError(f"cudaHostRegister failed for shared request indexes: {error}") + weakref.finalize(cpu_table, runtime.cudaHostUnregister, cpu_table.data_ptr()) lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") with lock: 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 39c08eb613..6b02cbed63 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,6 +38,10 @@ 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 6e64051081..0f5e8e9370 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,6 +37,10 @@ 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 713c6d0890..1962fa71b0 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -61,6 +61,15 @@ 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 137cef90a9..b21a4f0e4e 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -1,9 +1,6 @@ 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 @@ -13,31 +10,45 @@ logger = init_logger("lightllm.common.req_manager") -class _ReqIndexPool: - def __init__(self, max_request_num): - self.free_indexes = list(range(max_request_num - 1, -1, -1)) - self.marks = [0] * max_request_num +class _ReqNode: + def __init__(self, index): + self.index = index + self.next: "_ReqNode" = None - @property - def can_alloc_size(self): - return len(self.free_indexes) + +class _ReqLinkedList: + 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 def alloc(self): - if not self.free_indexes: + if self.root_node.next is None: logger.warning("alloc req index fail") return None - index = self.free_indexes.pop() - assert self.marks[index] == 0 - self.marks[index] = 1 - return index + 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 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.free_indexes.append(index) + self.can_alloc_size += 1 + return def is_all_free(self): - return len(self.free_indexes) == len(self.marks) + return self.can_alloc_size == len(self.marks) class ReqManager: @@ -48,16 +59,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 = _ReqIndexPool(max_request_num) + self.req_list = _ReqLinkedList(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 + # Allocation and cache metadata use this CPU mirror; inference reads the GPU table. + self.req_to_token_indexs_cpu = torch.zeros( + (max_request_num + 1, max_sequence_length), dtype=torch.int32, device="cpu", pin_memory=True + ) self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num) self.max_request_num = max_request_num self.HOLD_REQUEST_ID = max_request_num @@ -79,204 +90,26 @@ def init_hold_request_indexs(self): # HOLD_REQUEST_ID 对应的请求行供 DP padding、overlap microbatch 等占位请求使用。将该行 # 按 page_size 划分后，每一页都映射到 mem_manager 额外保留的同一个物理页；这样占位请求 # 无论访问哪一个逻辑位置，都会落到合法且不会参与正常分配的 KV cache 地址上。 - 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]) - ) - - 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: - return self._cpu_token_index_blocks[req_idx, start:end] - assert start % page_size == end % page_size == 0 - return self._cpu_token_index_blocks[req_idx, start // page_size : end // page_size] - - 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: - return torch.from_numpy(self._cpu_token_index_blocks[req_idx, start:end].copy()) - 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]) - - 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, start // page_size : (start + 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 + hold_row = self.req_to_token_indexs_cpu[self.HOLD_REQUEST_ID] + hold_page = torch.tensor( + self.mem_manager.HOLD_TOKEN_MEMINDEXES, + dtype=hold_row.dtype, + device=hold_row.device, ) - 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: - data[:] = index_blocks - else: - self.mem_manager.allocator.expand_token_index_blocks(index_blocks, data) - 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: - return max(0, target - req.hold_kv_len) - size = max(0, (target + page_size - 1) // page_size * page_size - req.hold_kv_len) - assert size % page_size == 0 - return size - - def classify_reqs_and_alloc_kv(self, candidates, available, batch_max_tokens, radix_cache=None): - """按优先级顺序接纳请求，再一次性分配所有必需的 token 容量。 - - prefill 和 decode 两组请求都会跨调度轮次保留已分配的 token 容量。 - decode 的可选额外预留仅使用满足所有必要分配后剩余的空闲容量。 - """ - prefill, decode, rejected, allocations, headroom = [], [], [], [], [] - prefill_tokens = required = 0 - page_size = self.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 = self.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, self.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: - self.alloc_req_token_indexes(allocations, radix_cache) - return prefill, decode, rejected, prefill_tokens - - 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 + hold_row.view(-1, self.mem_manager.page_size).copy_(hold_page) + self.req_to_token_indexs[self.HOLD_REQUEST_ID].copy_(hold_row) + return def alloc(self): return self.req_list.alloc() - def free(self, free_req_indexes: List[int], free_token_index_blocks): + def free(self, free_req_indexes: List[int], free_token_index): for req_index in free_req_indexes: self.req_list.free(req_index) if self.req_list.is_all_free(): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") - self.mem_manager.allocator.free_token_index_blocks(free_token_index_blocks) + self.mem_manager.free(free_token_index) def free_req(self, free_req_index: int): self.req_list.free(free_req_index) @@ -284,10 +117,10 @@ 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_index_blocks(self, free_token_index_blocks): - self.mem_manager.allocator.free_token_index_blocks(free_token_index_blocks) + def free_token(self, free_token_index): + self.mem_manager.free(free_token_index) return def free_all(self): - self.req_list = _ReqIndexPool(self.max_request_num) + self.req_list = _ReqLinkedList(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 b9ed18648b..549847f4d8 100644 --- a/lightllm/common/req_manager/req_sampling_params.py +++ b/lightllm/common/req_manager/req_sampling_params.py @@ -35,7 +35,6 @@ 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. self.req_to_next_token_scores = ( torch.zeros_like(self.req_to_next_token_ids, dtype=torch.float32) if get_env_start_args().mtp_dynamic_verify @@ -82,13 +81,16 @@ 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_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 + 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) token_id_counter( - prompt_ids=prompt_ids.cuda(non_blocking=True), - out_token_id_counter=self.req_to_out_token_id_counter[req.req_idx], + prompt_ids=prompt_ids, out_token_id_counter=self.req_to_out_token_id_counter[req.req_idx] ) + # TODO: 分析计数初始化后同步当前流的必要性，以及对请求接入和推理并行的影响。 + torch.cuda.current_stream().synchronize() return diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index e12039dcf8..9a1811a50a 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -6,7 +6,6 @@ ) 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 @@ -102,7 +101,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, _ = build_kv_indexes_and_input_ids(model_input, self.req_manager.req_to_token_indexs) + infer_state.mem_index = self._select_mem_indexes(model_input) # 这里调用的是 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 c66090bfe9..58929b1cc1 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -1,3 +1,4 @@ +import os import math import ctypes import asyncio diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index d2fab91e8d..c7fc111a64 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -29,7 +29,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 +from lightllm.server.core.objs.atomic_array_lock import AtomicShmArrayLock, AsyncLock, AtomicLockItem 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 0a70aea5cc..e9d462a256 100644 --- a/lightllm/server/router/batch.py +++ b/lightllm/server/router/batch.py @@ -1,5 +1,6 @@ import time -from typing import List +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple, Union 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 b627ed22ef..d72cb3edd2 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -10,9 +10,10 @@ 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 +from .model_infer.model_rpc import start_model_process, ModelRpcClient 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 48cb9584d9..15ce5f2501 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -6,16 +6,18 @@ import pickle from sortedcontainers import SortedDict -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union from lightllm.common.req_manager import ReqManager, HybridAttentionReqManager -from lightllm.server.core.objs import Req, FinishStatus, ShmReqManager +from lightllm.utils.infer_utils import mark_start, mark_end +from lightllm.server.core.objs import Req, SamplingParams, FinishStatus, ShmReqManager from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import ( HybridAttPagedRadixCache, HybridAttPagedTreeNode, ) from lightllm.utils.log_utils import init_logger +from lightllm.server.req_id_generator import convert_sub_id_to_group_id from lightllm.server.multimodal_params import MultimodalParams from lightllm.utils.envs_utils import get_env_start_args from lightllm.server.pd_io_struct import PDDecodeNodeInfo @@ -87,7 +89,6 @@ 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, _ = r assert r_id not in self.requests_mapping.keys() @@ -98,13 +99,11 @@ 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 下，建立一组请求间的主从关系 @@ -127,27 +126,26 @@ 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_blocks: List, req: "InferReq"): + def free_a_req_mem(self, free_token_index: List, req: "InferReq"): if self.radix_cache is None: - free_token_index_blocks.append(self.req_manager.get_cpu_token_index_blocks(req.req_idx, 0, req.hold_kv_len)) + free_token_index.append(self.req_manager.req_to_token_indexs_cpu[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_blocks=free_token_index_blocks, req=req) + self._free_req_mem_without_radix_insert(free_token_index=free_token_index, req=req) else: if not self.is_hybrid_att_model: - self._full_att_free_req(free_token_index_blocks=free_token_index_blocks, req=req) + self._full_att_free_req(free_token_index=free_token_index, req=req) else: - self._hybrid_att_free_req(free_token_index_blocks=free_token_index_blocks, req=req) + self._hybrid_att_free_req(free_token_index=free_token_index, req=req) assert len(req.hybrid_len_to_big_page_id) == 0 req.cur_kv_len = 0 req.hold_kv_len = 0 + req.move_to_gpu_index_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_blocks: List, req: "InferReq"): + def _free_req_mem_without_radix_insert(self, free_token_index: 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_blocks.append( - self.req_manager.get_cpu_token_index_blocks(req.req_idx, shared_kv_len, req.hold_kv_len) - ) + free_token_index.append(self.req_manager.req_to_token_indexs_cpu[req.req_idx][shared_kv_len : req.hold_kv_len]) if self.is_hybrid_att_model: # 释放请求尾部尚未移交给 radix cache 的 hybrid attention 小页状态。 @@ -165,31 +163,27 @@ def _free_req_mem_without_radix_insert(self, free_token_index_blocks: List, req: req.shared_kv_node = None return - def _full_att_free_req(self, free_token_index_blocks: List, req: "InferReq"): + def _full_att_free_req(self, free_token_index: List, req: "InferReq"): page_size = self.args.page_size cache_kv_len = req.cur_kv_len // page_size * page_size - 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) + # radix cache 只保存完整页，因此不足一页的有效 KV 和页内预留空间都不会参与插入。 + # 先释放截断后的 [cache_kv_len, hold_kv_len) 区间，再处理可缓存前缀的插入和去重。 + free_token_index.append(self.req_manager.req_to_token_indexs_cpu[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") + # The radix cache owns its indexes after this request row is reused. + value = self.req_manager.req_to_token_indexs_cpu[req.req_idx][:cache_kv_len].clone() + + prefix_len, _ = self.radix_cache.insert(key, value) old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len - free_token_index_blocks.append( - self.req_manager.get_cpu_token_index_blocks(req.req_idx, old_prefix_len, prefix_len) - ) + free_token_index.append(self.req_manager.req_to_token_indexs_cpu[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_blocks: List, req: "InferReq"): + def _hybrid_att_free_req(self, free_token_index: 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 @@ -209,32 +203,85 @@ def _hybrid_att_free_req(self, free_token_index_blocks: 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 - cache_kv_len = req.hybrid_cache_len - elif shared_kv_len < tail_big_page_token_num <= req.cur_kv_len: + free_token_index.append( + self.req_manager.req_to_token_indexs_cpu[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_cpu[req.req_idx][: req.cur_kv_len].clone() + 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_cpu[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_cpu[req.req_idx][tail_big_page_token_num : req.hold_kv_len] + ) + req.cur_kv_len = tail_big_page_token_num + 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_cpu[req.req_idx][: req.cur_kv_len].clone() + cur_page_num = tail_big_page_token_num // hash_page_size assert tail_big_page_token_num % hash_page_size == 0 - 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 + 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_cpu[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_cpu[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 - self._free_req_mem_without_radix_insert(free_token_index_blocks, req) + 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 return - 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, - ) + + assert False, f"error state: cur_kv_len: {req.cur_kv_len}" def _save_promptcache_kvbuffer(self): """ @@ -258,7 +305,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_blocks = [] + free_token_index = [] for request_id in finished_request_ids: req: InferReq = self.requests_mapping.pop(request_id) if self.args.diverse_mode: @@ -267,7 +314,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_blocks, req) + self.free_a_req_mem(free_token_index, req) free_req_index.append(req.req_idx) # logger.info(f"infer release req id {req.shm_req.request_id}") @@ -277,9 +324,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_blocks: - free_token_index_blocks = np.concatenate(free_token_index_blocks) - self.req_manager.free(free_req_index, free_token_index_blocks) + if free_token_index: + free_token_index = torch.cat(free_token_index) + self.req_manager.free(free_req_index, free_token_index) 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] @@ -303,12 +350,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_blocks = [] + free_token_index = [] for req in pause_reqs: if self.args.diverse_mode: # 发生暂停的时候，需要清除 diverse 模式下的主从关系 req.clear_master_slave_state() - self.free_a_req_mem(free_token_index_blocks, req) + self.free_a_req_mem(free_token_index, req) assert req.wait_pause is True req.wait_pause = False req.paused = True @@ -316,28 +363,27 @@ 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_blocks) != 0: - free_token_index_blocks = np.concatenate(free_token_index_blocks) - self.req_manager.free_token_index_blocks(free_token_index_blocks) + if len(free_token_index) != 0: + free_token_index = torch.cat(free_token_index) + self.req_manager.free_token(free_token_index) 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 = self.req_manager.get_need_alloc_token_num(req, req.get_cur_total_len()) + alloc_token_num = req._kv_cache_alloc_need(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_prefix_lens=req_prefix_lens) + req._hybrid_match_radix_cache() else: - req._match_radix_cache(req_prefix_lens=req_prefix_lens) + req._match_radix_cache() assert req.paused is True req.paused = False @@ -345,7 +391,6 @@ 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): @@ -489,8 +534,6 @@ 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 @@ -501,6 +544,8 @@ def __init__( # cur_kv_len is the logical length already written; hold_kv_len is the # page-aligned physical capacity owned by this request. self.hold_kv_len = 0 + # Length whose index copies have been enqueued on the inference stream. + self.move_to_gpu_index_len: int = 0 # 请求需要被暂停 self.wait_pause = False @@ -560,9 +605,9 @@ def __init__( if init_prefix_cache: if g_infer_context.is_hybrid_att_model: - self._hybrid_match_radix_cache(req_prefix_lens=req_prefix_lens) + self._hybrid_match_radix_cache() else: - self._match_radix_cache(req_prefix_lens=req_prefix_lens) + self._match_radix_cache() return def _init_all_state(self): @@ -597,14 +642,7 @@ def _init_all_state(self): return - 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): + def _match_radix_cache(self): assert ( g_infer_context.is_hybrid_att_model is False ), "current _match_radix_cache does not support hybrid attention models, to do..." @@ -615,12 +653,19 @@ def _match_radix_cache(self, req_prefix_lens): key = key[0 : len(key) - 1] # 最后一个不需要，因为需要一个额外的token，让其在prefill的时候输出下一个token的值 share_node, kv_len, value_tensor = g_infer_context.radix_cache.match_prefix(key, update_refs=True) if share_node is not None: - self._set_cache_prefix(share_node, value_tensor, req_prefix_lens) + self.shared_kv_node = share_node + ready_cache_len = share_node.node_prefix_total_len + g_infer_context.req_manager.req_to_token_indexs_cpu[self.req_idx, :ready_cache_len].copy_(value_tensor) + self.move_to_gpu_index_len = 0 + 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.shm_req.shm_cur_kv_len = self.cur_kv_len return - def _hybrid_match_radix_cache(self, req_prefix_lens): + def _hybrid_match_radix_cache(self): assert ( g_infer_context.is_hybrid_att_model is True ), "current _hybrid_match_radix_cache only supports hybrid attention models, to do..." @@ -642,41 +687,96 @@ def _hybrid_match_radix_cache(self, req_prefix_lens): share_node, kv_len, value_tensor = g_infer_context.radix_cache.match_prefix( key, block_hashs=block_hashs, update_refs=True ) - 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) - 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: + self.move_to_gpu_index_len = 0 assert self.tail_small_page_buffer_id is None - 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: - 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 - ) + if share_node.is_big_page_node(): + # 大页匹配 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) + ready_cache_len = share_node.node_prefix_total_len + g_infer_context.req_manager.req_to_token_indexs_cpu[self.req_idx, :ready_cache_len].copy_( + 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 + g_infer_context.req_manager.req_to_token_indexs_cpu[self.req_idx, :ready_cache_len].copy_( + 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, + ) + 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_cpu[ + self.req_idx, :cur_big_page_tokens + ].copy_(value_tensor[:cur_big_page_tokens]) + g_infer_context.req_manager.req_to_token_indexs_cpu[ + self.req_idx, cur_big_page_tokens:shared_kv_len + ].copy_(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) + 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 + g_infer_context.req_manager.req_to_token_indexs_cpu[ + self.req_idx, :ready_cache_len + ].copy_(value_tensor[: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 + ) self.shm_req.shm_cur_kv_len = self.cur_kv_len @@ -842,17 +942,39 @@ def _stop_sequences_matched(self, output_len: int): return True return False - def prefill_need_token_num(self, is_chuncked_prefill: bool) -> int: - """本轮计算量；调度通过计算预算检查后再检查 KV 容量。""" + 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。 + """ if is_chuncked_prefill: target_kv_len = self.get_chuncked_input_token_len() else: target_kv_len = self.get_cur_total_len() - return target_kv_len - self.cur_kv_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 容量)。 - def decode_need_token_num(self) -> int: - """普通 decode 为 1；MTP 保留原有的双窗口容量。""" - return 1 if self.mtp_step == 0 else 2 * (1 + self.mtp_step) + 第一个值沿用分页分配前的需求估算：普通 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 class InferReqUpdatePack: diff --git a/lightllm/server/router/model_infer/infer_req_ext.py b/lightllm/server/router/model_infer/infer_req_ext.py index 14a13ff66a..dc22556e18 100644 --- a/lightllm/server/router/model_infer/infer_req_ext.py +++ b/lightllm/server/router/model_infer/infer_req_ext.py @@ -152,7 +152,7 @@ def _mem_indexes(self) -> torch.Tensor: # 延迟导入，避免与 infer_batch 形成模块级循环依赖。 from lightllm.server.router.model_infer.infer_batch import g_infer_context - return g_infer_context.req_manager.req_to_token_indexs[self._req.req_idx] + return g_infer_context.req_manager.req_to_token_indexs_cpu[self._req.req_idx] def collect_prompt_logprobs(self) -> Optional[Tuple[np.ndarray, np.ndarray]]: """从 PromptLogprobsCaptureManager 提取 ``prompt_logprobs>0`` 的 top-k。 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 9c14559076..33d62008e3 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 +from typing import List, Tuple, Callable, Optional, Union from transformers.configuration_utils import PretrainedConfig from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.log_utils import init_logger @@ -13,6 +13,7 @@ from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack from lightllm.server.router.token_load import TokenLoad from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async from lightllm.common.basemodel.logprobs_manager import PromptLogprobsCaptureManager from lightllm.common.basemodel.moe_route_info_manager import MoeRouteInfoManager from lightllm.common.req_manager import HybridAttentionReqManager @@ -411,11 +412,11 @@ def _capture_prompt_logprobs_if_needed( return mgr = PromptLogprobsCaptureManager.get_instance() - mem_indexes = model_input.mem_indexes + mem_indexes = self.model._select_mem_indexes(model_input) 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: @@ -556,6 +557,8 @@ def _read_pd_trans_io_buffer_and_update_req_status(self): self.model.req_manager.req_sampling_params_manager.req_to_next_token_ids ) req_to_next_token_ids[req.req_idx, 0:1].fill_(obj.first_gen_token_id) + # TODO: 分析 PD 首 token 写入后同步当前流的必要性，以及对后续推理的阻塞影响。 + 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, @@ -641,6 +644,49 @@ def remaining_prefill_tokens(req: InferReq) -> int: ready_reqs.insert(0, target_req) return ready_reqs + # 一些可以复用的通用功能函数 + def _alloc_req_kv_mem( + self, + req_obj: InferReq, + alloc_token_num: int, + ) -> 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_cpu[req_obj.req_idx, old_hold_kv_len:new_hold_kv_len].copy_( + mem_indexes + ) + req_obj.hold_kv_len = new_hold_kv_len + return mem_indexes + + def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): + """Collect pending index tails, upload them in one batch, then advance request lengths.""" + req_manager = g_infer_context.req_manager + cpu_table, gpu_table = req_manager.req_to_token_indexs_cpu, req_manager.req_to_token_indexs + pending_copies = [] + src_tensors, dst_tensors = [], [] + for req in run_reqs: + start, end = req.move_to_gpu_index_len, req.hold_kv_len + assert 0 <= start <= end <= cpu_table.shape[1] + if start < end: + src_tensors.append(cpu_table[req.req_idx, start:end]) + dst_tensors.append(gpu_table[req.req_idx, start:end]) + pending_copies.append((req, end)) + if not pending_copies: + return + + cuda_memcpy_pinned_h2d_batch_async(dst_tensors, src_tensors) + gpu_table.record_stream(torch.cuda.current_stream()) + for req, end in pending_copies: + req.move_to_gpu_index_len = end + def _get_classed_reqs( self, req_ids: List[int] = None, @@ -666,7 +712,7 @@ def _get_classed_reqs( 4. prefill_reqs 需要进行prefill操作的请求 5. decode_reqs 需要进行decode操作的请求 - 为通过预算检查的两组请求预留 KV，保持尚未执行请求的跨轮容量。 + 返回前将两组请求尚未上传的 token 索引提交到推理流，并更新 move_to_gpu_index_len。 """ # 定期对 radix cache 进行 merge，防止查询插入的操作效率下降 self._timer_merge_radix_tree() @@ -688,70 +734,99 @@ 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() - def candidates(): - for req_obj in ready_reqs: - if req_obj.filter_mark: - finished_reqs.append(req_obj) - continue + for req_obj in ready_reqs: + 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 + 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 - if no_decode: + if no_decode: + is_decode = False + else: + is_decode = req_obj.cur_kv_len + 1 == req_obj.get_cur_total_len() + if is_decode and strict_prefill and req_obj.cur_kv_len + 1 == req_obj.shm_req.input_len: is_decode = False - else: - is_decode = req_obj.cur_kv_len + 1 == req_obj.get_cur_total_len() - if is_decode and strict_prefill and req_obj.cur_kv_len + 1 == req_obj.shm_req.input_len: - is_decode = False - if is_decode: - token_num = req_obj.decode_need_token_num() + if is_decode: + # 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) + decode_reqs.append(req_obj) + can_alloc_token_num -= alloc_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 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 - prefill_reqs, decode_reqs, rejected, prefill_tokens = g_infer_context.req_manager.classify_reqs_and_alloc_kv( - candidates(), can_alloc_token_num, self.batch_max_tokens, g_infer_context.radix_cache - ) - 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 - else: - finished_reqs.append(req_obj) - self.logger.info( - f"force early finish for PD decode req_id={req_obj.req_id} because token capacity is insufficient" + # 计算预算按本轮实际处理的 token 数累计，KV 预算按需要额外分配的页容量扣减。 + token_num, alloc_token_num = req_obj.prefill_need_token_num( + is_chuncked_prefill=not self.disable_chunked_prefill ) - else: - req_obj.wait_pause = True + 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) + 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 # 先由控制器确定请求需要写入的缓存层级，再按是否包含 CPU cache 决定是否发起 offload。 cache_controller = g_infer_context.cache_placement_controller @@ -778,7 +853,10 @@ def candidates(): 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， @@ -793,6 +871,10 @@ def candidates(): prefill_reqs.append(decode_req) decode_reqs = [] + if prefill_reqs or decode_reqs: + # 使用已有推理流提交映射表更新，与前后推理操作保持顺序。 + with torch.cuda.stream(g_infer_context.get_overlap_stream()): + self._copy_req_token_indexes_to_gpu(prefill_reqs + decode_reqs) return prefill_reqs, decode_reqs # 一些可以复用的通用功能函数 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 7ff3302d96..d5e6c2f210 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,4 +1,5 @@ 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, @@ -87,8 +88,9 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() - 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) + # TODO: 分析 prefill 元数据事件同步的必要性，以及对 CPU/GPU overlap 的影响。 + 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) # 第三阶段 event_pack.notify_forward_and_wait_post_handle() sync_event.synchronize() @@ -127,7 +129,6 @@ 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 @@ -160,28 +161,31 @@ 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, req_prefix_lens=req_prefix_lens) + self._master_req_to_radix_cache(master_req=req_obj) 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, req_prefix_lens=req_prefix_lens) + self._copy_master_req_to_slave_req(slave_req=req_obj) # 在拷贝后，请求独立了，与 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) - self.model.req_manager.copy_token_index_prefixes_to_gpu(req_prefix_lens) + # TODO: 分析 master/slave 状态更新后同步当前流的必要性，以及对 overlap 并行的影响。 + torch.cuda.current_stream().synchronize() return update_func_objs - def _master_req_to_radix_cache(self, master_req: InferReq, req_prefix_lens): + def _master_req_to_radix_cache(self, master_req: InferReq): 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.get_cpu_token_indexes(master_req.req_idx, 0, master_req.cur_kv_len) + value = self.model.req_manager.req_to_token_indexs_cpu[master_req.req_idx][: master_req.cur_kv_len].clone() 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(value[old_prefix_len:prefix_len]) + self.model.mem_manager.free( + self.model.req_manager.req_to_token_indexs_cpu[master_req.req_idx][old_prefix_len:prefix_len] + ) # 将原有共享节点替换为新共享节点，新共享节点对应的长度为当前的cur_kv_len self.radix_cache.dec_node_ref_counter(master_req.shared_kv_node) @@ -193,10 +197,11 @@ def _master_req_to_radix_cache(self, master_req: InferReq, req_prefix_lens): 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.write_token_index_prefix(master_req.req_idx, 0, value, req_prefix_lens=req_prefix_lens) + self.model.req_manager.req_to_token_indexs_cpu[master_req.req_idx, : master_req.cur_kv_len].copy_(value) + master_req.move_to_gpu_index_len = 0 return - def _copy_master_req_to_slave_req(self, slave_req: InferReq, req_prefix_lens): + def _copy_master_req_to_slave_req(self, slave_req: InferReq): master_req = slave_req.related_master_req assert master_req is not None @@ -205,11 +210,12 @@ def _copy_master_req_to_slave_req(self, slave_req: InferReq, req_prefix_lens): kv_len = master_req.cur_kv_len - 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) + self.model.req_manager.req_to_token_indexs_cpu[slave_req.req_idx, :kv_len].copy_( + self.model.req_manager.req_to_token_indexs_cpu[master_req.req_idx, :kv_len] + ) + slave_req.move_to_gpu_index_len = 0 + # TODO: 确认 CPU 镜像复制后是否需要同步当前流，以及对 diverse 并行流程的影响。 + # torch.cuda.current_stream().synchronize() 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 aad675cff4..f60c31190e 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,4 +1,5 @@ # 该文件用于提供在数据dp并行的推理模式下，共享kv cache trans相关的功能函数模块 +import time import numpy as np import dataclasses import torch @@ -8,6 +9,7 @@ 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 @@ -69,21 +71,22 @@ 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 = ( - g_infer_context.req_manager.get_need_alloc_token_num(req, target_kv_len) if trans_size > 0 else 0 - ) + alloc_token_num = req._kv_cache_alloc_need(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(): - mem_indexes = g_infer_context.req_manager.alloc_token_indexes( - req, target_kv_len, g_infer_context.radix_cache - ) + 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] 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 max_kv_len_mem_manager: MemoryManager = self.backend.mem_managers[max_kv_len_mem_manager_index] - max_kv_len_mem_indexes = max_kv_len_mem_manager.req_to_token_indexs[ + max_kv_len_mem_indexes = max_kv_len_mem_manager.req_to_token_indexs_cpu[ max_kv_len_req_idx, req.cur_kv_len : req.cur_kv_len + trans_size - ] + ].clone() trans_tasks.append( TransTask( req=req, 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 34c876f624..b7f5632923 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,25 +1,28 @@ import torch import numpy as np from typing import List, Tuple -from lightllm.server.router.model_infer.infer_batch import InferReq +from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context 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 + total_input_token_num = 0 input_ids = [] b_req_idx = [] b_seq_len = [] b_q_seq_len = [] batch_multimodal_params = [] b_ready_cache_len = [] - b_prefill_start_loc = [0] + b_mtp_index = [] 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) @@ -28,7 +31,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(len(input_token_ids) >= req.get_cur_total_len()) + b_prefill_has_output.append(False if len(input_token_ids) < req.get_cur_total_len() else True) seq_len = len(input_token_ids) input_token_len = seq_len - req.cur_kv_len @@ -37,10 +40,11 @@ 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_input_token_num += input_token_len 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 @@ -53,17 +57,16 @@ 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) - token_ids = torch.empty(b_prefill_start_loc[-1], dtype=torch.int64, device="cpu", pin_memory=True) + token_ids = torch.empty(total_input_token_num, 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, - ) + + b_req_idx, b_seq_len, b_mtp_index, b_ready_cache_len, b_q_seq_len = torch.from_numpy( + np.asarray((b_req_idx, b_seq_len, b_mtp_index, b_ready_cache_len, b_q_seq_len), dtype=np.int32) ).unbind() - b_is_decode_req = torch.tensor(b_is_decode_req, dtype=torch.bool, device="cpu", pin_memory=True) + b_is_decode_req = torch.tensor(b_is_decode_req, dtype=torch.bool, device="cpu") + b_prefill_start_loc = b_q_seq_len.cumsum(dim=0, dtype=torch.int32) - b_q_seq_len model_input = ModelInput( batch_size=b_seq_len.shape[0], @@ -83,7 +86,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> multimodal_params=batch_multimodal_params, ) - return model_input, list(req_objs) + return model_input, run_reqs def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[InferReq]]: @@ -92,41 +95,50 @@ 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}" - for step in range(req.mtp_step + 1): + 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): run_reqs.append(req) b_req_idx.append(req.req_idx) - b_seq_len.append(seq_len + step) - total_token_num += seq_len + step - b_mtp_index.append(step) + seq_len += 1 + b_seq_len.append(seq_len) + total_token_num += seq_len + b_mtp_index.append(step + 1) 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, b_seq_len, b_mtp_index = torch.from_numpy( - np.asarray((b_req_idx, b_seq_len, b_mtp_index), dtype=np.int32) + b_shared_seq_len = [req.get_radix_cache_shared_len() for req in run_reqs] + b_req_idx, b_seq_len, b_mtp_index, b_shared_seq_len = torch.from_numpy( + np.asarray((b_req_idx, b_seq_len, b_mtp_index, b_shared_seq_len), dtype=np.int32) ).unbind() b_position_delta = build_b_position_delta(multimodal_params) - b_shared_seq_len = torch.tensor( - [req.get_radix_cache_shared_len() for req in run_reqs], dtype=torch.int32, device="cpu" - ) b_shared_radix_node_id = torch.tensor( [-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=1, + max_q_seq_len=max_q_seq_len, 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 f4b2a8f76d..0703a08f98 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,3 +1,4 @@ +import threading import torch.distributed as dist import torch import dataclasses @@ -7,9 +8,11 @@ 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 @@ -92,6 +95,7 @@ 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: # 计算需要加载的页面（只加载未匹配的部分） @@ -99,9 +103,7 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): assert ready_page_num <= len(page_list) need_pages = page_list[ready_page_num:] # 只取需要的页面 - mem_indexes = g_infer_context.req_manager.alloc_token_indexes( - req, match_tokens, g_infer_context.radix_cache - ) + self.backend._alloc_req_kv_mem(req, need_token_num) if self.need_sync_compute_stream(): # TODO fa3 现在必须使用同步模式, 未来需要移除 @@ -111,7 +113,6 @@ 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.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 ) @@ -119,10 +120,11 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): # 这里需要进行pad操作，使操作的页面是完整的。 _start = page_len_start_list[ready_page_num] - _end = req.cur_kv_len - assert 0 <= _start <= _end, f"invalid pad range [{_start}, {_end}]" - mem_indexes_cuda = torch.cat( - [req_manager.req_to_token_indexs[req.req_idx, _start:_end], mem_indexes_cuda] + _end = req.hold_kv_len + assert 0 <= _start <= req.cur_kv_len <= _end, f"invalid pad range [{_start}, {_end}]" + # 旧前缀和新分配的索引已连续写入 pinned CPU 镜像，一次上传即可。 + mem_indexes_cuda = req_manager.req_to_token_indexs_cpu[req.req_idx, _start:_end].cuda( + non_blocking=True ) assert ( @@ -223,7 +225,6 @@ def _start_kv_cache_offload_task( assert len(token_hash_list) == len(page_len_list) if self.backend.is_master_in_dp: - find_index = bisect.bisect_right(page_len_list, req.cur_kv_len) move_block_size = find_index @@ -275,8 +276,9 @@ 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] + token_indexes = self.backend.model.req_manager.req_to_token_indexs_cpu[req.req_idx, :move_token_num].cuda( + non_blocking=True + ) 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 476eeda68f..51548ac16c 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,3 +1,4 @@ +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 @@ -101,14 +102,15 @@ 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 - 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 - ) - ) + 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_cpu[ + req_obj.req_idx, old_prefix_len : req_obj.hold_kv_len + ].clone() + self.model.mem_manager.free(mem_indexes) req_obj.cur_kv_len = old_prefix_len req_obj.hold_kv_len = old_prefix_len + req_obj.move_to_gpu_index_len = min(req_obj.move_to_gpu_index_len, old_prefix_len) if self.is_master_in_dp: req_obj.shm_req.shm_cur_kv_len = req_obj.cur_kv_len @@ -126,7 +128,14 @@ 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 - mem_indexes = self.model.req_manager.alloc_token_indexes(req_obj, input_len, self.radix_cache) + 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] 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 0d22f9f9a1..4357db5359 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,5 +1,6 @@ import torch.multiprocessing as mp -from typing import List +import random +from typing import List, Tuple 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 @@ -121,9 +122,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.get_cpu_token_indexes( - req_obj.req_idx, kv_start_index, kv_end_index - ).tolist() + mem_indexes = self.model.req_manager.req_to_token_indexs_cpu[ + 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 e5cd7d54fd..9a82cf85da 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, Optional +from typing import TYPE_CHECKING, List, Optional import torch diff --git a/lightllm/server/router/req_queue/base_queue.py b/lightllm/server/router/req_queue/base_queue.py index 2fbd8329df..0c25b8949a 100644 --- a/lightllm/server/router/req_queue/base_queue.py +++ b/lightllm/server/router/req_queue/base_queue.py @@ -1,5 +1,6 @@ import time -from typing import List +from typing import List, Dict +from lightllm.utils.infer_utils import calculate_time 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 07b8ef32f5..f7f7173ebb 100644 --- a/lightllm/server/router/req_queue/dp_base_queue.py +++ b/lightllm/server/router/req_queue/dp_base_queue.py @@ -1,3 +1,4 @@ +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 9ec3945bc9..6ecfffbea6 100644 --- a/lightllm/utils/custom_kernel_utis.py +++ b/lightllm/utils/custom_kernel_utis.py @@ -4,6 +4,29 @@ 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 + # TODO: 分析 D2H 完成等待的阻塞开销，以及是否能缩小当前流的同步范围。 + 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 f5880ea268..b77432b6ad 100644 --- a/skills/test_model/SKILL.md +++ b/skills/test_model/SKILL.md @@ -57,43 +57,3 @@ 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 ea4bfd3f18..4dec53ae4d 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,502 +875,5 @@ 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 58d7080edd..7205c34e53 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -568,21 +568,22 @@ 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)) - req_manager = self.model.req_manager - page_size = req_manager.mem_manager.page_size - allocations = [] + mem_manager = self.model.req_manager.mem_manager + page_size = mem_manager.page_size + cpu_table = self.model.req_manager.req_to_token_indexs_cpu + gpu_table = self.model.req_manager.req_to_token_indexs 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 - 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 + alloc_token_num = new_hold_kv_len - old_hold_kv_len + mem_indexes = mem_manager.alloc(alloc_token_num) + cpu_indexes = cpu_table[req_id, old_hold_kv_len:new_hold_kv_len] + cpu_indexes.copy_(mem_indexes) + gpu_table[req_id, old_hold_kv_len:new_hold_kv_len].copy_(cpu_indexes, non_blocking=True) + self._hold_kv_len_by_req[req_id] = new_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: @@ -706,6 +707,11 @@ 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 new file mode 100644 index 0000000000..107728c26d --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_fp8_page_size.py @@ -0,0 +1,45 @@ +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_input.py b/unit_tests/common/basemodel/test_model_input.py index 6672cccc31..68e02717f3 100644 --- a/unit_tests/common/basemodel/test_model_input.py +++ b/unit_tests/common/basemodel/test_model_input.py @@ -104,6 +104,7 @@ def test_prefill_rejects_position_delta(): model_input.check_input() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_padded_prefill_adds_non_decode_request_marker(): model_input = ModelInput( batch_size=1, @@ -127,6 +128,7 @@ def test_padded_prefill_adds_non_decode_request_marker(): req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), ) + model_input.to_cuda() padded_input = TpPartBaseModel._create_padded_prefill_model_input( model, model_input=model_input, @@ -137,6 +139,7 @@ def test_padded_prefill_adds_non_decode_request_marker(): assert padded_input.b_is_decode_req.tolist() == [True, False] +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_padded_prefill_builds_internal_request_for_empty_input(): model_input = ModelInput( batch_size=0, @@ -161,6 +164,7 @@ def test_padded_prefill_builds_internal_request_for_empty_input(): req_manager=SimpleNamespace(HOLD_REQUEST_ID=88), ) + model_input.to_cuda() padded_input = TpPartBaseModel._create_padded_prefill_model_input( model, model_input=model_input, @@ -176,6 +180,7 @@ def test_padded_prefill_builds_internal_request_for_empty_input(): assert torch.equal(padded_input.mtp_draft_input_hiddens, torch.zeros((1, 4))) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_padded_decode_builds_internal_request_from_empty_token_tensor(): model_input = ModelInput( batch_size=0, @@ -197,6 +202,7 @@ def test_padded_decode_builds_internal_request_from_empty_token_tensor(): req_manager=SimpleNamespace(HOLD_REQUEST_ID=88), ) + model_input.to_cuda() padded_input = TpPartBaseModel._create_padded_decode_model_input( model, model_input=model_input, diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index ca035fdaf2..ba95189e90 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -3,6 +3,7 @@ import torch import pytest +from lightllm.common.basemodel import basemodel from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput @@ -76,7 +77,7 @@ def test_prefill_unpad_restores_empty_output(): def _create_empty_decode_input(): - return ModelInput( + model_input = ModelInput( batch_size=0, total_token_num=0, max_q_seq_len=1, @@ -91,9 +92,12 @@ def _create_empty_decode_input(): is_prefill=False, multimodal_params=[], ) + model_input.to_cuda() + return model_input @torch.no_grad() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): execution_configs = ( # eager 普通模式：空 batch 只补一个 dummy request。 @@ -130,7 +134,7 @@ def pad_once(model_input, new_batch_size): graph_flags_at_att_init = [] - def create_infer_state(model_input, microbatch_index=0): + def create_infer_state(model_input): infer_state = SimpleNamespace( b_req_idx=model_input.b_req_idx, b_seq_len=model_input.b_seq_len, @@ -163,6 +167,7 @@ def create_infer_state(model_input, microbatch_index=0): assert output.logits.shape == (0, 4) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @pytest.mark.parametrize("overlap", [False, True]) @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize("graph_mode", ["eager", "capture", "replay"]) @@ -217,6 +222,7 @@ def forward(state): model_input.b_req_idx = torch.tensor([req for req, _ in request_rows], dtype=torch.int32) model_input.b_mtp_index = torch.tensor([step for _, step in request_rows], dtype=torch.int32) model_input.b_seq_len = 10 + model_input.b_mtp_index + model_input.to_cuda() expected = torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1) if overlap: output, empty_output = model._microbatch_overlap_decode_cuda(model_input, _create_empty_decode_input()) diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py index 847e2a6494..54f2ccacf1 100644 --- a/unit_tests/common/basemodel/test_overlap_utils.py +++ b/unit_tests/common/basemodel/test_overlap_utils.py @@ -82,6 +82,12 @@ 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( @@ -101,14 +107,14 @@ def fake_overlap_forward(input0, input1): outputs = model.microbatch_overlap_prefill(model_input0, model_input1) - assert events == ["to_cuda", "to_cuda", "forward"] + assert events == ["to_cuda", "gather", "to_cuda", "gather", "forward"] assert captured_inputs == [model_input0, model_input1] - assert captured_inputs[0].input_ids.tolist() == [10, 11] - assert captured_inputs[1].input_ids.tolist() == [12] + assert captured_inputs[0].input_ids.tolist() == [10, 91] + assert captured_inputs[1].input_ids.tolist() == [92] assert len(outputs) == 2 -def test_base_model_decode_defers_gather_until_padded_forward(monkeypatch): +def test_base_model_decode_accepts_two_inputs_and_skips_empty_gather(monkeypatch): model_input0 = _make_decode_input() model_input0.input_ids = None model_input1 = _make_empty_decode_input() @@ -122,6 +128,12 @@ 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( @@ -141,9 +153,10 @@ def fake_overlap_forward(input0, input1): model.microbatch_overlap_decode(model_input0, model_input1) - assert events == ["to_cuda", "to_cuda", "forward"] + assert events == ["to_cuda", "gather", "to_cuda", "forward"] assert captured_inputs == [model_input0, model_input1] - assert all(model_input.input_ids is None for model_input in captured_inputs) + assert captured_inputs[0].input_ids.tolist() == [40, 41, 42, 43, 44, 45] + assert captured_inputs[1].input_ids.numel() == 0 @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 46b9610c0d..55ef9fcf37 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 @@ -24,7 +24,6 @@ @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 fe1b968cf3..2f8544cf75 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,8 +117,6 @@ 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 14f9047289..11035c83a3 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,10 +1,12 @@ 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) @@ -24,13 +26,32 @@ def _run_both( a_raw, b_raw, ): - # 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( + """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( q=q, k=k, v=v, - initial_state=initial_state.clone(), + initial_state=state_new, cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, ssm_state_write_indices=ssm_state_write_indices, @@ -40,10 +61,8 @@ def _run_both( a_raw=a_raw, b_raw=b_raw, ) - 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) + + return o_old, o_new, fs_old, fs_new @pytest.mark.parametrize("mtp_step", [1, 2, 3]) @@ -63,11 +82,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.randperm(cache_slots, device="cuda", dtype=torch.int32)[:num_tokens].view(batch, seqlen) + ssm_idx = torch.randint(0, cache_slots, (batch, seqlen), device="cuda", dtype=torch.int32) cu_seqlens = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * seqlen num_accepted = torch.full((batch,), seqlen, device="cuda", dtype=torch.int32) - _run_both( + o_old, o_new, fs_old, fs_new = _run_both( q, k, v, @@ -82,6 +101,15 @@ 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): @@ -107,10 +135,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.randperm(cache_slots, device="cuda", dtype=torch.int32)[: batch * max_len].view(batch, max_len) + ssm_idx = torch.randint(0, cache_slots, (batch, max_len), device="cuda", dtype=torch.int32) num_accepted = lengths.to(device="cuda", dtype=torch.int32) - _run_both( + o_old, o_new, fs_old, fs_new = _run_both( q, k, v, @@ -125,6 +153,15 @@ 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 c1ef686f8e..55e7d0724a 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py +++ b/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py @@ -4,7 +4,19 @@ 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 +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]] 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_pad_decode_inputs.py b/unit_tests/common/basemodel/triton_kernel/test_pad_decode_inputs.py new file mode 100644 index 0000000000..15302ef9ad --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_pad_decode_inputs.py @@ -0,0 +1,92 @@ +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.triton_kernel.pad_decode_inputs import pad_decode_inputs + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _make_inputs(batch_size, has_position_delta): + return dict( + input_ids=torch.arange(batch_size, dtype=torch.int64, device="cuda") + 2 ** 40, + b_req_idx=torch.arange(batch_size, dtype=torch.int32, device="cuda") + 10, + b_mtp_index=torch.arange(batch_size, dtype=torch.int32, device="cuda") % 4, + b_seq_len=torch.arange(batch_size, dtype=torch.int32, device="cuda") + 5, + b_position_delta=( + torch.arange(batch_size, dtype=torch.int32, device="cuda") - 9 if has_position_delta else None + ), + b_shared_seq_len=torch.arange(batch_size, dtype=torch.int32, device="cuda") % 3, + b_shared_radix_node_id=torch.arange(batch_size, dtype=torch.int64, device="cuda") + 2 ** 40, + ) + + +def _reference(inputs, padded_batch_size, hold_request_id): + return ( + F.pad(inputs["input_ids"], (0, padded_batch_size), value=1), + F.pad(inputs["b_req_idx"], (0, padded_batch_size), value=hold_request_id), + F.pad(inputs["b_mtp_index"], (0, padded_batch_size), value=0), + F.pad(inputs["b_seq_len"], (0, padded_batch_size), value=2), + ( + F.pad(inputs["b_position_delta"], (0, padded_batch_size), value=0) + if inputs["b_position_delta"] is not None + else None + ), + F.pad(inputs["b_shared_seq_len"], (0, padded_batch_size), value=0), + F.pad(inputs["b_shared_radix_node_id"], (0, padded_batch_size), value=-1), + ) + + +@pytest.mark.parametrize("hold_request_id", [-1, 88]) +@pytest.mark.parametrize("has_position_delta", [False, True]) +@pytest.mark.parametrize( + "batch_size,new_batch_size", + [(0, 1), (0, 257), (1, 2), (3, 8), (255, 256), (256, 257), (257, 1025), (1025, 4096)], +) +def test_pad_decode_inputs_matches_torch(batch_size, new_batch_size, has_position_delta, hold_request_id): + inputs = _make_inputs(batch_size, has_position_delta) + originals = {name: value.clone() for name, value in inputs.items() if value is not None} + expected = _reference(inputs, new_batch_size - batch_size, hold_request_id) + + actual = pad_decode_inputs(**inputs, new_batch_size=new_batch_size, hold_request_id=hold_request_id) + + for output, reference in zip(actual, expected): + if reference is None: + assert output is None + else: + assert output.dtype == reference.dtype + assert output.device == reference.device + assert output.shape == reference.shape + assert torch.equal(output, reference) + for name, original in originals.items(): + assert torch.equal(inputs[name], original) + + +@pytest.mark.parametrize("has_position_delta", [False, True]) +def test_pad_decode_inputs_cuda_graph_replays_updated_inputs(has_position_delta): + inputs = _make_inputs(batch_size=3, has_position_delta=has_position_delta) + kwargs = dict(**inputs, new_batch_size=513, hold_request_id=88) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + pad_decode_inputs(**kwargs) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + outputs = pad_decode_inputs(**kwargs) + + inputs["input_ids"].add_(7) + inputs["b_req_idx"].add_(2) + inputs["b_shared_radix_node_id"].add_(11) + if has_position_delta: + inputs["b_position_delta"].sub_(3) + graph.replay() + + for output, reference in zip(outputs, _reference(inputs, padded_batch_size=510, hold_request_id=88)): + if reference is None: + assert output is None + else: + assert torch.equal(output, reference) diff --git a/unit_tests/common/basemodel/triton_kernel/test_pad_prefill_inputs.py b/unit_tests/common/basemodel/triton_kernel/test_pad_prefill_inputs.py new file mode 100644 index 0000000000..17ab85aee1 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_pad_prefill_inputs.py @@ -0,0 +1,86 @@ +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.triton_kernel.pad_prefill_inputs import pad_prefill_inputs + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _make_inputs(batch_size, handle_token_num): + return dict( + input_ids=torch.arange(handle_token_num, dtype=torch.int64, device="cuda") + 2 ** 40, + b_req_idx=torch.arange(batch_size, dtype=torch.int32, device="cuda") + 10, + b_mtp_index=torch.arange(batch_size, dtype=torch.int32, device="cuda") % 4, + b_seq_len=torch.arange(batch_size, dtype=torch.int32, device="cuda") + 5, + b_ready_cache_len=torch.arange(batch_size, dtype=torch.int32, device="cuda") % 3, + b_is_decode_req=torch.arange(batch_size, device="cuda") % 2 == 0, + ) + + +def _reference(inputs, padded_token_num, hold_request_id): + return ( + F.pad(inputs["input_ids"], (0, padded_token_num), value=1), + F.pad(inputs["b_req_idx"], (0, 1), value=hold_request_id), + F.pad(inputs["b_mtp_index"], (0, 1), value=0), + F.pad(inputs["b_seq_len"], (0, 1), value=padded_token_num), + F.pad(inputs["b_ready_cache_len"], (0, 1), value=0), + F.pad(inputs["b_is_decode_req"], (0, 1), value=False), + ) + + +@pytest.mark.parametrize("hold_request_id", [-1, 88]) +@pytest.mark.parametrize( + "batch_size,handle_token_num,padded_token_num", + [ + (0, 0, 1), + (0, 0, 129), + (1, 1, 3), + (3, 17, 1), + (257, 2051, 11), + (1025, 4099, 257), + (3, 65537, 7), + ], +) +def test_pad_prefill_inputs_matches_torch(batch_size, handle_token_num, padded_token_num, hold_request_id): + inputs = _make_inputs(batch_size, handle_token_num) + originals = {name: value.clone() for name, value in inputs.items()} + expected = _reference(inputs, padded_token_num, hold_request_id) + + actual = pad_prefill_inputs( + **inputs, + new_handle_token_num=handle_token_num + padded_token_num, + hold_request_id=hold_request_id, + ) + + for output, reference in zip(actual, expected): + assert output.dtype == reference.dtype + assert output.device == reference.device + assert output.shape == reference.shape + assert torch.equal(output, reference) + for name, original in originals.items(): + assert torch.equal(inputs[name], original) + + +def test_pad_prefill_inputs_cuda_graph_replays_updated_inputs(): + inputs = _make_inputs(batch_size=3, handle_token_num=17) + kwargs = dict(**inputs, new_handle_token_num=24, hold_request_id=88) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + pad_prefill_inputs(**kwargs) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + outputs = pad_prefill_inputs(**kwargs) + + inputs["input_ids"].add_(7) + inputs["b_req_idx"].add_(2) + inputs["b_is_decode_req"].logical_not_() + graph.replay() + + for output, reference in zip(outputs, _reference(inputs, padded_token_num=7, hold_request_id=88)): + assert torch.equal(output, reference) 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 15a4447466..0e66266dea 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,3 +48,34 @@ 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 80cc75ee33..ec67af6077 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py +++ b/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py @@ -81,6 +81,39 @@ 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 new file mode 100644 index 0000000000..0dab062934 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py @@ -0,0 +1,75 @@ +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 772644c28f..3e14d5baaa 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,19 +17,14 @@ ], ) 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) + assert torch.equal(in_tensor, out_tensor[0:token_num, :]) if __name__ == "__main__": diff --git a/unit_tests/common/test_paged_mem_manager.py b/unit_tests/common/test_paged_mem_manager.py new file mode 100644 index 0000000000..68bc514991 --- /dev/null +++ b/unit_tests/common/test_paged_mem_manager.py @@ -0,0 +1,32 @@ +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 new file mode 100644 index 0000000000..594157560b --- /dev/null +++ b/unit_tests/common/test_pinned_memory_non_blocking_copy.py @@ -0,0 +1,39 @@ +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 new file mode 100644 index 0000000000..84fe266613 --- /dev/null +++ b/unit_tests/common/test_quantized_paged_mem_manager.py @@ -0,0 +1,33 @@ +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_cpu.py b/unit_tests/common/test_req_manager_cpu.py new file mode 100644 index 0000000000..fd351e7d49 --- /dev/null +++ b/unit_tests/common/test_req_manager_cpu.py @@ -0,0 +1,645 @@ +from types import SimpleNamespace +from multiprocessing.reduction import ForkingPickler +from contextlib import nullcontext +import builtins + +import pytest +import torch + +# Use the production import order to avoid the existing package initialization cycle. +from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.req_manager.base import ReqManager, _ReqLinkedList +from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async +from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext +from lightllm.server.router.model_infer import infer_batch +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend + + +@pytest.fixture(autouse=True) +def reset_batch_copy_availability(monkeypatch): + from lightllm.common import cuda_wrapper + + monkeypatch.setattr(cuda_wrapper, "_cuda_memcpy_batch_available", True) + + +def _make_manager(device="cpu"): + manager = ReqManager.__new__(ReqManager) + manager.max_request_num = 2 + manager.req_list = _ReqLinkedList(2) + manager.req_to_token_indexs_cpu = torch.full( + (3, 16), -1, dtype=torch.int32, pin_memory=torch.device(device).type == "cuda" + ) + manager.req_to_token_indexs = torch.full((3, 16), -1, dtype=torch.int32, device=device) + return manager + + +def test_cpu_updates_own_allocator_values_and_defer_gpu_publication(): + manager = _make_manager() + indexes = torch.tensor([20, 21, 22, 23], dtype=torch.int32) + manager.req_to_token_indexs_cpu[0, 4:8].copy_(indexes) + indexes.fill_(99) + manager.req_to_token_indexs_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) + + assert manager.req_to_token_indexs_cpu[0, :8].tolist() == [0, 1, 2, 3, 20, 21, 22, 23] + assert manager.req_to_token_indexs.eq(-1).all() + + +def test_releasing_request_memory_resets_uploaded_length(): + manager = _make_manager() + manager.req_to_token_indexs_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) + context = InferenceContext() + context.req_manager = manager + req = SimpleNamespace( + req_idx=0, + cur_kv_len=6, + hold_kv_len=8, + move_to_gpu_index_len=8, + shm_req=SimpleNamespace(shm_cur_kv_len=6), + ) + released = [] + context.free_a_req_mem(released, req) + + assert released[0].tolist() == list(range(8)) + assert req.cur_kv_len == req.hold_kv_len == req.move_to_gpu_index_len == 0 + + +def test_shared_cpu_mirror_exposes_updates_before_gpu_publication(): + manager = _make_manager() + source = manager.req_to_token_indexs_cpu.share_memory_() + remote = ForkingPickler.loads(ForkingPickler.dumps(source)) + manager.req_to_token_indexs_cpu[0, :2].copy_(torch.tensor([8, 9], dtype=torch.int32)) + + assert remote[0, :2].tolist() == [8, 9] + assert manager.req_to_token_indexs.eq(-1).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("registration_fails", [False, True]) +def test_shared_export_preserves_pinned_mirror_and_async_upload(monkeypatch, registration_fails): + from lightllm.common.kv_cache_mem_manager import mem_manager as impl + + req_manager = _make_manager("cuda") + mem_manager = impl.MemoryManager.__new__(impl.MemoryManager) + runtime = torch.cuda.cudart() + registrations, cleanups, remote_tables = [], [], [] + + def register(ptr, size, flags): + registrations.append((ptr, size)) + return 1 if registration_fails else runtime.cudaHostRegister(ptr, size, flags) + + original_finalize = impl.weakref.finalize + + def finalize(obj, *args): + cleanup = original_finalize(obj, *args) + cleanups.append(cleanup) + return cleanup + + def serialize(manager): + assert manager.req_to_token_indexs_cpu.is_pinned() + remote_tables.append(ForkingPickler.loads(ForkingPickler.dumps(manager.req_to_token_indexs_cpu))) + return memoryview(b"test") + + monkeypatch.setattr( + torch.cuda, + "cudart", + lambda: SimpleNamespace( + cudaError=runtime.cudaError, + cudaHostRegister=register, + cudaHostUnregister=runtime.cudaHostUnregister, + ), + ) + monkeypatch.setattr(impl.weakref, "finalize", finalize) + monkeypatch.setattr(impl, "ForkingPickler", SimpleNamespace(dumps=serialize)) + monkeypatch.setattr(impl, "kv_trans_use_p2p", lambda: False) + monkeypatch.setattr(impl, "FileLock", lambda path: nullcontext()) + monkeypatch.setattr(impl, "get_node_world_size", lambda: 1) + monkeypatch.setattr(impl, "get_current_rank_in_node", lambda: 0) + monkeypatch.setattr(impl, "get_service_shm_name", lambda name: name) + monkeypatch.setattr( + impl, + "create_or_link_shm", + lambda name, expected_size, **kwargs: SimpleNamespace( + name=name, size=expected_size, buf=bytearray(expected_size) + ), + ) + + try: + if registration_fails: + with pytest.raises(RuntimeError, match="cudaHostRegister failed"): + mem_manager.write_to_shm(req_manager) + assert not cleanups and not remote_tables + return + + mem_manager.write_to_shm(req_manager) + mem_manager.write_to_shm(req_manager) + source = req_manager.req_to_token_indexs_cpu + assert mem_manager.req_to_token_indexs_cpu is source + assert source.is_shared() and source.is_pinned() + assert len(registrations) == len(cleanups) == 1 + source[0, :4].copy_(torch.arange(40, 44, dtype=torch.int32)) + assert all(table[0, :4].tolist() == [40, 41, 42, 43] for table in remote_tables) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + cuda_memcpy_pinned_h2d_batch_async([req_manager.req_to_token_indexs[0, :4]], [source[0, :4]]) + stream.synchronize() + assert req_manager.req_to_token_indexs[0, :4].cpu().tolist() == [40, 41, 42, 43] + finally: + torch.cuda.synchronize() + for cleanup in cleanups: + assert cleanup() == runtime.cudaError.success + + +def test_prefix_match_invalidates_uploaded_length(monkeypatch): + manager = _make_manager() + req = InferReq.__new__(InferReq) + req.req_idx = 0 + req.args = SimpleNamespace(page_size=4) + req.cur_kv_len = 0 + req.hold_kv_len = req.move_to_gpu_index_len = 4 + req.sampling_param = SimpleNamespace(disable_prompt_cache=False) + req.shm_req = SimpleNamespace(shm_prompt_ids=SimpleNamespace(arr=list(range(8)))) + req.get_cur_total_len = lambda: 8 + node = SimpleNamespace(node_prefix_total_len=4) + indexes = torch.arange(32, 36, dtype=torch.int32) + context = SimpleNamespace( + is_hybrid_att_model=False, + req_manager=manager, + radix_cache=SimpleNamespace(match_prefix=lambda key, update_refs: (node, 4, indexes)), + ) + monkeypatch.setattr(infer_batch, "g_infer_context", context) + + req._match_radix_cache() + + assert req.move_to_gpu_index_len == 0 + assert req.cur_kv_len == req.hold_kv_len == 4 + assert manager.req_to_token_indexs_cpu[0, :4].tolist() == [32, 33, 34, 35] + + +def test_radix_insertion_owns_cpu_indexes_before_request_row_reuse(): + manager = _make_manager() + manager.req_to_token_indexs_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) + cached = [] + + def insert(key, indexes): + cached.append(indexes) + return 0, None + + context = InferenceContext() + context.args = SimpleNamespace(page_size=4) + context.req_manager = manager + context.radix_cache = SimpleNamespace(insert=insert) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=6, + hold_kv_len=8, + shared_kv_node=None, + get_input_token_ids=lambda: list(range(6)), + ) + released = [] + context._full_att_free_req(released, req) + assert released[0].tolist() == [4, 5, 6, 7] + manager.req_to_token_indexs_cpu[0, :8].fill_(99) + + assert cached[0].tolist() == [0, 1, 2, 3] + assert manager.req_to_token_indexs.eq(-1).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_async_batch_reads_source_after_prior_stream_writes_and_preserves_untouched_ranges(): + manager = _make_manager("cuda") + incoming_indexes = torch.tensor([40, 41, 42, 43], dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + torch.cuda._sleep(2_000_000) + manager.req_to_token_indexs_cpu[0, 4:8].copy_(incoming_indexes, non_blocking=True) + cuda_memcpy_pinned_h2d_batch_async( + [manager.req_to_token_indexs[0, 4:8]], + [manager.req_to_token_indexs_cpu[0, 4:8]], + ) + stream.synchronize() + + assert manager.req_to_token_indexs[0, 4:8].cpu().tolist() == [40, 41, 42, 43] + assert manager.req_to_token_indexs[0, :4].eq(-1).all() + assert manager.req_to_token_indexs[0, 8:].eq(-1).all() + assert manager.req_to_token_indexs[1:].eq(-1).all() + manager.req_to_token_indexs_cpu[0, 4:8].copy_(torch.tensor([80, 81, 82, 83], dtype=torch.int32)) + with torch.cuda.stream(stream): + cuda_memcpy_pinned_h2d_batch_async( + [manager.req_to_token_indexs[0, 4:8]], + [manager.req_to_token_indexs_cpu[0, 4:8]], + ) + stream.synchronize() + assert manager.req_to_token_indexs[0, 4:8].cpu().tolist() == [80, 81, 82, 83] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("default_stream", [False, True]) +def test_native_batch_uses_stream_order_and_handles_disjoint_ranges(monkeypatch, default_stream): + from cuda.bindings import runtime as cuda_runtime + + manager = _make_manager("cuda") + manager.req_to_token_indexs_cpu.copy_(torch.arange(48, dtype=torch.int32).reshape(3, 16)) + ranges = [(0, 4, 8), (1, 1, 5), (0, 12, 16)] + expected = torch.full_like(manager.req_to_token_indexs_cpu, -1) + for row, start, end in ranges: + expected[row, start:end].copy_(manager.req_to_token_indexs_cpu[row, start:end]) + calls = [] + native_copy = cuda_runtime.cudaMemcpyBatchAsync + + def record_native(*args): + calls.append((args[3], args[2], args[4][0].srcAccessOrder)) + return native_copy(*args) + + monkeypatch.setattr(cuda_runtime, "cudaMemcpyBatchAsync", record_native) + stream = torch.cuda.default_stream() if default_stream else torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + if default_stream: + from lightllm.common import cuda_wrapper as base + + with pytest.raises(AssertionError, match="非默认 CUDA 流"): + cuda_memcpy_pinned_h2d_batch_async( + [manager.req_to_token_indexs[row, start:end] for row, start, end in ranges], + [manager.req_to_token_indexs_cpu[row, start:end] for row, start, end in ranges], + ) + assert calls == [] + assert base._cuda_memcpy_batch_available + return + torch.cuda._sleep(2_000_000) + cuda_memcpy_pinned_h2d_batch_async( + [manager.req_to_token_indexs[row, start:end] for row, start, end in ranges], + [manager.req_to_token_indexs_cpu[row, start:end] for row, start, end in ranges], + ) + result = manager.req_to_token_indexs.clone() + stream.synchronize() + + assert calls == [(3, [16, 16, 16], cuda_runtime.cudaMemcpySrcAccessOrder.cudaMemcpySrcAccessOrderStream)] + assert torch.equal(result.cpu(), expected) + + +def test_empty_batch_does_not_call_cuda(monkeypatch): + original_import = builtins.__import__ + + def check_import(name, *args, **kwargs): + if name == "cuda.bindings": + pytest.fail("empty batch must not import CUDA bindings") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", check_import) + cuda_memcpy_pinned_h2d_batch_async([], []) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("failure", ["missing_binding", "missing_api", "native_error", "exception"]) +def test_batch_fallback_uses_current_stream_without_synchronizing(monkeypatch, failure): + from cuda.bindings import runtime as cuda_runtime + from lightllm.common import cuda_wrapper as base + + manager = _make_manager("cuda") + manager.req_to_token_indexs_cpu.copy_(torch.arange(48, dtype=torch.int32).reshape(3, 16)) + ranges = [(0, 4, 8), (1, 1, 5), (0, 12, 16)] + expected = torch.full_like(manager.req_to_token_indexs_cpu, -1) + for row, start, end in ranges: + expected[row, start:end].copy_(manager.req_to_token_indexs_cpu[row, start:end]) + + batch_attempts = [] + if failure == "missing_binding": + original_import = builtins.__import__ + + def missing_import(name, *args, **kwargs): + if name == "cuda.bindings": + batch_attempts.append(name) + raise ModuleNotFoundError("cuda.bindings is unavailable") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", missing_import) + elif failure == "missing_api": + monkeypatch.delattr(cuda_runtime, "cudaMemcpyBatchAsync") + else: + + def failed_batch(*args): + batch_attempts.append(args) + if failure == "exception": + raise RuntimeError("batch API could not be loaded") + return (cuda_runtime.cudaError_t.cudaErrorNotSupported,) + + monkeypatch.setattr(cuda_runtime, "cudaMemcpyBatchAsync", failed_batch) + + loop_calls = [] + torch_copy = torch.Tensor.copy_ + + def record_copy(dst, src, *args, **kwargs): + if dst.is_cuda and src.device.type == "cpu": + assert kwargs["non_blocking"] is True + loop_calls.append((dst.numel() * dst.element_size(), torch.cuda.current_stream(dst.device).cuda_stream)) + return torch_copy(dst, src, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "copy_", record_copy) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + dsts = [manager.req_to_token_indexs[row, start:end] for row, start, end in ranges] + srcs = [manager.req_to_token_indexs_cpu[row, start:end] for row, start, end in ranges] + + def unexpected_sync(*args, **kwargs): + pytest.fail("copy submission must not synchronize") + + with torch.cuda.stream(stream): + torch.cuda._sleep(2_000_000) + manager.req_to_token_indexs.fill_(-1) + with monkeypatch.context() as no_sync: + no_sync.setattr(torch.cuda.Stream, "synchronize", unexpected_sync) + no_sync.setattr(torch.cuda, "synchronize", unexpected_sync) + cuda_memcpy_pinned_h2d_batch_async(dsts, srcs) + result = manager.req_to_token_indexs.clone() + stream.synchronize() + + assert torch.equal(result.cpu(), expected) + assert [size for size, _ in loop_calls] == [16] * 3 + assert all(copy_stream == stream.cuda_stream for _, copy_stream in loop_calls) + assert not base._cuda_memcpy_batch_available + attempts_after_failure = len(batch_attempts) + manager.req_to_token_indexs_cpu.fill_(99) + with torch.cuda.stream(stream): + cuda_memcpy_pinned_h2d_batch_async(dsts, srcs) + stream.synchronize() + assert len(batch_attempts) == attempts_after_failure + assert len(loop_calls) == 6 + assert manager.req_to_token_indexs[0, 4:8].cpu().tolist() == [99] * 4 + + +@pytest.mark.parametrize("entry", ["forward", "microbatch_overlap_prefill", "microbatch_overlap_decode"]) +def test_model_entries_do_not_manage_request_upload_progress(entry): + model = TpPartBaseModel.__new__(TpPartBaseModel) + events = [] + model.req_manager = SimpleNamespace() + model.args = SimpleNamespace(enable_prefill_decode_mixed=False) + model._prefill = lambda model_input: events.append("infer") + model._decode = lambda model_input: events.append("infer") + model._microbatch_overlap_prefill_cuda = lambda *inputs: events.append("infer") + model._microbatch_overlap_decode_cuda = lambda *inputs: events.append("infer") + model_input = SimpleNamespace( + is_prefill=True, + input_ids=torch.ones(1, dtype=torch.int64), + to_cuda=lambda: events.append("input"), + ) + + if entry == "forward": + model.forward(model_input) + assert events == ["input", "infer"] + else: + getattr(model, entry)(model_input, model_input) + assert events == ["input", "input", "infer"] + + +def _make_scheduler(monkeypatch, manager, reqs, budget=16, mixed=False): + from lightllm.server.router.model_infer.mode_backend import base_backend + + args = SimpleNamespace(page_size=4, enable_cpu_cache=False, enable_prefill_decode_mixed=mixed, run_mode="normal") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + context = SimpleNamespace( + req_manager=manager, + radix_cache=None, + get_can_alloc_token_num=lambda: 32, + get_overlap_stream=lambda: stream, + 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) + manager.mem_manager = SimpleNamespace(alloc=lambda size: torch.arange(100, 100 + size, dtype=torch.int32)) + backend = ModeBackend.__new__(ModeBackend) + backend.args = args + backend.support_overlap = False + backend.disable_chunked_prefill = False + backend.batch_max_tokens = budget + backend.is_master_in_dp = True + backend._timer_merge_radix_tree = lambda: None + backend._filter_not_ready_reqs = lambda req_ids: reqs + backend._reorder_pd_high_priority_reqs = lambda reqs: reqs + backend._reorder_long_prefill_reqs = lambda reqs: reqs + for req in reqs: + req.args = args + copies = [] + copy_indexes = backend._copy_req_token_indexes_to_gpu + + def record_copy(run_reqs): + ranges = [ + (req.req_idx, req.move_to_gpu_index_len, req.hold_kv_len) + for req in run_reqs + if req.move_to_gpu_index_len < req.hold_kv_len + ] + if ranges: + copies.append(ranges) + copy_indexes(run_reqs) + + backend._copy_req_token_indexes_to_gpu = record_copy + return backend, stream, copies + + +def _make_scheduled_req(row, total_len=12, moved=4): + req = InferReq.__new__(InferReq) + req.req_idx = req.req_id = row + req.cur_kv_len = 3 + req.hold_kv_len = 4 + req.move_to_gpu_index_len = moved + req.mtp_step = 0 + req.filter_mark = req.wait_pause = req.paused = req.infer_aborted = False + req.finish_status = infer_batch.FinishStatus() + req.shm_req = SimpleNamespace(input_len=total_len) + req.is_slave_req = lambda: False + req.get_cur_total_len = lambda: total_len + req.get_chuncked_input_token_len = lambda: min(total_len, req.cur_kv_len + 3) + return req + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_scheduler_submits_native_batch_on_inference_stream_from_default_stream(monkeypatch): + from cuda.bindings import runtime as cuda_runtime + + manager = _make_manager("cuda") + req = _make_scheduled_req(0, moved=0) + backend, stream, _ = _make_scheduler(monkeypatch, manager, [req]) + submitted_streams = [] + native_copy = cuda_runtime.cudaMemcpyBatchAsync + + def record_batch(*args): + submitted_streams.append(args[-1]) + return native_copy(*args) + + monkeypatch.setattr(cuda_runtime, "cudaMemcpyBatchAsync", record_batch) + with torch.cuda.stream(torch.cuda.default_stream()): + backend._get_classed_reqs(req_ids=[0]) + assert torch.cuda.current_stream().cuda_stream == 0 + stream.synchronize() + + assert submitted_streams == [stream.cuda_stream] + assert req.move_to_gpu_index_len == req.hold_kv_len + assert torch.equal( + manager.req_to_token_indexs[0, : req.hold_kv_len].cpu(), manager.req_to_token_indexs_cpu[0, : req.hold_kv_len] + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_scheduler_uploads_only_new_tail_and_reuses_uploaded_page(monkeypatch): + manager = _make_manager("cuda") + manager.req_to_token_indexs_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) + manager.req_to_token_indexs[0, :4].copy_(manager.req_to_token_indexs_cpu[0, :4]) + req = _make_scheduled_req(0) + backend, stream, copies = _make_scheduler(monkeypatch, manager, [req]) + + with torch.cuda.stream(stream): + prefill, decode = backend._get_classed_reqs(req_ids=[0]) + stream.synchronize() + assert prefill == [req] and decode == [] + assert copies == [[(0, 4, 8)]] + assert req.move_to_gpu_index_len == req.hold_kv_len == 8 + assert manager.req_to_token_indexs[0, :8].cpu().tolist() == [0, 1, 2, 3, 100, 101, 102, 103] + + req.cur_kv_len = 5 # The next chunk still fits into the reserved page. + with torch.cuda.stream(stream): + backend._get_classed_reqs(req_ids=[0]) + assert copies == [[(0, 4, 8)]] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_scheduler_leaves_budget_skipped_request_unpublished(monkeypatch): + manager = _make_manager("cuda") + manager.req_to_token_indexs_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) + req = _make_scheduled_req(0, moved=0) + req.hold_kv_len = 8 + backend, stream, copies = _make_scheduler(monkeypatch, manager, [req], budget=2) + + assert backend._get_classed_reqs(req_ids=[0]) == ([], []) + assert req.move_to_gpu_index_len == 0 + assert copies == [] + assert manager.req_to_token_indexs.eq(-1).all() + + backend.batch_max_tokens = 3 + with torch.cuda.stream(stream): + backend._get_classed_reqs(req_ids=[0]) + stream.synchronize() + assert copies == [[(0, 0, 8)]] + assert req.move_to_gpu_index_len == 8 + + +def test_pd_transfer_failure_truncates_uploaded_length(monkeypatch): + from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl + + manager = _make_manager() + manager.req_to_token_indexs_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) + req = SimpleNamespace( + req_idx=0, + req_id=0, + infer_aborted=False, + pd_task_num=1, + pd_task_failed_num=1, + pd_task_success_num=0, + finish_status=infer_batch.FinishStatus(infer_batch.FinishStatus.FINISHED_ERROR), + shared_kv_node=SimpleNamespace(node_prefix_total_len=4), + cur_kv_len=6, + hold_kv_len=8, + move_to_gpu_index_len=8, + ) + monkeypatch.setattr(decode_impl, "g_infer_context", SimpleNamespace(requests_mapping={0: req})) + freed = [] + backend = decode_impl.PDDecodeNode.__new__(decode_impl.PDDecodeNode) + backend.is_master_in_dp = False + backend.model = SimpleNamespace( + req_manager=manager, mem_manager=SimpleNamespace(free=lambda indexes: freed.extend(indexes.tolist())) + ) + + assert backend._filter_not_ready_reqs([0]) == [req] + assert freed == [4, 5, 6, 7] + assert req.cur_kv_len == req.hold_kv_len == req.move_to_gpu_index_len == 4 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("mixed", [False, True]) +def test_scheduler_publishes_prefill_and_decode_before_return(monkeypatch, mixed): + manager = _make_manager("cuda") + for row in (0, 1): + manager.req_to_token_indexs_cpu[row, :4].copy_(torch.arange(row * 4, row * 4 + 4, dtype=torch.int32)) + prefill_req = _make_scheduled_req(0, moved=0) + decode_req = _make_scheduled_req(1, total_len=4, moved=0) + backend, stream, copies = _make_scheduler(monkeypatch, manager, [prefill_req, decode_req], mixed=mixed) + + with torch.cuda.stream(stream): + prefill, decode = backend._get_classed_reqs(req_ids=[0, 1]) + stream.synchronize() + + assert copies == [[(0, 0, 8), (1, 0, 4)]] + assert prefill_req.move_to_gpu_index_len == 8 + assert decode_req.move_to_gpu_index_len == 4 + assert prefill == ([prefill_req, decode_req] if mixed else [prefill_req]) + assert decode == ([] if mixed else [decode_req]) + for req in (prefill_req, decode_req): + assert torch.equal( + manager.req_to_token_indexs[req.req_idx, : req.hold_kv_len].cpu(), + manager.req_to_token_indexs_cpu[req.req_idx, : req.hold_kv_len], + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("loop_fails", [False, True]) +def test_request_upload_lengths_follow_fallback_copy_result(monkeypatch, loop_fails): + from cuda.bindings import runtime as cuda_runtime + + manager = _make_manager("cuda") + reqs = [_make_scheduled_req(0, moved=0), _make_scheduled_req(1, total_len=4, moved=0)] + backend, stream, _ = _make_scheduler(monkeypatch, manager, reqs) + monkeypatch.setattr( + cuda_runtime, + "cudaMemcpyBatchAsync", + lambda *args: (cuda_runtime.cudaError_t.cudaErrorInvalidValue,), + ) + + if loop_fails: + torch_copy = torch.Tensor.copy_ + + def failed_copy(dst, src, *args, **kwargs): + if dst.is_cuda and src.device.type == "cpu": + raise RuntimeError("individual copy failed") + return torch_copy(dst, src, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "copy_", failed_copy) + with torch.cuda.stream(stream), pytest.raises(RuntimeError, match="individual copy failed"): + backend._get_classed_reqs(req_ids=[0, 1]) + assert [req.move_to_gpu_index_len for req in reqs] == [0, 0] + else: + with torch.cuda.stream(stream): + backend._get_classed_reqs(req_ids=[0, 1]) + stream.synchronize() + for req in reqs: + assert req.move_to_gpu_index_len == req.hold_kv_len + assert torch.equal( + manager.req_to_token_indexs[req.req_idx, : req.hold_kv_len].cpu(), + manager.req_to_token_indexs_cpu[req.req_idx, : req.hold_kv_len], + ) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +@pytest.mark.parametrize("batch_available", [False, True]) +def test_copy_uses_current_stream_on_current_device(monkeypatch, batch_available): + from lightllm.common import cuda_wrapper as base + + monkeypatch.setattr(base, "_cuda_memcpy_batch_available", batch_available) + manager = _make_manager("cuda:1") + manager.req_to_token_indexs_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) + stream = torch.cuda.Stream(device=1) + stream.wait_stream(torch.cuda.current_stream(1)) + with torch.cuda.device(0): + with torch.cuda.stream(stream): + cuda_memcpy_pinned_h2d_batch_async( + [manager.req_to_token_indexs[0, :4]], + [manager.req_to_token_indexs_cpu[0, :4]], + ) + assert torch.cuda.current_device() == 1 + assert torch.cuda.current_stream().cuda_stream == stream.cuda_stream + assert torch.cuda.current_device() == 0 + stream.synchronize() + assert manager.req_to_token_indexs[0, :4].cpu().tolist() == list(range(4)) 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..ea057a481d --- /dev/null +++ b/unit_tests/common/test_req_manager_page.py @@ -0,0 +1,514 @@ +from types import MethodType, SimpleNamespace + +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend import generic_pre_process +from lightllm.common.req_manager import ReqManager +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_cpu = torch.full((3, 8), -1, dtype=torch.int32) + 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_cpu[2].tolist() == [100, 101, 102, 103] * 2 + assert req_manager.req_to_token_indexs_cpu[: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 = ReqManager.__new__(ReqManager) + req_manager.mem_manager = mem_manager + req_manager.max_request_num = 2 + req_manager.req_to_token_indexs_cpu = torch.full((2, 16), -1, dtype=torch.int32) + req_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, + get_overlap_stream=lambda: 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 + + def copy_indexes(run_reqs): + for req in run_reqs: + start, end = req.move_to_gpu_index_len, req.hold_kv_len + req_manager.req_to_token_indexs[req.req_idx, start:end].copy_( + req_manager.req_to_token_indexs_cpu[req.req_idx, start:end] + ) + req.move_to_gpu_index_len = end + + backend._copy_req_token_indexes_to_gpu = copy_indexes + return context, backend + + +def _make_req(req_idx): + req = SimpleNamespace(req_idx=req_idx, cur_kv_len=0, hold_kv_len=0, move_to_gpu_index_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_cpu[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_cpu[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_cpu[0, :4].tolist() == [0, 1, 2, 3] + assert generic_pre_process.g_infer_context.req_manager.req_to_token_indexs_cpu[1, :4].tolist() == [4, 5, 6, 7] + assert req0.hold_kv_len == req1.hold_kv_len == 4 + + +def test_alloc_req_kv_mem_defers_gpu_table_update(monkeypatch): + context, backend = _make_context(monkeypatch) + req = _make_req(0) + + backend._alloc_req_kv_mem(req, alloc_token_num=4) + backend._alloc_req_kv_mem(req, alloc_token_num=4) + + assert context.req_manager.req_to_token_indexs_cpu[0, :8].tolist() == list(range(8)) + assert context.req_manager.req_to_token_indexs.eq(-1).all() + assert req.move_to_gpu_index_len == 0 + + +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._reorder_pd_high_priority_reqs = lambda reqs: reqs + backend._reorder_long_prefill_reqs = lambda reqs: 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_cpu[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 + prefill_reqs, decode_reqs = backend._get_classed_reqs(req_ids=[0, 1]) + + assert prefill_reqs == [reqs[0]] + assert decode_reqs == [] + 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 + + +def test_decode_scheduler_reserves_cpu_request_table(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._reorder_pd_high_priority_reqs = lambda reqs: reqs + backend._reorder_long_prefill_reqs = lambda reqs: 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_cpu[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] + + prefill_reqs, decode_reqs = backend._get_classed_reqs(req_ids=[0]) + + assert prefill_reqs == [] + assert decode_reqs == [req] + assert req.hold_kv_len == 16 + assert context.req_manager.mem_manager.alloc_sizes == [12] + + +@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_cpu = torch.full((2, 64), -1, dtype=torch.int32) + 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._reorder_pd_high_priority_reqs = lambda reqs: reqs + backend._reorder_long_prefill_reqs = lambda reqs: 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_cpu[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_cpu[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_cpu[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_cpu[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_cpu[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_cpu=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_cpu=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_cpu=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, + move_to_gpu_index_len=0, + 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_cpu=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) + + 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_cpu=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) + + 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 27484e1b85..8e2085fa00 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, monkeypatch): +def test_parallel_block_decode_commits_target_hiddens_directly(model_class): model = model_class.__new__(model_class) model._cos_cached = torch.arange(24).view(6, 4) model._sin_cached = model._cos_cached + 100 @@ -70,14 +70,7 @@ def test_parallel_block_decode_commits_target_hiddens_directly(model_class, monk target_hiddens = torch.arange(6, dtype=torch.float32).view(2, 3) mem_indexes = torch.tensor([7, 11]) - 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) + model._select_mem_indexes = lambda _: mem_indexes 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 823ca9d4fd..3cae97ac51 100644 --- a/unit_tests/server/core/objs/test_req.py +++ b/unit_tests/server/core/objs/test_req.py @@ -1,6 +1,7 @@ import pytest import easydict -from lightllm.server.core.objs.req import Req +from types import SimpleNamespace +from lightllm.server.core.objs.req import Req, ChunkedPrefillReq, SamplingParams 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 @@ -79,6 +80,17 @@ 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 fee9fb621b..a75af88847 100644 --- a/unit_tests/server/httpserver/test_running_request_lifecycle.py +++ b/unit_tests/server/httpserver/test_running_request_lifecycle.py @@ -418,3 +418,12 @@ 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 8a7f016400..4c2bee94dd 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -270,5 +270,53 @@ def test_case10(): assert tree.root_node.ref_counter == 1 +def test_page_aligned_insert_and_match(): + tree = RadixCache(100, 99, page_size=4) + values = torch.arange(100, 110, dtype=torch.int64) + + prefix_len, _ = tree.insert(torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), values) + assert prefix_len == 0 + assert tree.get_tree_total_tokens_num() == 8 + + # A mismatch inside a page cannot produce a partial-page cache hit. + node, matched_len, matched_values = tree.match_prefix(torch.tensor([1, 2, 0, 4, 5, 6, 7, 8])) + assert node is None + assert matched_len == 0 + assert matched_values is None + + node, matched_len, matched_values = tree.match_prefix(torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 11])) + assert node is not None + assert matched_len == 8 + assert matched_values.tolist() == list(range(100, 108)) + + # The second sequence shares exactly one page and then branches by its + # complete second-page key. + prefix_len, _ = tree.insert( + torch.tensor([1, 2, 3, 4, 5, 6, 0, 8]), + torch.arange(200, 208, dtype=torch.int64), + ) + assert prefix_len == 4 + assert tree.get_tree_total_tokens_num() == 12 + + +def test_page_size_must_match_mem_manager(): + mem_manager = SimpleNamespace(page_size=8) + + with pytest.raises(ValueError, match="must match mem_manager page_size 8"): + RadixCache(100, 100, mem_manager=mem_manager, page_size=4) + + +def test_page_key_bytes_does_not_share_tensor_memory(): + token_ids = torch.tensor([1, 2, 3, 4], dtype=torch.int64) + expected_key = token_ids.clone().numpy().tobytes() + key = TreeNode(page_size=4).get_child_key(token_ids) + + token_ids[0] = 100 + + assert isinstance(key, bytes) + assert key == expected_key + assert key != TreeNode(page_size=4).get_child_key(token_ids) + + if __name__ == "__main__": pytest.main() 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 new file mode 100644 index 0000000000..2f58371e78 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_diverse_backend.py @@ -0,0 +1,78 @@ +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_cpu = torch.full((2, 8), -1, dtype=torch.int32) + req_to_token_indexs_cpu[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, + move_to_gpu_index_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_cpu=req_to_token_indexs_cpu, + ), + ) + backend.radix_cache = _FakeRadixCache() + backend.is_master_in_dp = False + + backend._copy_master_req_to_slave_req(slave_req) + + assert req_to_token_indexs_cpu[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 slave_req.move_to_gpu_index_len == 0 + 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 574f17a938..42519a7415 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,8 +1,31 @@ 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( @@ -30,6 +53,8 @@ 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 == [] @@ -43,6 +68,8 @@ 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 == [] @@ -57,6 +84,7 @@ 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), @@ -83,6 +111,7 @@ 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) ( @@ -103,6 +132,7 @@ 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)] ( @@ -125,6 +155,7 @@ 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 a546f6fcbf..b584f900dc 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,17 +42,15 @@ 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( - get_cpu_token_index_blocks=lambda row, start, end: torch.arange(10 + start, 10 + end).numpy() - ) + context.req_manager = SimpleNamespace(req_to_token_indexs_cpu=torch.tensor([[10, 11, 12, 13, 14, 15, 16]])) 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_index_blocks = [] + free_token_indexes = [] - context._free_req_mem_without_radix_insert(free_token_index_blocks, req) + context._free_req_mem_without_radix_insert(free_token_indexes, req) - assert free_token_index_blocks[0].tolist() == [12, 13, 14, 15, 16] + assert free_token_indexes[0].tolist() == [12, 13, 14, 15, 16] assert released_refs == [shared_node] assert req.shared_kv_node is None @@ -62,7 +60,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_blocks, req: inserted_reqs.append(req) + context._full_att_free_req = lambda free_token_index, req: inserted_reqs.append(req) req = SimpleNamespace( cache_tiers=(CacheTier.GPU, CacheTier.CPU, CacheTier.DISK), cur_kv_len=3, @@ -138,9 +136,7 @@ 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( - get_cpu_token_index_blocks=lambda row, start, end: torch.arange(10 + start, 10 + end).numpy() - ) + context.req_manager = SimpleNamespace(req_to_token_indexs_cpu=torch.tensor([[10, 11, 12]])) 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), @@ -153,12 +149,92 @@ 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_index_blocks = [] + free_token_indexes = [] - context._free_req_mem_without_radix_insert(free_token_index_blocks, req) + context._free_req_mem_without_radix_insert(free_token_indexes, req) - assert free_token_index_blocks[0].tolist() == [10, 11, 12] + assert free_token_indexes[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 assert req.hybrid_len_to_big_page_id == {} + + +def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch): + table = torch.full((1, 136), -1, dtype=torch.int32) + table[0, :4] = torch.arange(4, dtype=torch.int32) + next_index = 4 + alloc_sizes = [] + loaded_indexes = [] + + def alloc(need_size): + nonlocal next_index + alloc_sizes.append(need_size) + mem_indexes = torch.arange(next_index, next_index + need_size, dtype=torch.int32) + next_index += need_size + return mem_indexes + + def alloc_req_kv_mem(req, alloc_token_num): + 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 + req.hold_kv_len = new_hold_kv_len + return mem_indexes + + operator = SimpleNamespace( + load_cpu_cache_to_gpu=lambda mem_indexes, **kwargs: loaded_indexes.extend(mem_indexes.tolist()) + ) + backend = SimpleNamespace( + is_master_in_dp=False, + radix_cache=None, + model=SimpleNamespace( + mem_manager=SimpleNamespace(operator=operator), + req_manager=SimpleNamespace(req_to_token_indexs_cpu=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) + module.backend = backend + module.init_sync_group = object() + module.need_sync_compute_stream = lambda: False + module.cpu_cache_client = SimpleNamespace() + context = SimpleNamespace( + req_manager=SimpleNamespace(req_to_token_indexs_cpu=table, mem_manager=SimpleNamespace(alloc=alloc)), + get_can_alloc_token_num=lambda: 200, + ) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=4, + hold_kv_len=4, + shm_req=SimpleNamespace( + input_len=256, + disk_prompt_cache_len=0, + cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: [7]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [132]), + ), + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + ) + monkeypatch.setattr(multi_level_kv_cache_impl, "g_infer_context", context) + monkeypatch.setattr(multi_level_kv_cache_impl.dist, "barrier", lambda group: None) + uploaded_indexes = [] + + def upload(tensor, **kwargs): + uploaded_indexes.append(tensor) + return tensor + + def unexpected_copy(*args, **kwargs): + pytest.fail("CPU mirror upload must not re-pin indexes or concatenate GPU tensors") + + monkeypatch.setattr(torch.Tensor, "cuda", upload) + monkeypatch.setattr(torch.Tensor, "pin_memory", unexpected_copy) + monkeypatch.setattr(torch, "cat", unexpected_copy) + + module.load_cpu_cache_to_reqs([req]) + + assert len(uploaded_indexes) == 2 # 页面编号和完整索引区间各上传一次。 + assert alloc_sizes == [128] + assert req.cur_kv_len == 132 + assert req.hold_kv_len == 132 + assert loaded_indexes == list(range(132)) + assert table[0, :132].tolist() == list(range(132)) 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 new file mode 100644 index 0000000000..102af90f44 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py @@ -0,0 +1,145 @@ +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_cpu=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_cpu=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_cpu=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 24ca0c9ecd..73e133cc18 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,7 +3,6 @@ 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 @@ -25,14 +24,13 @@ 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), + decode_need_token_num=MagicMock(return_value=(1, 1)), ) @@ -54,11 +52,7 @@ def _classify_without_token_capacity(monkeypatch, req, support_overlap=True): backend._reorder_pd_high_priority_reqs = MagicMock(side_effect=lambda reqs: reqs) backend._reorder_long_prefill_reqs = MagicMock(side_effect=lambda reqs: 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 new file mode 100644 index 0000000000..9e44ed774b --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py @@ -0,0 +1,52 @@ +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 new file mode 100644 index 0000000000..dff0e892a9 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py @@ -0,0 +1,33 @@ +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/req_queue/test_pd_queue_selection.py b/unit_tests/server/router/req_queue/test_pd_queue_selection.py new file mode 100644 index 0000000000..7c8cd5cbf0 --- /dev/null +++ b/unit_tests/server/router/req_queue/test_pd_queue_selection.py @@ -0,0 +1,109 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +from lightllm.server.router.req_queue import _get_req_queue_class +from lightllm.server.router.batch import Batch +from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_decode import PDDecodeQueue +from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_prefill import PDPrefillQueue + + +def _make_args(run_mode: str): + return SimpleNamespace( + run_mode=run_mode, + diverse_mode=False, + output_constraint_mode="none", + first_token_constraint_mode=False, + disable_chunked_prefill=False, + ) + + +def test_pd_prefill_uses_prefill_queue(): + args = _make_args("prefill") + + assert _get_req_queue_class(args, router=None, dp_size_in_node=1) is PDPrefillQueue + + +def test_pd_decode_uses_decode_queue(): + args = _make_args("decode") + + assert _get_req_queue_class(args, router=None, dp_size_in_node=1) is PDDecodeQueue + + +def test_pd_prefill_peak_tokens_sum_request_tuples(): + queue = PDPrefillQueue.__new__(PDPrefillQueue) + queue.dp_index = 0 + queue.args = SimpleNamespace(page_size=16) + queue.is_busy = lambda: False + queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) + reqs = [ + SimpleNamespace( + request_id="req-0", + input_len=10, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), + get_tuple_tokens=MagicMock(return_value=(11, 34)), + ), + SimpleNamespace( + request_id="req-1", + input_len=20, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), + get_tuple_tokens=MagicMock(return_value=(21, 34)), + ), + SimpleNamespace( + request_id="req-2", + input_len=100, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=1), + get_tuple_tokens=MagicMock(return_value=(101, 34)), + ), + ] + batch = Batch(batch_id=1, reqs=reqs, dp_size_in_node=2) + + assert queue._caclu_batch_estimated_peak_token_num(batch) == 100 + reqs[0].get_tuple_tokens.assert_called_once_with(False, 10) + reqs[1].get_tuple_tokens.assert_called_once_with(False, 10) + reqs[2].get_tuple_tokens.assert_not_called() + + +def test_pd_decode_sums_request_tuples_for_non_decode_requests(): + queue = PDDecodeQueue.__new__(PDDecodeQueue) + queue.args = SimpleNamespace(page_size=16) + queue.dp_index = 0 + queue.is_busy = lambda: False + queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) + reqs = [ + SimpleNamespace( + request_id="req-0", + input_len=20, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), + is_infer_decode=lambda: False, + get_tuple_tokens=MagicMock(return_value=(21, 34)), + ), + SimpleNamespace( + request_id="req-1", + input_len=100, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=1), + is_infer_decode=lambda: False, + get_tuple_tokens=MagicMock(return_value=(101, 34)), + ), + ] + batch = Batch(batch_id=1, reqs=reqs, dp_size_in_node=2) + + assert queue._caclu_batch_estimated_peak_token_num(batch) == 55 + reqs[0].get_tuple_tokens.assert_called_once_with(False, 10) + reqs[1].get_tuple_tokens.assert_not_called() + + +def test_pd_decode_uses_tuple_estimation(): + queue = PDDecodeQueue.__new__(PDDecodeQueue) + queue.args = SimpleNamespace(page_size=16) + queue.dp_index = 0 + queue.is_busy = lambda: False + queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) + req = SimpleNamespace( + request_id="req-0", + sample_params=SimpleNamespace(suggested_dp_index=0), + is_infer_decode=lambda: True, + get_tuple_tokens=lambda is_busy, ema_req_out_len: (16, 32), + ) + batch = Batch(batch_id=1, reqs=[req], dp_size_in_node=1) + + assert queue._caclu_batch_estimated_peak_token_num(batch) == 48 diff --git a/unit_tests/server/test_dp_request_capacity.py b/unit_tests/server/test_dp_request_capacity.py index aea1372403..f16001f8c7 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,expected", [("full_attention", 4), ("diverse", 32), ("cache_fetch", 32)]) -def test_capacity_depends_on_mode(monkeypatch, mode, expected): +@pytest.mark.parametrize("mode", ["full_attention", "diverse", "cache_fetch"]) +def test_modes_requiring_existing_capacity_keep_global_default(monkeypatch, mode): 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_capacity_depends_on_mode(monkeypatch, mode, expected): diverse_mode=mode == "diverse", enable_dp_prompt_cache_fetch=mode == "cache_fetch", ) - assert get_running_max_req_size_per_dp(args) == expected + assert get_running_max_req_size_per_dp(args) == 32 @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 new file mode 100644 index 0000000000..def9608eb8 --- /dev/null +++ b/unit_tests/server/test_pd_start_args.py @@ -0,0 +1,75 @@ +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 new file mode 100644 index 0000000000..9b30f4c65f --- /dev/null +++ b/unit_tests/test_static_benchmark_paged_kv.py @@ -0,0 +1,158 @@ +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)) + +import static_benchmark + +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(monkeypatch, page_size: int = 4): + mem_manager = FakeMemManager(page_size=page_size) + req_manager = SimpleNamespace( + mem_manager=mem_manager, + req_to_token_indexs_cpu=torch.full((4, 32), -1, dtype=torch.int32), + 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(monkeypatch): + executor, mem_manager, req_manager = make_executor(monkeypatch, 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 torch.equal(req_manager.req_to_token_indexs, req_manager.req_to_token_indexs_cpu) + 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_cpu[0, :4], page_size=4) + assert_complete_physical_pages(req_manager.req_to_token_indexs_cpu[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_cpu[0, :8], page_size=4) + assert_complete_physical_pages(req_manager.req_to_token_indexs_cpu[1, :8], page_size=4) + + +def test_mtp_expanded_rows_allocate_once_per_request(monkeypatch): + executor, mem_manager, req_manager = make_executor(monkeypatch, 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_cpu[0, :8], page_size=4) + assert_complete_physical_pages(req_manager.req_to_token_indexs_cpu[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 new file mode 100644 index 0000000000..116e299e24 --- /dev/null +++ b/unit_tests/utils/test_config_utils_fixed_kv.py @@ -0,0 +1,33 @@ +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 8a7e82fe83..f7c690bbe3 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_180_seconds(monkeypatch): +def test_pd_cache_high_priority_max_age_defaults_to_36_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() == 180 + assert get_pd_cache_high_priority_max_age_seconds() == 36 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_2048(monkeypatch): +def test_pd_cache_high_priority_min_prompt_tokens_defaults_to_4096(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() == 2048 + assert get_pd_cache_high_priority_min_prompt_tokens() == 4096 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_20_seconds(monkeypatch): +def test_pd_node_resource_wait_timeout_defaults_to_10_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() == 20 + assert get_pd_node_resource_wait_timeout_seconds() == 10 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 85153f43b3..7530c0c4d2 100644 --- a/unit_tests/utils/test_sgl_utils.py +++ b/unit_tests/utils/test_sgl_utils.py @@ -107,7 +107,7 @@ def test_fa3_decode_preserves_actual_length_before_graph_capture(monkeypatch, pa state = fa3_module.Fa3DecodeAttState( backend=SimpleNamespace( model=model, - infer_page_size=page_size, + page_size=page_size, uses_causal_attention=lambda: True, uses_dynamic_spec_verify_layout=lambda: False, get_page_table_view=lambda att_batch_size, max_kv_len, microbatch_index: torch.zeros( diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py index e3898aa801..f4f85d4b04 100644 --- a/unit_tests/utils/test_speculative_utils.py +++ b/unit_tests/utils/test_speculative_utils.py @@ -125,7 +125,7 @@ def test_fa3_prefill_state_owns_causality(monkeypatch, state_class): 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, infer_page_size=1), + backend=SimpleNamespace(uses_causal_attention=lambda: False, page_size=1, infer_page_size=1), infer_state=infer_state, ) From 20e1a79b0414f2519854156d93043fd0fae682e6 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 14 Oct 2026 07:14:29 +0000 Subject: [PATCH 04/22] style: revert unrelated blank-line changes Restore blank-line spacing from the PR base without changing code behavior. --- lightllm/common/basemodel/basemodel.py | 1 + lightllm/common/kv_cache_mem_manager/mem_manager.py | 1 + lightllm/server/router/model_infer/infer_batch.py | 1 + lightllm/server/router/model_infer/mode_backend/base_backend.py | 2 ++ .../router/model_infer/mode_backend/diverse_backend/impl.py | 1 + .../router/model_infer/mode_backend/multi_level_kv_cache.py | 1 + unit_tests/server/router/dynamic_prompt/test_radix_cache.py | 1 + unit_tests/test_static_benchmark_paged_kv.py | 1 - 8 files changed, 8 insertions(+), 1 deletion(-) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index b539be887c..68a2bf27e6 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -641,6 +641,7 @@ 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" diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 6c812ff21f..6f4718faab 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -29,6 +29,7 @@ class MemoryManager: + operator_class = NormalMemOperator def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 15ce5f2501..567823bfab 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -350,6 +350,7 @@ 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 = [] for req in pause_reqs: if self.args.diverse_mode: 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 33d62008e3..8ec1153b4a 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -746,6 +746,7 @@ def _get_classed_reqs( can_alloc_token_num = g_infer_context.get_can_alloc_token_num() for req_obj in ready_reqs: + if req_obj.filter_mark: finished_reqs.append(req_obj) continue @@ -982,6 +983,7 @@ def _sample_and_scatter_token( b_prefill_has_output_cpu: torch.Tensor = None, mask_func: Optional[Callable] = None, ): + if mask_func is not None: assert len(run_reqs) == logits.shape[0] mask_func(run_reqs, logits) 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 d5e6c2f210..c5090f467f 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 @@ -40,6 +40,7 @@ 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 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 0703a08f98..103b25d92b 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 @@ -225,6 +225,7 @@ def _start_kv_cache_offload_task( assert len(token_hash_list) == len(page_len_list) if self.backend.is_master_in_dp: + find_index = bisect.bisect_right(page_len_list, req.cur_kv_len) move_block_size = find_index 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 4c2bee94dd..c5d9623e62 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -92,6 +92,7 @@ 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")) diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py index 9b30f4c65f..99e99fb0f1 100644 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ b/unit_tests/test_static_benchmark_paged_kv.py @@ -43,7 +43,6 @@ def make_executor(monkeypatch, page_size: int = 4): req_to_token_indexs_cpu=torch.full((4, 32), -1, dtype=torch.int32), 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) From 9c7d44acfc7b9e4e549b96e1959f0f0e217d62b7 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 14 Oct 2026 08:58:26 +0000 Subject: [PATCH 05/22] perf: reuse request buffers and publish DP indexes on GPU Keep pinned CPU request indexes local and publish pending GPU indexes before DP cache transfers. Reuse GPU index tables for cache loading and offloading, and cache prompt ID buffers per CUDA stream to avoid unnecessary allocation and synchronization. Document the lifecycle and cross-stream risks of borrowing inference index tables. Update regression tests for stream isolation, async counting, DP index publication, and cache loading. Validation: 42 focused tests passed; git diff --check passed. --- docs/kv_cache_page_size.md | 5 +- .../kv_cache_mem_manager/mem_manager.py | 11 -- .../common/req_manager/req_sampling_params.py | 34 +++++- .../server/router/model_infer/infer_batch.py | 2 +- .../model_infer/mode_backend/base_backend.py | 1 - .../dp_backend/dp_shared_kv_trans.py | 14 ++- .../mode_backend/multi_level_kv_cache.py | 27 +++-- unit_tests/common/test_req_manager_cpu.py | 81 +++---------- unit_tests/common/test_req_sampling_params.py | 109 ++++++++++++++++++ .../mode_backend/test_multi_level_kv_cache.py | 15 ++- .../test_paged_side_allocations.py | 82 ++++++++++++- 11 files changed, 283 insertions(+), 98 deletions(-) create mode 100644 unit_tests/common/test_req_sampling_params.py diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index 8d45aeb1cf..fc4818e463 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -30,7 +30,10 @@ prefill/decode 请求的 成功后将长度更新到 `hold_kv_len`。`DuringApiCall` 源访问属性保证返回前读取完 CPU 数据， 后续计算依靠同流顺序读取 GPU 表；长度表示异步拷贝已排队，不表示设备已经完成拷贝。 暂停/释放时重置长度，前缀重映射时回退长度以重新上传。 -DP 通过共享 CPU 镜像读取其他 rank 的分配元数据，无需等待下一次推理。 +DP 在发布请求 KV 信息前，让专用传输流等待 overlap 流已有的索引更新， +再将各 TP rank 的待上传索引拷贝到 GPU 表； +当前流等待传输流，再由同步 NCCL barrier 保证上传完成。其他 DP rank 在 barrier 后 +通过 GPU 请求表读取分配元数据，不跨进程共享 CPU 镜像。 请求筛选、地址计算、批量上传及长度更新均在这个 backend 函数内完成；`ReqManager` 只持有映射表。 独立 benchmark 和模型初始化检查共用 `cuda_memcpy_batch_async` 底层调用。 diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 6f4718faab..38c870ce4e 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -1,7 +1,6 @@ import math import os import re -import weakref import torch import torch.distributed as dist import torch.multiprocessing as mp @@ -249,16 +248,6 @@ def write_to_shm(self, req_manager): # 这个地方是一个不太优雅的设计，但是暂时这么做，可以让dp shared kv swap模块直接访问 req_manager 中的 req_to_token_indexs # 避免过多无用的数据复制和传输开销。 self.req_to_token_indexs: torch.Tensor = req_manager.req_to_token_indexs - # DP transfers need allocation metadata even before the next model forward publishes it. - self.req_to_token_indexs_cpu: torch.Tensor = req_manager.req_to_token_indexs_cpu.share_memory_() - # share_memory_ 会更换存储并丢失 pinned 属性，注册共享存储以继续直接异步上传。 - if not self.req_to_token_indexs_cpu.is_pinned(): - runtime = torch.cuda.cudart() - cpu_table = self.req_to_token_indexs_cpu - error = runtime.cudaHostRegister(cpu_table.data_ptr(), cpu_table.numel() * cpu_table.element_size(), 0) - if error != runtime.cudaError.success: - raise RuntimeError(f"cudaHostRegister failed for shared request indexes: {error}") - weakref.finalize(cpu_table, runtime.cudaHostUnregister, cpu_table.data_ptr()) lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") with lock: diff --git a/lightllm/common/req_manager/req_sampling_params.py b/lightllm/common/req_manager/req_sampling_params.py index 549847f4d8..476569f36d 100644 --- a/lightllm/common/req_manager/req_sampling_params.py +++ b/lightllm/common/req_manager/req_sampling_params.py @@ -81,16 +81,23 @@ 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(), + # 每次独立申请 pinned 内存，避免缓冲池循环复用时覆盖尚未完成的异步上传源。 + # 上传源释放后由 PyTorch pinned allocator 跟踪拷贝完成再复用；上传和计数 + # 在同一流上按序执行，后续采样由推理入口的流依赖保证，无需在此阻塞 CPU。 + prompt_ids_numpy = req.shm_req.get_prompt_ids_numpy() + prompt_ids_cpu = torch.empty( + prompt_ids_numpy.shape, dtype=torch.int32, - ).cuda(non_blocking=True) + device="cpu", + pin_memory=True, + ) + prompt_ids_cpu.numpy()[:] = prompt_ids_numpy + # 当前流的 GPU 暂存区覆盖最大请求长度，仅上传实际 prompt 部分。 + prompt_ids = self._get_prompt_ids_gpu_buffer()[: prompt_ids_cpu.numel()] + prompt_ids.copy_(prompt_ids_cpu, non_blocking=True) token_id_counter( prompt_ids=prompt_ids, out_token_id_counter=self.req_to_out_token_id_counter[req.req_idx] ) - # TODO: 分析计数初始化后同步当前流的必要性，以及对请求接入和推理并行的影响。 - torch.cuda.current_stream().synchronize() return @@ -150,3 +157,18 @@ def gen_cpu_out_token_counter_sampling_params(self, req_objs: List["InferReq"]): p_token_counts_tensor.cuda(non_blocking=True), p_cumsum_seq_len_tensor.cuda(non_blocking=True), ) + + def _get_prompt_ids_gpu_buffer(self) -> torch.Tensor: + """获取当前 CUDA 流的 int32 暂存区，容量为最大请求长度加 64。""" + buffers = getattr(self, "_prompt_ids_gpu_buffers", None) + if buffers is None: + buffers = {} + self._prompt_ids_gpu_buffers = buffers + stream = torch.cuda.current_stream() + buffer = buffers.get(stream) + if buffer is None: + buffer = torch.empty(get_env_start_args().max_req_total_len + 64, dtype=torch.int32, device=stream.device) + buffers[stream] = buffer + # 每个流最多缓存一个 GPU buffer；使用者须在当前流上提交拷贝和读取， + # 下一次覆盖会排在此前读取之后，不同流使用独立 buffer，避免互相覆盖。 + return buffer diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 567823bfab..41df035fd5 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -459,7 +459,7 @@ def __init__( vocab_size: int, ) -> None: self.shm_param = shm_req.sample_params - self.disable_prompt_cache = self.shm_param.disable_prompt_cache or self.shm_param.prompt_logprobs >= 0 + self.disable_prompt_cache = self.shm_param.disable_prompt_cache if self.shm_param.top_k == -1: self.shm_param.top_k = vocab_size 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 8ec1153b4a..c0214be576 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -683,7 +683,6 @@ def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): return cuda_memcpy_pinned_h2d_batch_async(dst_tensors, src_tensors) - gpu_table.record_stream(torch.cuda.current_stream()) for req, end in pending_copies: req.move_to_gpu_index_len = end 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 f60c31190e..3c4200a161 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 @@ -23,6 +23,7 @@ def __init__(self, max_req_num: int, dp_size_in_node: int, backend): self.backend: DPChunkedPrefillBackend = backend self.max_req_num = max_req_num + self.copy_stream = torch.cuda.Stream() # 0 代表 kv_len, 1 代表 radix_cache_len self.shared_req_infos = ShmArray( @@ -38,6 +39,15 @@ def fill_reqs_info(self, reqs: List[InferReq]): """ 填充请求的 kv 信息到共享内存中 """ + # 所有 TP rank 先发布 CPU 索引，确保其他 DP rank 读取 GPU 表时数据已经就绪。 + current_stream = torch.cuda.current_stream() + # 常规索引上传使用 overlap stream，本次上传需排在此前的索引更新之后。 + self.copy_stream.wait_stream(g_infer_context.get_overlap_stream()) + with torch.cuda.stream(self.copy_stream): + self.backend._copy_req_token_indexes_to_gpu(reqs) + current_stream.wait_stream(self.copy_stream) + # 同步 NCCL barrier 会等待当前流完成，保证其他 rank 读取到已上传的索引。 + # TODO: 分析传输流依赖及 NCCL barrier 等待对 DP shared 并行效率的影响。 dist.barrier(group=self.backend.node_nccl_group) if self.backend.is_master_in_dp: self.shared_req_infos.arr[0 : len(reqs), self.dp_rank_in_node, self._KV_LEN_INDEX] = [ @@ -84,9 +94,9 @@ def build_shared_kv_trans_tasks( 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 max_kv_len_mem_manager: MemoryManager = self.backend.mem_managers[max_kv_len_mem_manager_index] - max_kv_len_mem_indexes = max_kv_len_mem_manager.req_to_token_indexs_cpu[ + max_kv_len_mem_indexes = max_kv_len_mem_manager.req_to_token_indexs[ max_kv_len_req_idx, req.cur_kv_len : req.cur_kv_len + trans_size - ].clone() + ] trans_tasks.append( TransTask( req=req, 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 103b25d92b..51c95d134e 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 @@ -103,7 +103,8 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): assert ready_page_num <= len(page_list) need_pages = page_list[ready_page_num:] # 只取需要的页面 - self.backend._alloc_req_kv_mem(req, need_token_num) + allocated_mem_indexes = self.backend._alloc_req_kv_mem(req, need_token_num) + assert allocated_mem_indexes is not None if self.need_sync_compute_stream(): # TODO fa3 现在必须使用同步模式, 未来需要移除 @@ -122,9 +123,14 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): _end = req.hold_kv_len assert 0 <= _start <= req.cur_kv_len <= _end, f"invalid pad range [{_start}, {_end}]" - # 旧前缀和新分配的索引已连续写入 pinned CPU 镜像，一次上传即可。 - mem_indexes_cuda = req_manager.req_to_token_indexs_cpu[req.req_idx, _start:_end].cuda( - non_blocking=True + # 借用已有 GPU 索引表，避免为加载临时分配索引显存；旧前缀和新分配的 + # 索引从 pinned CPU 镜像一次上传，copy_ 与 KV 加载在当前流上按序执行。 + # 风险：加载完成前须保持 CPU 索引不变、请求槽位有效，且不能有其他流 + # 冲突访问此区间；修改初始化或流调度时需重新检查依赖，不能只依赖写入相同值。 + # 此处不更新推理流的上传进度。 + mem_indexes_cuda = req_manager.req_to_token_indexs[req.req_idx, _start:_end] + mem_indexes_cuda.copy_( + req_manager.req_to_token_indexs_cpu[req.req_idx, _start:_end], non_blocking=True ) assert ( @@ -277,9 +283,16 @@ def _start_kv_cache_offload_task( move_token_num = page_len_list[item_size - 1] assert req.cur_kv_len >= move_token_num - token_indexes = self.backend.model.req_manager.req_to_token_indexs_cpu[req.req_idx, :move_token_num].cuda( - non_blocking=True - ) + req_manager = self.backend.model.req_manager + # 借用已有 GPU 索引表作为 offload 输入，避免 .cuda() 额外分配临时索引显存。 + # 风险警告：此处跨流借用推理索引表，“写入相同值”不能替代同步保证。 + # 安全性依赖结束请求的延迟清理，确保此前的推理访问已结束；CPU 索引前缀在 + # 上传期间必须保持不变，卸载完成事件确认前不得释放 KV 或复用 req_idx。 + # copy_ 与 offload 在同一卸载流上按序执行；这里只借用表，不更新推理流的上传进度。 + # 后续若用于运行中请求、前缀重映射，或修改 abort/清理/流调度逻辑，必须重新 + # 检查跨流依赖，必要时增加事件等待或改用卸载流独立暂存区，请谨慎修改。 + token_indexes = req_manager.req_to_token_indexs[req.req_idx, :move_token_num] + token_indexes.copy_(req_manager.req_to_token_indexs_cpu[req.req_idx, :move_token_num], non_blocking=True) mem_manager = self.backend.model.mem_manager diff --git a/unit_tests/common/test_req_manager_cpu.py b/unit_tests/common/test_req_manager_cpu.py index fd351e7d49..b8d20f8657 100644 --- a/unit_tests/common/test_req_manager_cpu.py +++ b/unit_tests/common/test_req_manager_cpu.py @@ -1,5 +1,4 @@ from types import SimpleNamespace -from multiprocessing.reduction import ForkingPickler from contextlib import nullcontext import builtins @@ -64,52 +63,27 @@ def test_releasing_request_memory_resets_uploaded_length(): assert req.cur_kv_len == req.hold_kv_len == req.move_to_gpu_index_len == 0 -def test_shared_cpu_mirror_exposes_updates_before_gpu_publication(): - manager = _make_manager() - source = manager.req_to_token_indexs_cpu.share_memory_() - remote = ForkingPickler.loads(ForkingPickler.dumps(source)) - manager.req_to_token_indexs_cpu[0, :2].copy_(torch.tensor([8, 9], dtype=torch.int32)) - - assert remote[0, :2].tolist() == [8, 9] - assert manager.req_to_token_indexs.eq(-1).all() - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("registration_fails", [False, True]) -def test_shared_export_preserves_pinned_mirror_and_async_upload(monkeypatch, registration_fails): +def test_shared_export_keeps_cpu_mirror_local_and_pinned(monkeypatch): from lightllm.common.kv_cache_mem_manager import mem_manager as impl req_manager = _make_manager("cuda") mem_manager = impl.MemoryManager.__new__(impl.MemoryManager) - runtime = torch.cuda.cudart() - registrations, cleanups, remote_tables = [], [], [] - - def register(ptr, size, flags): - registrations.append((ptr, size)) - return 1 if registration_fails else runtime.cudaHostRegister(ptr, size, flags) - - original_finalize = impl.weakref.finalize - - def finalize(obj, *args): - cleanup = original_finalize(obj, *args) - cleanups.append(cleanup) - return cleanup + source = req_manager.req_to_token_indexs_cpu + source_ptr = source.data_ptr() + exports = [] def serialize(manager): - assert manager.req_to_token_indexs_cpu.is_pinned() - remote_tables.append(ForkingPickler.loads(ForkingPickler.dumps(manager.req_to_token_indexs_cpu))) + assert manager.req_to_token_indexs is req_manager.req_to_token_indexs + assert not hasattr(manager, "req_to_token_indexs_cpu") + exports.append(manager) return memoryview(b"test") - monkeypatch.setattr( - torch.cuda, - "cudart", - lambda: SimpleNamespace( - cudaError=runtime.cudaError, - cudaHostRegister=register, - cudaHostUnregister=runtime.cudaHostUnregister, - ), - ) - monkeypatch.setattr(impl.weakref, "finalize", finalize) + def forbid_cpu_export(*args, **kwargs): + pytest.fail("CPU mirror export and host registration must not run") + + monkeypatch.setattr(torch.Tensor, "share_memory_", forbid_cpu_export) + monkeypatch.setattr(torch.cuda, "cudart", forbid_cpu_export) monkeypatch.setattr(impl, "ForkingPickler", SimpleNamespace(dumps=serialize)) monkeypatch.setattr(impl, "kv_trans_use_p2p", lambda: False) monkeypatch.setattr(impl, "FileLock", lambda path: nullcontext()) @@ -124,32 +98,13 @@ def serialize(manager): ), ) - try: - if registration_fails: - with pytest.raises(RuntimeError, match="cudaHostRegister failed"): - mem_manager.write_to_shm(req_manager) - assert not cleanups and not remote_tables - return + mem_manager.write_to_shm(req_manager) + mem_manager.write_to_shm(req_manager) - mem_manager.write_to_shm(req_manager) - mem_manager.write_to_shm(req_manager) - source = req_manager.req_to_token_indexs_cpu - assert mem_manager.req_to_token_indexs_cpu is source - assert source.is_shared() and source.is_pinned() - assert len(registrations) == len(cleanups) == 1 - source[0, :4].copy_(torch.arange(40, 44, dtype=torch.int32)) - assert all(table[0, :4].tolist() == [40, 41, 42, 43] for table in remote_tables) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - cuda_memcpy_pinned_h2d_batch_async([req_manager.req_to_token_indexs[0, :4]], [source[0, :4]]) - stream.synchronize() - assert req_manager.req_to_token_indexs[0, :4].cpu().tolist() == [40, 41, 42, 43] - finally: - torch.cuda.synchronize() - for cleanup in cleanups: - assert cleanup() == runtime.cudaError.success + assert len(exports) == 4 + assert req_manager.req_to_token_indexs_cpu is source + assert source.data_ptr() == source_ptr + assert source.is_pinned() and not source.is_shared() def test_prefix_match_invalidates_uploaded_length(monkeypatch): diff --git a/unit_tests/common/test_req_sampling_params.py b/unit_tests/common/test_req_sampling_params.py new file mode 100644 index 0000000000..c9b7688ca5 --- /dev/null +++ b/unit_tests/common/test_req_sampling_params.py @@ -0,0 +1,109 @@ +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.common.req_manager import req_sampling_params +from lightllm.common.req_manager.req_sampling_params import ReqSamplingParamsManager +from lightllm.common.basemodel.triton_kernel.gen_sampling_params import token_id_counter + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def test_prompt_ids_gpu_buffer_reuses_per_stream_and_isolates_async_copies(monkeypatch): + monkeypatch.setattr(req_sampling_params, "get_env_start_args", lambda: SimpleNamespace(max_req_total_len=32)) + manager = ReqSamplingParamsManager.__new__(ReqSamplingParamsManager) + stream0, stream1 = torch.cuda.Stream(), torch.cuda.Stream() + source0 = torch.full((96,), 11, dtype=torch.int32, pin_memory=True) + source1 = torch.full((96,), 22, dtype=torch.int32, pin_memory=True) + assert not hasattr(manager, "_prompt_ids_gpu_buffers") + with torch.cuda.stream(stream0): + first = manager._get_prompt_ids_gpu_buffer() + assert manager._get_prompt_ids_gpu_buffer() is first + assert first.shape == (96,) + assert first.dtype == torch.int32 and first.device == stream0.device + torch.cuda._sleep(100_000_000) + first.copy_(source0, non_blocking=True) + completed0 = torch.cuda.Event() + completed0.record() + with torch.cuda.stream(stream1): + other = manager._get_prompt_ids_gpu_buffer() + assert manager._get_prompt_ids_gpu_buffer() is other + assert other.data_ptr() != first.data_ptr() + assert other.shape == (96,) + torch.cuda._sleep(100_000_000) + other.copy_(source1, non_blocking=True) + completed1 = torch.cuda.Event() + completed1.record() + with torch.cuda.stream(stream0): + assert manager._get_prompt_ids_gpu_buffer() is first + completed0.synchronize() + completed1.synchronize() + assert torch.equal(first.cpu(), source0) + assert torch.equal(other.cpu(), source1) + assert len(manager._prompt_ids_gpu_buffers) == 2 + + +@pytest.mark.parametrize("mode", ["gpu_counter", "pin_mem_counter"]) +def test_prompt_counter_initialization_reuses_stream_buffer(monkeypatch, mode): + monkeypatch.setattr(req_sampling_params, "get_env_start_args", lambda: SimpleNamespace(max_req_total_len=8192)) + manager = ReqSamplingParamsManager.__new__(ReqSamplingParamsManager) + manager.penalty_counter_mode = mode + manager.req_to_next_token_ids = torch.zeros((8, 1), dtype=torch.int64, device="cuda") + manager.req_to_next_token_scores = None + for name in ( + "req_to_presence_penalty", + "req_to_frequency_penalty", + "req_to_repetition_penalty", + "req_to_exponential_decay_length_penalty", + ): + setattr(manager, name, torch.zeros(8, device="cuda")) + manager.req_to_out_token_id_counter = torch.zeros( + (8, 16), + dtype=torch.int32, + device="cuda" if mode == "gpu_counter" else "cpu", + pin_memory=mode == "pin_mem_counter", + ) + # 提前编译计数 kernel，随后延迟 GPU，检验连续复用暂存区时的异步执行顺序。 + token_id_counter(torch.zeros(1, dtype=torch.int32, device="cuda"), manager.req_to_out_token_id_counter[0]) + torch.cuda.synchronize() + + def forbid_synchronize(*args, **kwargs): + pytest.fail("request initialization must not synchronize its CUDA stream") + + monkeypatch.setattr(torch.cuda.Stream, "synchronize", forbid_synchronize) + stream = torch.cuda.Stream() + buffer_ptrs = [] + with torch.cuda.stream(stream): + torch.cuda._sleep(100_000_000) + for index in range(8): + params = SimpleNamespace( + presence_penalty=1.0, + frequency_penalty=0.0, + repetition_penalty=1.0, + input_penalty=True, + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (0, 1.0)), + ) + prompt = np.full(4096, index + 1, dtype=np.int64) + req = SimpleNamespace( + req_idx=index, + sampling_param=SimpleNamespace(shm_param=params), + shm_req=SimpleNamespace( + alloc_shm_numpy_len=8192, + get_prompt_ids_numpy=lambda prompt=prompt: prompt, + ), + get_last_gen_token=lambda: 0, + ) + manager.init_req_sampling_params(req) + buffer_ptrs.append(manager._get_prompt_ids_gpu_buffer().data_ptr()) + completed = torch.cuda.Event() + completed.record() + completed.synchronize() + expected = torch.zeros((8, 16), dtype=torch.int32) + expected[torch.arange(8), torch.arange(1, 9)] = 4096 + assert torch.equal(manager.req_to_out_token_id_counter.cpu(), expected) + assert len(set(buffer_ptrs)) == len(manager._prompt_ids_gpu_buffers) == 1 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 b584f900dc..2fe5cdd73e 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 @@ -162,6 +162,7 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch): table = torch.full((1, 136), -1, dtype=torch.int32) + gpu_table = torch.full_like(table, -2) table[0, :4] = torch.arange(4, dtype=torch.int32) next_index = 4 alloc_sizes = [] @@ -181,15 +182,17 @@ def alloc_req_kv_mem(req, alloc_token_num): req.hold_kv_len = new_hold_kv_len return mem_indexes - operator = SimpleNamespace( - load_cpu_cache_to_gpu=lambda mem_indexes, **kwargs: loaded_indexes.extend(mem_indexes.tolist()) - ) + def load(mem_indexes, **kwargs): + assert mem_indexes.data_ptr() == gpu_table[0, :132].data_ptr() + loaded_indexes.extend(mem_indexes.tolist()) + + operator = SimpleNamespace(load_cpu_cache_to_gpu=load) backend = SimpleNamespace( is_master_in_dp=False, radix_cache=None, model=SimpleNamespace( mem_manager=SimpleNamespace(operator=operator), - req_manager=SimpleNamespace(req_to_token_indexs_cpu=table), + req_manager=SimpleNamespace(req_to_token_indexs_cpu=table, req_to_token_indexs=gpu_table), ), _alloc_req_kv_mem=alloc_req_kv_mem, ) @@ -232,9 +235,11 @@ def unexpected_copy(*args, **kwargs): module.load_cpu_cache_to_reqs([req]) - assert len(uploaded_indexes) == 2 # 页面编号和完整索引区间各上传一次。 + assert len(uploaded_indexes) == 1 # 仅页面编号使用 .cuda()；索引上传到已有 GPU 表。 assert alloc_sizes == [128] assert req.cur_kv_len == 132 assert req.hold_kv_len == 132 assert loaded_indexes == list(range(132)) assert table[0, :132].tolist() == list(range(132)) + assert gpu_table[0, :132].tolist() == list(range(132)) + assert gpu_table[0, 132:].eq(-2).all() 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 index 102af90f44..6af1b1b854 100644 --- 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 @@ -8,6 +8,7 @@ # 在独立运行本测试文件时形成模块级循环导入。 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 import base_backend 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 @@ -101,6 +102,85 @@ def test_pd_prefill_rejects_unaligned_transfer_page_size(): backend._prefill_chuncked_handle_func(req, next_token_id=1, next_token_prob=1.0, output_len=0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("is_master_in_dp", [False, True]) +def test_dp_cache_fetch_publishes_gpu_indexes_before_barrier_returns(monkeypatch, is_master_in_dp): + cpu_table = torch.empty((2, 16), dtype=torch.int32, pin_memory=True) + torch.arange(32, out=cpu_table.view(-1)) + gpu_table = torch.full((2, 16), -1, dtype=torch.int32, device="cuda") + gpu_table[1, :2].copy_(cpu_table[1, :2], non_blocking=True) + req_manager = SimpleNamespace(req_to_token_indexs_cpu=cpu_table, req_to_token_indexs=gpu_table) + reqs = [ + SimpleNamespace(req_idx=1, cur_kv_len=4, hold_kv_len=4, move_to_gpu_index_len=2), + SimpleNamespace(req_idx=0, cur_kv_len=4, hold_kv_len=8, move_to_gpu_index_len=0), + ] + stream = torch.cuda.Stream() + overlap_stream = torch.cuda.Stream() + current_stream = torch.cuda.current_stream() + # 测试刚创建的 GPU 表先完成初始化；生产表在模型初始化期间已经就绪。 + overlap_stream.wait_stream(current_stream) + # 模拟 overlap 流上尚未完成的旧索引写入，验证 DP 上传会排在其后。 + with torch.cuda.stream(overlap_stream): + torch.cuda._sleep(20_000_000) + gpu_table[0, :8].fill_(-2) + context = SimpleNamespace(req_manager=req_manager, get_overlap_stream=lambda: overlap_stream) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) + monkeypatch.setattr(base_backend, "g_infer_context", context) + backend = SimpleNamespace(node_nccl_group=object(), is_master_in_dp=is_master_in_dp) + order = [] + + def upload(run_reqs): + assert torch.cuda.current_stream() == stream + order.append("upload") + base_backend.ModeBackend._copy_req_token_indexes_to_gpu(backend, run_reqs) + + original_synchronize = torch.cuda.Stream.synchronize + original_wait_stream = torch.cuda.Stream.wait_stream + + def forbid_synchronize(*args): + pytest.fail("DP shared must rely on NCCL barrier rather than explicitly synchronizing a stream") + + def wait_stream(waiting_stream, source_stream): + if waiting_stream == stream: + assert source_stream == overlap_stream + order.append("copy_waits_overlap") + else: + assert waiting_stream == current_stream and source_stream == stream + order.append("current_waits_copy") + original_wait_stream(waiting_stream, source_stream) + + def barrier(group): + assert group is backend.node_nccl_group + assert order == ["copy_waits_overlap", "upload", "current_waits_copy"] + # 模拟同步 NCCL barrier 在 C++ 层等待当前流完成。 + original_synchronize(current_stream) + assert torch.equal(gpu_table[1, :4].cpu(), cpu_table[1, :4]) + assert torch.equal(gpu_table[0, :8].cpu(), cpu_table[0, :8]) + assert gpu_table[0, 8:].eq(-1).all() + assert gpu_table[1, 4:].eq(-1).all() + order.append("barrier") + + backend._copy_req_token_indexes_to_gpu = upload + monkeypatch.setattr(torch.cuda.Stream, "synchronize", forbid_synchronize) + monkeypatch.setattr(torch.cuda.Stream, "wait_stream", wait_stream) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", barrier) + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = backend + module.copy_stream = stream + module.dp_rank_in_node = 0 + module.shared_req_infos = SimpleNamespace(arr=np.full((2, 2, 2), -1, dtype=np.int64)) + + module.fill_reqs_info(reqs) + + assert order == ["copy_waits_overlap", "upload", "current_waits_copy", "barrier"] + assert [req.move_to_gpu_index_len for req in reqs] == [4, 8] + if is_master_in_dp: + assert module.shared_req_infos.arr[:, 0].tolist() == [[4, 1], [4, 0]] + else: + assert (module.shared_req_infos.arr == -1).all() + assert cpu_table.is_pinned() and not cpu_table.is_shared() + + 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) @@ -121,7 +201,7 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat dp_world_size=1, rank_in_dp=0, node_nccl_group=None, - mem_managers=[SimpleNamespace(req_to_token_indexs_cpu=source_table)] * 2, + 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) From 4d50e730bf63367542482a297f05a907835e9c6f Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 14 Oct 2026 09:03:28 +0000 Subject: [PATCH 06/22] refactor: simplify request index allocation with deque Replace linked nodes with a deque while preserving allocation order, recently freed index reuse, capacity tracking, and duplicate-release checks. Validation: 51 request-manager tests passed; formatting and git diff --check passed. --- lightllm/common/req_manager/base.py | 29 +++++++--------------- unit_tests/common/test_req_manager_cpu.py | 30 +++++++++++++++++++++++ 2 files changed, 39 insertions(+), 20 deletions(-) diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index b21a4f0e4e..5da1b207e3 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -1,3 +1,4 @@ +from collections import deque from typing import List import torch @@ -10,39 +11,27 @@ logger = init_logger("lightllm.common.req_manager") -class _ReqNode: - def __init__(self, index): - self.index = index - self.next: "_ReqNode" = None - - class _ReqLinkedList: def __init__(self, max_request_num): - self.nodes = [_ReqNode(i) for i in range(max_request_num)] + self._free_indexes = deque(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 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 + index = self._free_indexes.popleft() + assert self.marks[index] == 0 + self.marks[index] = 1 self.can_alloc_size -= 1 - return get_node.index + 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._free_indexes.appendleft(index) self.marks[index] = 0 self.can_alloc_size += 1 return diff --git a/unit_tests/common/test_req_manager_cpu.py b/unit_tests/common/test_req_manager_cpu.py index b8d20f8657..473624f95d 100644 --- a/unit_tests/common/test_req_manager_cpu.py +++ b/unit_tests/common/test_req_manager_cpu.py @@ -33,6 +33,36 @@ def _make_manager(device="cpu"): return manager +def test_req_index_allocation_preserves_order_capacity_and_reuse(): + indexes = _ReqLinkedList(3) + assert indexes.is_all_free() and indexes.can_alloc_size == 3 + assert [indexes.alloc() for _ in range(3)] == [0, 1, 2] + assert indexes.can_alloc_size == 0 and not indexes.is_all_free() + assert indexes.alloc() is None + + indexes.free(1) + indexes.free(0) + assert indexes.can_alloc_size == 2 + assert [indexes.alloc(), indexes.alloc()] == [0, 1] + assert indexes.alloc() is None + for index in range(3): + indexes.free(index) + assert indexes.is_all_free() and indexes.can_alloc_size == 3 + + +def test_req_index_free_rejects_unallocated_and_duplicate_indexes(): + indexes = _ReqLinkedList(2) + with pytest.raises(AssertionError): + indexes.free(0) + assert indexes.alloc() == 0 + indexes.free(0) + with pytest.raises(AssertionError): + indexes.free(0) + assert indexes.can_alloc_size == 2 + assert [indexes.alloc(), indexes.alloc()] == [0, 1] + assert indexes.alloc() is None + + def test_cpu_updates_own_allocator_values_and_defer_gpu_publication(): manager = _make_manager() indexes = torch.tensor([20, 21, 22, 23], dtype=torch.int32) From cff1cd87629492c84a5f93b5b016b2bf0880c032 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 14 Oct 2026 09:06:41 +0000 Subject: [PATCH 07/22] perf: remove redundant PD first-token stream synchronization Rely on the inference entrypoint stream dependency to order first-token writes before GPU reads. Document that CPU request updates do not require waiting for the CUDA write and that future stream changes must preserve this dependency. Validation: 19 PD and paged-allocation tests passed; git diff --check passed. --- .../server/router/model_infer/mode_backend/base_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 c0214be576..4d1143dcd4 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -557,8 +557,8 @@ def _read_pd_trans_io_buffer_and_update_req_status(self): self.model.req_manager.req_sampling_params_manager.req_to_next_token_ids ) req_to_next_token_ids[req.req_idx, 0:1].fill_(obj.first_gen_token_id) - # TODO: 分析 PD 首 token 写入后同步当前流的必要性，以及对后续推理的阻塞影响。 - torch.cuda.current_stream().synchronize() + # handle 只更新 CPU 状态；推理入口通过 overlap_stream.wait_stream(current_stream) + # 保证读取首 token 前写入已完成，无需在此阻塞 CPU。修改写入或推理流时须保留此依赖。 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, From 15942d969d3980d55dd68250c6bff8bb33bd114a Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 14 Oct 2026 09:25:34 +0000 Subject: [PATCH 08/22] fix: disable prefix cache for prompt logprobs requests --- .../server/router/model_infer/infer_batch.py | 3 +- .../test_prompt_logprobs_cache.py | 102 ++++++++++++++++++ 2 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 41df035fd5..a40823b25f 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -459,7 +459,8 @@ def __init__( vocab_size: int, ) -> None: self.shm_param = shm_req.sample_params - self.disable_prompt_cache = self.shm_param.disable_prompt_cache + # Cached KV skips prompt logits, so logprobs requests must recompute the full prompt. + 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 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py new file mode 100644 index 0000000000..b22f9a7403 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py @@ -0,0 +1,102 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +# Follow the production import order to avoid the existing package initialization cycle. +from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 +from lightllm.server.core.objs.sampling_params import SamplingParams +from lightllm.server.router.model_infer import infer_batch +from lightllm.server.router.model_infer.infer_batch import InferReq, InferSamplingParams +from lightllm.server.router.model_infer.mode_backend import base_backend + + +def _make_req_with_cached_prefix(monkeypatch, prompt_logprobs, disable_prompt_cache=False, hybrid=False): + sample_params = SamplingParams() + sample_params.top_k = -1 + sample_params.prompt_logprobs = prompt_logprobs + sample_params.disable_prompt_cache = disable_prompt_cache + req = InferReq.__new__(InferReq) + req.req_idx = 0 + req.args = SimpleNamespace( + page_size=2, linear_att_hash_page_size=2, linear_att_page_block_num=2, max_req_total_len=16 + ) + req.cur_kv_len = req.hold_kv_len = req.move_to_gpu_index_len = 0 + req.shared_kv_node = req.tail_small_page_buffer_id = None + req.shm_req = SimpleNamespace( + sample_params=sample_params, + input_len=6, + prompt_cache_len=0, + shm_prompt_ids=SimpleNamespace(arr=torch.arange(6).numpy()), + hybrid_token_hash_list=SimpleNamespace(size=2, get_all=lambda: [10, 20]), + ) + req.sampling_param = InferSamplingParams(req.shm_req, vocab_size=8) + req.get_cur_total_len = lambda: 6 + node = SimpleNamespace(node_prefix_total_len=4, big_page_buffer_idx=0, is_big_page_node=lambda: True) + match_prefix = Mock(return_value=(node, 4, torch.arange(4, dtype=torch.int32))) + manager = SimpleNamespace( + req_to_token_indexs_cpu=torch.zeros((1, 8), dtype=torch.int32), + restore_big_page_state=Mock(), + init_hybrid_attention_state=Mock(), + ) + monkeypatch.setattr( + infer_batch, + "g_infer_context", + SimpleNamespace( + is_hybrid_att_model=hybrid, req_manager=manager, radix_cache=SimpleNamespace(match_prefix=match_prefix) + ), + ) + return req, match_prefix + + +@pytest.mark.parametrize("hybrid", [False, True]) +@pytest.mark.parametrize( + "prompt_logprobs, disable_prompt_cache, expected_prefix_len", + [(-1, False, 4), (0, False, 0), (3, False, 0), (-1, True, 0)], +) +def test_prompt_logprobs_requests_recompute_cached_prefix( + monkeypatch, hybrid, prompt_logprobs, disable_prompt_cache, expected_prefix_len +): + req, match_prefix = _make_req_with_cached_prefix(monkeypatch, prompt_logprobs, disable_prompt_cache, hybrid) + + if hybrid: + req._hybrid_match_radix_cache() + else: + req._match_radix_cache() + + assert req.cur_kv_len == req.hold_kv_len == expected_prefix_len + assert req.shm_req.prompt_cache_len == expected_prefix_len + assert req.shm_req.shm_cur_kv_len == expected_prefix_len + assert match_prefix.call_count == (1 if expected_prefix_len else 0) + + +@pytest.mark.parametrize("prompt_logprobs", [0, 3]) +def test_cached_prompt_still_captures_logprobs_for_every_prompt_position(monkeypatch, prompt_logprobs): + req, _ = _make_req_with_cached_prefix(monkeypatch, prompt_logprobs) + req._match_radix_cache() + req.prompt_selected_logprobs = SimpleNamespace(add_chunk=Mock()) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + backend.disable_chunked_prefill = True + backend.is_master_in_dp = True + q_len = req.get_cur_total_len() - req.cur_kv_len + indexes = torch.arange(q_len, dtype=torch.int32) + backend.model = SimpleNamespace(return_all_prompt_logics=True, _select_mem_indexes=lambda _: indexes) + capture_manager = SimpleNamespace(is_buffer_initialized=lambda: True, capture=Mock()) + monkeypatch.setattr(base_backend.PromptLogprobsCaptureManager, "get_instance", lambda: capture_manager) + logits = torch.arange(q_len * 8, dtype=torch.float32).reshape(q_len, 8) + + backend._capture_prompt_logprobs_if_needed(model_input=None, run_reqs=[req], prompt_logits=logits) + + if prompt_logprobs == 0: + req.prompt_selected_logprobs.add_chunk.assert_called_once() + start, end, logprobs, ranks = req.prompt_selected_logprobs.add_chunk.call_args.args + assert (start, end) == (1, req.shm_req.input_len) + assert logprobs.shape == ranks.shape == (5,) + expected = logits[:5].log_softmax(dim=-1)[torch.arange(5), torch.arange(1, 6)] + torch.testing.assert_close(logprobs, expected) + else: + capture_manager.capture.assert_called_once() + captured = capture_manager.capture.call_args.kwargs + assert captured["mem_indexes"].tolist() == list(range(5)) + assert captured["top_token_ids"].shape == captured["top_logprobs"].shape == (5, prompt_logprobs) From 43303f8472cea0a766054f1b185b738870bbdfed Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 14 Oct 2026 09:31:47 +0000 Subject: [PATCH 09/22] docs: clarify stream-ordered KV index copy lifetime --- docs/kv_cache_page_size.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index fc4818e463..69acd152af 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -27,8 +27,10 @@ KV 存储和请求表，但资源的申请、缓存和释放以完整物理页 `_get_classed_reqs` 返回前通过 backend 的 `_copy_req_token_indexes_to_gpu` 统一收集 prefill/decode 请求的 `[move_to_gpu_index_len, hold_kv_len)` 区间，通过一次原生 `cudaMemcpyBatchAsync` 提交， -成功后将长度更新到 `hold_kv_len`。`DuringApiCall` 源访问属性保证返回前读取完 CPU 数据， -后续计算依靠同流顺序读取 GPU 表；长度表示异步拷贝已排队，不表示设备已经完成拷贝。 +成功后将长度更新到 `hold_kv_len`。原生批量拷贝使用 `cudaMemcpySrcAccessOrderStream` 源访问属性， +源数据读取遵循当前非默认推理流的顺序。调用方必须保持对应的 pinned CPU 源区间有效且内容不变， +直到传输完成；API 返回时源数据读取可能仍未完成，修改、释放或复用源区间前须确认传输已完成。 +后续计算依靠同流顺序读取 GPU 表；`move_to_gpu_index_len` 表示异步拷贝已排队，不表示设备已经完成拷贝。 暂停/释放时重置长度，前缀重映射时回退长度以重新上传。 DP 在发布请求 KV 信息前，让专用传输流等待 overlap 流已有的索引更新， 再将各 TP rank 的待上传索引拷贝到 GPU 表； From 997404ac99f4dfb3c4a6fbc4acc0624fd6087c82 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Thu, 22 Oct 2026 16:30:38 +0000 Subject: [PATCH 10/22] Refactor request token indexes and paged KV allocation Replace the global CPU request table with per-request pinned-cache and pure-GPU modes. Keep token-level interfaces over page-stride allocation and append-only request indexes. Centralize batch H2D stream ordering and pinned-source lifetime management, and update cache, PD/DP, metadata and benchmark callers. Keep lifecycle tests runnable without CUDA and cover real GPU uploads separately. Fix formatting and lint issues. Validation: 228 tests passed on GPUs 6/7; 72 passed and 156 CUDA-dependent tests skipped without visible GPUs. Black, Flake8 and diff checks passed. --- docs/kv_cache_page_size.md | 152 ++++- lightllm/common/basemodel/basemodel.py | 6 +- lightllm/common/cuda_wrapper.py | 132 +++- .../common/kv_cache_mem_manager/allocator.py | 125 ++-- .../kv_cache_mem_manager/mem_manager.py | 2 +- lightllm/common/req_manager/base.py | 32 +- lightllm/server/api_cli.py | 11 + lightllm/server/core/objs/start_args_type.py | 1 + .../server/router/model_infer/infer_batch.py | 70 +-- .../router/model_infer/infer_req_ext.py | 10 +- .../router/model_infer/infer_req_ext1.py | 215 +++++++ .../model_infer/mode_backend/base_backend.py | 59 +- .../mode_backend/diverse_backend/impl.py | 29 +- .../dp_backend/dp_shared_kv_trans.py | 18 +- .../mode_backend/multi_level_kv_cache.py | 32 +- .../pd/decode_node_impl/decode_impl.py | 18 +- .../pd/prefill_node_impl/prefill_impl.py | 4 +- .../benchmark_event_synchronize_gil.py | 233 +++++++ test/benchmark/benchmark_pinned_empty_h2d.py | 248 ++++++++ .../static_inference/static_benchmark.py | 44 +- unit_tests/common/cpu_index_helpers.py | 70 +++ unit_tests/common/test_kv_cache_allocator.py | 333 ++++++++++ .../common/test_pinned_memory_stream.py | 353 +++++++++++ unit_tests/common/test_req_manager_cpu.py | 316 ++++++---- unit_tests/common/test_req_manager_page.py | 84 ++- unit_tests/common/test_req_token_index.py | 594 ++++++++++++++++++ .../mode_backend/test_diverse_backend.py | 97 ++- .../mode_backend/test_multi_level_kv_cache.py | 114 +++- .../test_paged_side_allocations.py | 51 +- .../test_prompt_logprobs_cache.py | 33 +- unit_tests/test_static_benchmark_paged_kv.py | 45 +- 31 files changed, 3042 insertions(+), 489 deletions(-) create mode 100644 lightllm/server/router/model_infer/infer_req_ext1.py create mode 100644 test/benchmark/benchmark_event_synchronize_gil.py create mode 100644 test/benchmark/benchmark_pinned_empty_h2d.py create mode 100644 unit_tests/common/cpu_index_helpers.py create mode 100644 unit_tests/common/test_kv_cache_allocator.py create mode 100644 unit_tests/common/test_pinned_memory_stream.py create mode 100644 unit_tests/common/test_req_token_index.py diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index 69acd152af..04186270d4 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -20,24 +20,140 @@ KV 存储和请求表，但资源的申请、缓存和释放以完整物理页 真实推理索引由模型执行层从请求表选取。 `ReqManager` 只保留请求 ID、请求表及其原有生命周期职责，不提供 page allocator 接口。 -`ReqManager.req_to_token_indexs_cpu` 是与 GPU 请求表同形状、同 dtype 的 CPU 镜像。 -调度分配、前缀匹配和 diverse 共享直接拷贝到镜像；缓存插入、释放和 PD 元数据直接读取镜像。 -缓存持有的索引必须复制，避免请求槽位复用后改变缓存内容。 -每个 `InferReq` 用 `move_to_gpu_index_len` 记录已经提交到推理流的 GPU 索引长度。 -`_get_classed_reqs` 返回前通过 backend 的 `_copy_req_token_indexes_to_gpu` 统一收集 -prefill/decode 请求的 -`[move_to_gpu_index_len, hold_kv_len)` 区间，通过一次原生 `cudaMemcpyBatchAsync` 提交， -成功后将长度更新到 `hold_kv_len`。原生批量拷贝使用 `cudaMemcpySrcAccessOrderStream` 源访问属性， -源数据读取遵循当前非默认推理流的顺序。调用方必须保持对应的 pinned CPU 源区间有效且内容不变， -直到传输完成；API 返回时源数据读取可能仍未完成，修改、释放或复用源区间前须确认传输已完成。 -后续计算依靠同流顺序读取 GPU 表；`move_to_gpu_index_len` 表示异步拷贝已排队，不表示设备已经完成拷贝。 -暂停/释放时重置长度，前缀重映射时回退长度以重新上传。 -DP 在发布请求 KV 信息前，让专用传输流等待 overlap 流已有的索引更新， -再将各 TP rank 的待上传索引拷贝到 GPU 表； -当前流等待传输流，再由同步 NCCL barrier 保证上传完成。其他 DP rank 在 barrier 后 -通过 GPU 请求表读取分配元数据，不跨进程共享 CPU 镜像。 -请求筛选、地址计算、批量上传及长度更新均在这个 backend 函数内完成；`ReqManager` 只持有映射表。 -独立 benchmark 和模型初始化检查共用 `cuda_memcpy_batch_async` 底层调用。 +## KV 页分配 + +`KvCacheAllocator(size, page_size)` 的 `size`、申请数量、释放索引和共享剩余容量均使用 token 单位。 +内部 pinned `mem_state` 每页仅保存一个物理 token 页首索引，形如 `[0, page_size, 2 * page_size, ...]`， +游标按页计数。申请和释放必须覆盖完整页；释放时按 `page_size` 取页首，重新申请时展开回 token 索引。 +分配结果直接写入 token 级 pinned 返回 tensor：`page_size=1` 使用 `copy_`， +其他页大小使用广播 `torch.add(..., out=...)`，不再分配中间展开 tensor。 +返回缓冲区按需创建，单次申请超过 `1024 * 1024` 个 token 时独立分配；其他非空申请从 +`1024 * 1024` 个 `int32` 元素的 pinned 缓冲区切片，仅保留未使用的尾部，不回绕覆盖旧结果。 +尾部不足时新建缓冲区，旧切片自行持有原存储；空申请不创建大缓冲区。 +`mem_state` 的容量随 `page_size` 缩小，分配器额外持有的返回缓冲区底层存储最多为 4 MiB， +不再按 KV 总容量预分配三倍 token 表。在途结果及 PyTorch pinned allocator 缓存的存储另计。 + +## 请求 Token 索引 + +`ReqManager` 只持有 GPU `req_to_token_indexs`，不再分配全局 CPU 页表。 +每个 `InferReq` 通过成员 `token_index` 持有一个 `ReqTokenIndex`，实现位于 `server/router/model_infer/infer_req_ext1.py`。 +索引操作直接使用 `req.token_index.append_token_indexes(...)` 等接口，`InferReq` 不再提供同名转发方法。 +基类仅在构造函数中初始化公共对象和状态，其余方法只定义抽象接口；两种子类各自实现读写、上传和重置逻辑。 +扩展构造函数接收 `infer_req` 和必填的 `index_stream`，仅提取并保存固定的 `req_idx`、GPU 请求表行视图及索引流， +避免反向持有请求对象形成循环引用；索引流由外部选择，Helper 内部不查询默认流或捕获构造时的当前流。 +CPU 缓存实现同时记录输入长度 `input_len`，helper 在共享请求初始化后、前缀匹配前构造。 +工厂 `create_req_token_index(infer_req, index_stream)` 从请求参数中选择实现，流参数直接传给子类， +不再单独传 GPU 行和模式。在线服务和静态 benchmark 在统一创建入口传入 `torch.cuda.current_stream()`， +Helper 保存该流，后续不随调用点的当前流切换： + +| `--req_token_index_mode` | 存储与优点 | 代价 | +| --- | --- | --- | +| `pin_mem_cache`（默认） | `PinMemCacheReqTokenIndex` 按请求直接申请 pinned CPU token 缓存；CPU 读取无需 D2H，GPU 更新可批量提交 | 占用 pinned CPU 内存，需要维护上传进度 | +| `pure_gpu` | `PureGpuReqTokenIndex` 直接维护 GPU 请求表，无常驻 CPU 索引缓存 | CPU 读取需要 D2H 并等待结果，频繁的 CPU 索引操作可能更慢 | + +两种模式均保留相同的 GPU 请求表，显存占用基本相同。追加后的所需长度为 `end = token_index_len + mem_indexes.numel()`。 +CPU 缓存首次按 `max(input_len, end)` +向上对齐到 1K token 申请；后续容量不足时，在本次所需长度 `end` 上多预留 2K token，再向上对齐到 1K， +容量上限为 GPU 请求表行长度。缓冲区在首次非空写入时按需申请，重置时保留输入长度，丢弃已申请的缓冲区引用； +由新建的 `torch.empty(..., pin_memory=True, requires_grad=False)` +独立持有，不借用 KV allocator 的返回缓冲区。索引缓存按 token 保存，不再按页压缩或展开。 + +统一接口为 `append_token_indexes(mem_indexes)`、`get_cpu_token_indexes(start, end)`、 +`get_free_token_indexes(start, end)` 和 `get_gpu_token_indexes(start, end)`，均使用 token 单位， +追加接口将完整输入接到已记录索引的尾部，非空输入在两种模式下均断言为 CPU Tensor；空输入直接返回。 +CPU 缓存模式直接复制到自身的 pinned 缓冲区，无需先 pin 输入；GPU 模式仅在输入未 pinned 时内部转换，已有 pinned 输入直接复用。 +前缀匹配和索引重建路径直接传入 CPU 索引，不再自行 pin。读取接口支持页内任意区间；提交给 allocator 的释放区间仍须覆盖完整页。 +CPU getter 返回独立快照，缓存、PD 元数据和异步任务可跨请求释放持有；GPU getter 返回当前请求的 +GPU 表视图，不应跨请求槽位复用持有。Helper 保存外部传入的索引流，GPU 索引读写固定在该流， +调用方保证 token 区间有效；GPU getter 返回前让调用流等待索引流，没有新增上传时也需承接已有写入的依赖。 +调用流就是索引流时不额外等待；这里只排入 GPU 依赖，不阻塞 CPU，也不改变调用方的当前流。 +模型和 KV 传输流在各自入口统一等待默认流的流程继续保留，Helper 不维护逐请求 CUDA 写事件。 + +释放专用接口 `get_free_token_indexes` 在 `pin_mem_cache` 模式返回 CPU 快照，在 `pure_gpu` +模式复用 GPU getter 的范围校验和流等待，再返回 GPU 请求表切片，省掉完整 token 索引的 D2H 快照。 +流等待只保证已有索引写入就绪；重写请求槽位前，仍须在同一流按序消费旧视图，或让索引流等待消费流。 +批量释放先合并视图，再归还请求 ID，暂停恢复也遵循此顺序。 +allocator 仍维护 CPU 空闲页表，因此最终按 `page_size` 取页首并阻塞 D2H，返回时才能重新分配； +此接口不取消该同步，也不能代替缓存和元数据需要的独立 CPU 快照。 + +### 上传与源内存生命周期 + +`token_index_len` 是当前连续记录的 token 索引长度，包含页内预留槽位；每次追加从该位置开始， +并按输入元素个数增长。CPU 缓存扩容时保留 +`[0, token_index_len)` 的已有映射。它与 CPU 缓存容量 `_cpu_buffer.numel()`、已算好 KV 的长度 +`cur_kv_len` 不同。 +`gpu_index_submitted_len` 是最新 GPU 索引写入已提交的前缀右端点，后续索引流操作可按序使用 +`[0, gpu_index_submitted_len)`。它表示提交进度，不表示 GPU 传输已完成，也不是累计上传量； +CPU 缓存模式的追加接口只更新 CPU 索引和 `token_index_len`，不修改 GPU 提交进度。 +需要替换前缀时，先保存仍需保留的索引（包括预留尾部），再调用 `reset()` 并按顺序追加，重建完整连续映射。 +首次前缀匹配时索引为空；暂停恢复前已由统一释放流程重置，两处匹配入口直接追加命中前缀，不重复 `reset()`。 +释放或暂停请求时也由索引扩展的 `reset()` 统一清零两个长度。 +CPU 模式重置时将 `_cpu_buffer` 设为 `None`，下一次追加重新申请，随后从零开始提交 GPU 更新。 +单独的 GPU getter 直接上传 CPU 缓存视图：普通追加不会改写已有区间，扩容和重置使用新存储， +PyTorch `copy_(non_blocking=True)` 会登记 pinned 源的流使用，旧存储在 H2D 完成前不会被回收复用。 +因此不再为单个请求上传额外申请 pinned 快照或执行一次 CPU 拷贝，也不需要在 reset 中同步 CUDA。 + +`get_upload_indexes()` 不接受区间参数，CPU 模式返回内部计算的 +`[gpu_index_submitted_len, token_index_len)` 视图，不改变提交进度；无待上传内容时返回空 tensor， +纯 GPU 模式返回 `None`。CPU 缓存模式的 GPU getter 校验 `end <= token_index_len` 后按需提交完整待上传区间， +纯 GPU 模式校验 `end <= gpu_index_submitted_len`；读取均不允许推进超出已有索引的提交进度。 +批量入口遇到存在未提交区间、但没有 CPU 上传源的情况会直接断言，不再替 GPU 模式补写提交进度。 + +例如 CPU 缓存模式追加 64 个索引后，两个字段分别为 `64/0`；提交到 GPU 后为 `64/64`； +再追加 32 个索引后为 `96/64`，下次仅需提交 `[64,96)`。需要替换映射时，`reset()` 将两者清为 `0/0`， +重新追加 96 个完整索引后为 `96/0`，再从头提交 GPU 更新。纯 GPU 模式在追加接口中 +直接提交写表，所以追加后两者相同。 +`_get_classed_reqs` 返回前，backend 收集各请求待上传的 pinned 视图，直接使用 +`cudaMemcpyBatchAsync` 提交，不再申请聚合缓冲区或额外复制 CPU 索引。 +同一批请求必须共用一个索引流，backend 对此进行断言；随后只准备源、目标视图列表， +在索引流上下文调用批量拷贝并更新提交进度。返回前仅在调用流不同于索引流时排入等待， +没有新上传时也承接已有写入的依赖。backend 不再持有上传流，也不逐条交接多条索引流。 +原生 API 不接受 legacy 默认流，适配由 `PinnedH2DBatchCopier` 内部处理：按设备延迟创建并复用非默认上传流， +先让上传流等待默认流，提交后再让默认流等待上传流；非默认索引流直接提交，PyTorch 回退直接使用当前流。 +这些操作只排入 GPU 依赖，不阻塞 CPU。 +模型入口原有的推理流等待默认流继续保留；普通增量追加不等待上一批模型执行结束。 + +原生拷贝使用 `cudaMemcpySrcAccessOrderStream`，需要源内存在实际传输结束前有效且不变。 +仅调用 `index_stream.wait_stream(upload_stream)` 不能保护 CPU 源的生命周期，因为 CPU 会立即返回。 +`cuda_wrapper.py` 将原生接口开关、内部上传流、后台线程池和提交/保活方法统一封装在 `PinnedH2DBatchCopier` 中， +原有 `cuda_memcpy_pinned_h2d_batch_async` 函数保留为委托入口。原生提交成功后，在同一上传流上记录一个 +`torch.cuda.Event(blocking=True)`，交给进程内复用的单线程 `ThreadPoolExecutor` 等待。 +线程池在首次需要保活时加锁创建，工作线程在首次 submit 时启动；模块导入、空批次和 PyTorch 回退均不创建线程池， +也不因线程池初始化而初始化 CUDA。 +后台任务持有 `list(src_pinned_tensors)`，独立于调用方可变的列表；源切片会保留完整 pinned storage， +即使调用方立即丢弃源 Tensor 或清空自己的列表，底层内存仍不会被回收复用。 +后台 `event.synchronize()` 休眠等待并释放 GIL，完成后清空源列表，不向 Future 返回 Tensor； +主推理线程不等待后台任务，默认索引流由封装承接内部上传流的依赖，调用流按需等待索引流。 +索引扩展不维护逐请求事件，backend 不维护保活列表，也不需要额外的小片段 H2D 或 C++ 扩展。 +该方式只保护引用生命周期，不阻止 CPU 主动改写数据。请求缓存只追加，reset 和扩容换用新存储， +因此可直接保活缓存视图；上传期间不能改写已提交的区间，无需独立聚合快照。 +单线程按提交顺序处理；多上传流共用线程池时可能延后回收，不改变 GPU 工作顺序。 +正常进程退出时由线程池等待未完成任务，不取消保活任务。 + +事件创建、记录、后台任务提交及完成等待属于信任域；异常直接传播，不增加同步收尾、源引用隔离或恢复逻辑。 +原生 batch API 不可用时回退到普通 PyTorch `copy_`，由其内部登记 host allocator，无需创建后台保活任务。 + +GPU 模式在追加接口中直接向构造时指定的索引流提交 GPU 更新，槽位复用后的新旧写入遵循同一流顺序； +追加后由调用流等待索引流，D2H 读取也固定在索引流并阻塞到 CPU 快照就绪。 +已 pinned 的输入直接用于 H2D，省去中间快照的申请和 CPU 拷贝；普通 CPU 输入在追加接口内部 pin 后上传。 +KV allocator 返回的 pinned 切片不会被后续申请覆盖；前缀匹配和重建路径直接传入 CPU 索引，由 Helper 按需处理。 +PyTorch `copy_(non_blocking=True)` 内部向 host allocator 登记源 storage 的流使用， +因此调用方可立即释放输入引用，无需手工记录事件；但在 H2D 完成前不能主动改写源内容。 + +### 请求结束与跨流使用 + +暂停或释放时重置请求 Helper；前缀重映射通过重置后重新追加，重新建立 CPU 模式的上传进度。 +PD Decode 传输失败时，等待全部任务返回后将 `cur_kv_len` 回退到原共享前缀长度（无共享前缀则为 0）， +保留 `hold_kv_len` 和索引状态，由后续统一释放回收尾部页面，避免失败 KV 进入缓存。 +CPU KV Cache 在卸载流上下文中调用 GPU getter；Helper 在请求索引流上准备数据，再让当前卸载流等待索引流，调用方无需额外同步。 +纯 GPU 模式直接使用请求表视图，不再执行无谓的 D2H/H2D。 +DP 发布前复用统一上传入口，所有 TP rank 通过默认流上的 NCCL barrier 确认索引可见后再发布。 +diverse 模式会重写旧前缀，在重写入口统一让默认流等待推理流，防止覆盖上一批模型仍在读取的映射。 +KV 卸载完成事件和请求延迟释放机制仍保留，它们负责实际 KV 任务和槽位生命周期。 +模型预热不创建请求对象，直接将 allocator 返回的 token 索引拷入 GPU 请求表。 +静态 benchmark 使用包含 `req_idx`、`req_manager`、`args` 和 `shm_req.input_len` 的轻量请求对象创建索引 Helper， +不进入真实 `InferReq` 的共享内存生命周期。 +HOLD 请求直接在 GPU 表初始化为预留物理页的循环索引。 ## 生命周期 diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 68a2bf27e6..c89a88d90c 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -1000,8 +1000,7 @@ def _check_max_len_infer(self): alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size mem_indexes = self.mem_manager.alloc(alloc_token_num) req_idx = int(b_req_idx[0]) - self.req_manager.req_to_token_indexs_cpu[req_idx, : len(mem_indexes)].copy_(mem_indexes) - self.req_manager.req_to_token_indexs.copy_(self.req_manager.req_to_token_indexs_cpu, non_blocking=True) + self.req_manager.req_to_token_indexs[req_idx, :alloc_token_num].copy_(mem_indexes, non_blocking=True) b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = self.batch_max_tokens b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda") @@ -1083,8 +1082,7 @@ def _autotune_warmup(self): alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size mem_indexes = self.mem_manager.alloc(alloc_token_num) req_idx = int(b_req_idx[0]) - self.req_manager.req_to_token_indexs_cpu[req_idx, : len(mem_indexes)].copy_(mem_indexes) - self.req_manager.req_to_token_indexs.copy_(self.req_manager.req_to_token_indexs_cpu, non_blocking=True) + self.req_manager.req_to_token_indexs[req_idx, :alloc_token_num].copy_(mem_indexes, non_blocking=True) b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = input_len b_ready_cache_len = torch.zeros(1, dtype=torch.int32, device="cuda") diff --git a/lightllm/common/cuda_wrapper.py b/lightllm/common/cuda_wrapper.py index b089227004..5d37e8327f 100644 --- a/lightllm/common/cuda_wrapper.py +++ b/lightllm/common/cuda_wrapper.py @@ -16,13 +16,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""This file is a pure Python wrapper for the cudart library. -It avoids the need to compile a separate shared library, and is -convenient for use when we just need to call a few functions. -""" +"""Python cudart wrapper with asynchronous pinned-source lifetime management.""" import ctypes +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass +from threading import Lock from typing import Any, Dict, List, Optional # this line makes it possible to directly load `libcudart.so` using `ctypes` @@ -38,48 +37,125 @@ cudaError_t = ctypes.c_int cudaMemcpyKind = ctypes.c_int -_cuda_memcpy_batch_available = True -def cuda_memcpy_pinned_h2d_batch_async(dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor]): - """在当前非默认流上，将 pinned CPU Tensor 批量异步拷贝到 GPU Tensor，仅支持 H2D。 +class PinnedH2DBatchCopier: + """每个推理进程共享批量 H2D 状态，统一管理异步源保活与异常回退。""" - 使用 Stream 模式；调用方保证源为 pinned 内存、目标在当前 CUDA 设备上， - 各对 Tensor 连续且形状和 dtype 一致，并保持源数据有效且不变直到传输完成。 - """ - global _cuda_memcpy_batch_available + _native_batch_available = True + # 原生 batch 禁止 legacy 默认流，仅在默认流调用时按设备延迟创建上传流。 + _upload_streams: Dict[torch.device, torch.cuda.Stream] = {} + # 只在原生 H2D 成功提交、需要后台保活时创建线程池，导入和 copy_ 回退均不创建。 + _release_pool: Optional[ThreadPoolExecutor] = None + _release_pool_lock = Lock() + + @classmethod + def copy_async(cls, dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor]): + """按当前流顺序批量异步 H2D，并保护 pinned 源直到传输结束。 + + 源由 PyTorch pinned allocator 申请，各对 Tensor 连续且形状和 dtype 一致。 + 默认流的原生 batch 适配封装在内部；调用方只遵循当前流顺序，不阻塞 CPU。 + 上传完成前不能主动改写源，目标须保持有效到传输结束。 + """ + if not dst_gpu_tensors: + return + stream = torch.cuda.current_stream() + if cls._native_batch_available and stream.cuda_stream == 0: + # 原生 API 不接受 legacy 默认流；上传流先等待当前流，提交后再交回依赖。 + upload_stream = cls._get_upload_stream(stream.device) + upload_stream.wait_stream(stream) + with torch.cuda.stream(upload_stream): + cls._copy_on_stream(dst_gpu_tensors, src_pinned_tensors, upload_stream) + stream.wait_stream(upload_stream) + else: + cls._copy_on_stream(dst_gpu_tensors, src_pinned_tensors, stream) + + @classmethod + def _get_upload_stream(cls, device: torch.device) -> torch.cuda.Stream: + if device not in cls._upload_streams: + cls._upload_streams[device] = torch.cuda.Stream(device=device) + return cls._upload_streams[device] + + @classmethod + def _copy_on_stream( + cls, dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor], stream: torch.cuda.Stream + ): + """在已选定的当前流提交，沿用原生保活和 PyTorch 回退。""" + if cls._native_batch_available and cls._try_native_batch_h2d(dst_gpu_tensors, src_pinned_tensors, stream): + cls._protect_pinned_sources(src_pinned_tensors, stream) + return - if not dst_gpu_tensors: - return - stream = torch.cuda.current_stream() - # CUDA 批量拷贝接口明确禁止使用 legacy 默认流（句柄为 0）。 - # 在异常回退前检查，避免默认流调用错误被吞掉并关闭全局 batch 开关。 - assert stream.cuda_stream != 0, "cuda_memcpy_pinned_h2d_batch_async 必须在非默认 CUDA 流上调用" + # PyTorch copy_ 已经登记 host allocator，回退路径无需额外保活或登记。 + for dst, src in zip(dst_gpu_tensors, src_pinned_tensors): + dst.copy_(src, non_blocking=True) - if _cuda_memcpy_batch_available: + @classmethod + def _try_native_batch_h2d( + cls, dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor], stream: torch.cuda.Stream + ) -> bool: + """尝试原生批量提交；接口准备或提交失败时禁用该路径，让调用方回退到 copy_。""" try: from cuda.bindings import runtime as cuda_runtime attrs = cuda_runtime.cudaMemcpyAttributes() - # 使用 Stream 模式，让源数据读取遵循当前流顺序。 - # 调用方需保持源数据有效且不变，直到传输完成。 + # Stream 模式让源读取遵循上传流顺序；提交成功后由完成事件和后台任务保活。 attrs.srcAccessOrder = cuda_runtime.cudaMemcpySrcAccessOrder.cudaMemcpySrcAccessOrderStream dst_ptrs = [tensor.data_ptr() for tensor in dst_gpu_tensors] src_ptrs = [tensor.data_ptr() for tensor in src_pinned_tensors] byte_sizes = [tensor.numel() * tensor.element_size() for tensor in src_pinned_tensors] - (error,) = cuda_runtime.cudaMemcpyBatchAsync( - dst_ptrs, src_ptrs, byte_sizes, len(byte_sizes), [attrs], [0], 1, stream.cuda_stream + batch_copy = cuda_runtime.cudaMemcpyBatchAsync + except Exception as error: + cls._native_batch_available = False + logger.warning( + "Native batch H2D API unavailable; " + f"disabling native batch copy and falling back to PyTorch copy_: {error}" ) + return False + + try: + (error,) = batch_copy(dst_ptrs, src_ptrs, byte_sizes, len(byte_sizes), [attrs], [0], 1, stream.cuda_stream) if error != cuda_runtime.cudaError_t.cudaSuccess: raise RuntimeError(f"cudaMemcpyBatchAsync failed: {error.name}") except Exception as error: - _cuda_memcpy_batch_available = False - logger.warning(f"Disabling CUDA batch copies; falling back to individual copies: {error}") - else: - return + cls._native_batch_available = False + logger.warning( + "Native batch H2D submission failed; " + f"disabling native batch copy and falling back to PyTorch copy_: {error}" + ) + return False + return True + + @classmethod + def _protect_pinned_sources(cls, src_pinned_tensors: List[torch.Tensor], stream: torch.cuda.Stream): + """每批 raw H2D 在上传流记录一个完成事件，由复用线程池接管源 Tensor 引用。""" + # blocking=True 让后台线程休眠等待；Event.synchronize 会释放 GIL，不阻塞主推理线程。 + # 事件排在 raw batch 之后，覆盖此前所有源读取；调用方的 wait_stream 负责 GPU 数据依赖。 + event = torch.cuda.Event(blocking=True) + event.record(stream) + # 必须复制容器：调用方可能立即清空自己的列表，切片本身仍持有整块 pinned storage。 + cls._get_release_pool().submit(cls._wait_and_release_pinned_sources, event, list(src_pinned_tensors)) + + @classmethod + def _get_release_pool(cls) -> ThreadPoolExecutor: + if cls._release_pool is None: + # 首次创建加锁，避免多个提交线程同时建立等待线程池；后续调用直接复用。 + with cls._release_pool_lock: + if cls._release_pool is None: + cls._release_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="lightllm-pinned-release") + return cls._release_pool + + @classmethod + def _wait_and_release_pinned_sources(cls, event: torch.cuda.Event, sources: List[torch.Tensor]): + """只管理 CPU 源的生命周期，不修改请求状态，也不向流提交新的 GPU 工作。""" + event.synchronize() + # 普通工作线程可安全释放 Tensor；不受 CUDA host callback 禁止调用 CUDA API 的限制。 + # 不把 sources 返回给 Future，避免任务结束后仍被结果对象持有。 + sources.clear() + - for dst, src in zip(dst_gpu_tensors, src_pinned_tensors): - dst.copy_(src, non_blocking=True) +def cuda_memcpy_pinned_h2d_batch_async(dst_gpu_tensors: List[torch.Tensor], src_pinned_tensors: List[torch.Tensor]): + """按当前流顺序批量异步 H2D，默认流适配、源内存保活和回退由封装统一管理。""" + PinnedH2DBatchCopier.copy_async(dst_gpu_tensors, src_pinned_tensors) class cudaIpcMemHandle_t(ctypes.Structure): diff --git a/lightllm/common/kv_cache_mem_manager/allocator.py b/lightllm/common/kv_cache_mem_manager/allocator.py index 1331943b1a..839aa49a33 100644 --- a/lightllm/common/kv_cache_mem_manager/allocator.py +++ b/lightllm/common/kv_cache_mem_manager/allocator.py @@ -8,20 +8,13 @@ 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 + """内部按页维护空闲索引，对外的容量和申请、释放参数均以 token 为单位。""" - 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 + def __init__(self, size: int, page_size: int) -> None: + assert page_size >= 1, "page_size must be positive" + self.page_size = page_size + self.page_offsets = torch.arange(page_size, dtype=torch.int32, device="cpu", requires_grad=False) + self._init_buffers(size) rank_in_node = get_current_rank_in_node() # 用共享内存进行共享，router 模块读取进行精确的调度估计；基础层会统一添加服务前缀以防止实例冲突。 @@ -29,75 +22,87 @@ def __init__(self, size: int) -> None: 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: + def alloc(self, need_size: int) -> torch.Tensor: + assert need_size >= 0 and need_size % self.page_size == 0, "KV cache allocations must cover complete pages" + need_page_num = need_size // self.page_size + if need_page_num > 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 + page_start = self.mark_start + page_end = page_start + need_page_num + dst = self._alloc_return_tensor(need_size) + page_starts = self.mem_state[page_start:page_end] + if self.page_size == 1: + dst.copy_(page_starts) + else: + # 将每个 page 的页首索引展开为页内连续的 token 级别索引，直接写入 dst。 + torch.add(page_starts[:, None], self.page_offsets, out=dst.view(-1, self.page_size)) + self.mark_start = page_end self.can_use_mem_size -= need_size self.shared_can_use_token_num.set_value(self.can_use_mem_size) - - # 利用缓冲区返回，避免异步情况下的内存竞争 - if self._return_start + need_size > self._mem_state_return.shape[0]: - 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 + return dst 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)}" - + """接收按完整页展开的 token 索引，每页仅回收一个页首索引。""" + free_token_num = len(free_index) + assert free_token_num % self.page_size == 0, "KV cache releases must cover complete pages" + free_page_num = free_token_num // self.page_size + page_end = self.mark_start + page_start = page_end - free_page_num + assert page_start >= 0, f"error free state start: {self.mark_start} free len {len(free_index)}" + + page_starts = free_index[:: self.page_size] if isinstance(free_index, list): - self.mem_state.numpy()[start:end] = free_index + self.mem_state.numpy()[page_start:page_end] = page_starts else: - # 从 gpu 到 cpu 的拷贝操作是流内阻塞操作 - self.mem_state[start:end] = free_index - - self.mark_start -= len(free_index) + # GPU 输入在默认流按页取首索引后阻塞 D2H；返回时 CPU 空闲页表已可重新分配。 + self.mem_state[page_start:page_end].copy_(page_starts) - self.can_use_mem_size += len(free_index) + self.mark_start = page_start + self.can_use_mem_size += free_token_num self.shared_can_use_token_num.set_value(self.can_use_mem_size) - if self.can_use_mem_size == len(self.mem_state): + if self.can_use_mem_size == self.size: 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.mem_state = torch.arange(0, self.size, self.page_size, out=self.mem_state, requires_grad=False) self.mark_start = 0 - self.mark_end = len(self.mem_state) - self.can_use_mem_size = len(self.mem_state) + self.mark_end = self.mem_state.numel() + 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 - """ - self.size = new_size + """仅供测试：按新的 token 容量重新初始化空闲页表。""" + self._init_buffers(new_size) + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + + def _init_buffers(self, size: int) -> None: + assert size >= 0 and size % self.page_size == 0, "KV cache size must cover complete pages" + self.size = size + # 每项保存一个物理页的首 token 索引；[mark_start, mark_end) 是空闲页区间。 self.mem_state = torch.arange( - 0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True + 0, size, self.page_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.mark_end = self.mem_state.numel() + self._return_buffer = None + self.can_use_mem_size = size + + def _alloc_return_tensor(self, need_size: int) -> torch.Tensor: + """返回后续申请不会覆写的 pinned 区间，仅保留 slab 尚未使用的尾部。""" + buffer_token_num = 1024 * 1024 + if need_size == 0 or need_size > buffer_token_num: + return torch.empty(need_size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True) + if self._return_buffer is None or self._return_buffer.numel() < need_size: + self._return_buffer = torch.empty( + buffer_token_num, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True + ) + # 只保留未使用的尾部，已返回的切片不会被后续分配覆盖。 + dst = self._return_buffer[:need_size] + self._return_buffer = self._return_buffer[need_size:] + return dst diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 38c870ce4e..1f7dd600e9 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -47,7 +47,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, page_size=self.page_size) self._init_buffers( self.size, diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index 5da1b207e3..f334c24a00 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -49,15 +49,20 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana # 的那个batch size 进行运行，所有 padding 的请求都会使用预留的这个请求管理 id 进行处理 # 这样让 DP 的实现更为简化一些。 self.req_list = _ReqLinkedList(max_request_num) - page_size = get_env_start_args().page_size - max_sequence_length = (max_sequence_length + page_size - 1) // page_size * page_size + args = get_env_start_args() + page_size = args.page_size + self.page_size = page_size + aligned_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" - ) - # Allocation and cache metadata use this CPU mirror; inference reads the GPU table. - self.req_to_token_indexs_cpu = torch.zeros( - (max_request_num + 1, max_sequence_length), dtype=torch.int32, device="cpu", pin_memory=True + (max_request_num + 1, aligned_sequence_length), dtype=torch.int32, device="cuda" ) + self.req_token_index_mode = args.req_token_index_mode + if self.req_token_index_mode == "pin_mem_cache": + logger.info("请求 token 索引模式 pin_mem_cache：按请求申请 pinned CPU 缓存，索引更新批量上传 GPU") + elif self.req_token_index_mode == "pure_gpu": + logger.info("请求 token 索引模式 pure_gpu：直接维护 GPU 表，CPU 读取索引时执行 D2H 同步") + else: + raise ValueError(f"未知请求 token 索引模式：{self.req_token_index_mode}") self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num) self.max_request_num = max_request_num self.HOLD_REQUEST_ID = max_request_num @@ -79,14 +84,13 @@ 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_cpu[self.HOLD_REQUEST_ID] - hold_page = torch.tensor( - self.mem_manager.HOLD_TOKEN_MEMINDEXES, - dtype=hold_row.dtype, - device=hold_row.device, + assert self.mem_manager.page_size == self.page_size + hold_row = self.req_to_token_indexs[self.HOLD_REQUEST_ID] + token_offsets = torch.arange( + hold_row.numel(), dtype=hold_row.dtype, device=hold_row.device, requires_grad=False ) - hold_row.view(-1, self.mem_manager.page_size).copy_(hold_page) - self.req_to_token_indexs[self.HOLD_REQUEST_ID].copy_(hold_row) + hold_page_start = self.mem_manager.HOLD_TOKEN_MEMINDEXES[0] + hold_row.copy_(token_offsets.remainder(self.page_size).add_(hold_page_start)) return def alloc(self): diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index d4c6504cce..1332156db2 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -497,6 +497,17 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""KV cache page size in tokens. Values greater than 1 make each request reserve page-aligned contiguous KV slots and make paged attention/cache reuse operate on full pages.""", ) + parser.add_argument( + "--req_token_index_mode", + type=str, + choices=["pin_mem_cache", "pure_gpu"], + default="pin_mem_cache", + help="""Storage mode for request token indexes. pin_mem_cache (default): allocates a pinned CPU cache + per request, allowing CPU reads without D2H transfers and batched GPU uploads, at the cost of pinned + memory usage and upload progress tracking. pure_gpu: maintains the GPU request table directly without + a persistent CPU index cache; CPU reads require D2H transfers and may incur synchronization waits. + Both modes retain the same GPU request table and have approximately the same GPU memory usage.""", + ) parser.add_argument( "--vit_att_backend", type=str, diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index ae84204fec..967b83c479 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -176,6 +176,7 @@ class StartArgs: default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "flashinfer"]} ) page_size: int = field(default=1) + req_token_index_mode: str = field(default="pin_mem_cache", metadata={"choices": ["pin_mem_cache", "pure_gpu"]}) vit_att_backend: List[str] = field( default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "sdpa", "xformers"]} ) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index a40823b25f..5d1f377245 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -24,6 +24,7 @@ from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient from lightllm.server.multi_level_kv_cache import CachePlacementController, CacheTier from lightllm.server.router.model_infer.infer_req_ext import FinalTokenMetadataExt, PromptSelectedLogprobsExt +from lightllm.server.router.model_infer.infer_req_ext1 import ReqTokenIndex, create_req_token_index if TYPE_CHECKING: from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend @@ -128,7 +129,7 @@ def add_reqs(self, requests: List[Tuple[int, int, Any, int]], init_prefix_cache: def free_a_req_mem(self, free_token_index: List, req: "InferReq"): if self.radix_cache is None: - free_token_index.append(self.req_manager.req_to_token_indexs_cpu[req.req_idx][0 : req.hold_kv_len]) + free_token_index.append(req.token_index.get_free_token_indexes(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) else: @@ -139,13 +140,14 @@ def free_a_req_mem(self, free_token_index: List, req: "InferReq"): assert len(req.hybrid_len_to_big_page_id) == 0 req.cur_kv_len = 0 req.hold_kv_len = 0 - req.move_to_gpu_index_len = 0 + # 待释放列表已持有 CPU 快照或 GPU 行视图；reset 清理扩展状态，不改写 GPU 表。 + req.token_index.reset() 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"): 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_cpu[req.req_idx][shared_kv_len : req.hold_kv_len]) + free_token_index.append(req.token_index.get_free_token_indexes(shared_kv_len, req.hold_kv_len)) if self.is_hybrid_att_model: # 释放请求尾部尚未移交给 radix cache 的 hybrid attention 小页状态。 @@ -167,16 +169,16 @@ def _full_att_free_req(self, free_token_index: 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_cpu[req.req_idx][cache_kv_len : req.hold_kv_len]) + # 先收集 [cache_kv_len, hold_kv_len) 待释放区间，再处理可缓存前缀的插入和去重。 + free_token_index.append(req.token_index.get_free_token_indexes(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") - # The radix cache owns its indexes after this request row is reused. - value = self.req_manager.req_to_token_indexs_cpu[req.req_idx][:cache_kv_len].clone() + # radix cache 跨请求槽位复用持有索引，因此使用独立 CPU 快照。 + value = req.token_index.get_cpu_token_indexes(0, cache_kv_len) prefix_len, _ = self.radix_cache.insert(key, value) 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_cpu[req.req_idx][old_prefix_len:prefix_len]) + free_token_index.append(req.token_index.get_free_token_indexes(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) @@ -203,13 +205,11 @@ 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_cpu[req.req_idx][req.hybrid_cache_len : req.hold_kv_len] - ) + free_token_index.append(req.token_index.get_free_token_indexes(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_cpu[req.req_idx][: req.cur_kv_len].clone() + value = req.token_index.get_cpu_token_indexes(0, req.cur_kv_len) 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 @@ -222,7 +222,7 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): 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_cpu[req.req_idx][old_prefix_len:prefix_len]) + free_token_index.append(req.token_index.get_free_token_indexes(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) @@ -230,15 +230,13 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): 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_cpu[req.req_idx][tail_big_page_token_num : req.hold_kv_len] - ) + free_token_index.append(req.token_index.get_free_token_indexes(tail_big_page_token_num, req.hold_kv_len)) req.cur_kv_len = tail_big_page_token_num 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_cpu[req.req_idx][: req.cur_kv_len].clone() + value = req.token_index.get_cpu_token_indexes(0, req.cur_kv_len) 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] @@ -251,7 +249,7 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): 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_cpu[req.req_idx][old_prefix_len:prefix_len]) + free_token_index.append(req.token_index.get_free_token_indexes(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) @@ -259,9 +257,7 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): return if shared_kv_len <= req.cur_kv_len: - free_token_index.append( - self.req_manager.req_to_token_indexs_cpu[req.req_idx][shared_kv_len : req.hold_kv_len] - ) + free_token_index.append(req.token_index.get_free_token_indexes(shared_kv_len, req.hold_kv_len)) # 该分支不会把 prefill 阶段累积的 big page id 插入 radix cache（典型为 pause/abort # 在 prefill 跨过 big page 边界后、到达末尾前触发），需在此显式释放，避免泄漏。 @@ -325,6 +321,7 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool self.shm_req_manager.put_back_req_obj(req.shm_req) if free_token_index: + # 同一服务的索引模式一致；pure_gpu 在默认流合并 GPU 视图，随后才归还请求槽位。 free_token_index = torch.cat(free_token_index) self.req_manager.free(free_req_index, free_token_index) @@ -365,6 +362,7 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): logger.debug(f"infer paused req id {req.req_id}") if len(free_token_index) != 0: + # 恢复请求会重写请求表，必须先在默认流消费暂停时借用的 GPU 视图。 free_token_index = torch.cat(free_token_index) self.req_manager.free_token(free_token_index) return self @@ -540,14 +538,13 @@ def __init__( self.args = get_env_start_args() self.req_id = req_id self.req_idx = req_idx + self.req_manager = g_infer_context.req_manager self.shm_index = shm_index self.multimodal_params = multimodal_params self.vocab_size = vocab_size # cur_kv_len is the logical length already written; hold_kv_len is the # page-aligned physical capacity owned by this request. self.hold_kv_len = 0 - # Length whose index copies have been enqueued on the inference stream. - self.move_to_gpu_index_len: int = 0 # 请求需要被暂停 self.wait_pause = False @@ -599,6 +596,7 @@ def __init__( self.get_chuncked_input_token_ids = self.get_chuncked_input_token_ids_for_hybrid_att self._init_all_state() + self.token_index: ReqTokenIndex = create_req_token_index(self, index_stream=torch.cuda.current_stream()) self.generator = None if self.sampling_param.shm_param.seed != -1: @@ -657,8 +655,8 @@ def _match_radix_cache(self): if share_node is not None: self.shared_kv_node = share_node ready_cache_len = share_node.node_prefix_total_len - g_infer_context.req_manager.req_to_token_indexs_cpu[self.req_idx, :ready_cache_len].copy_(value_tensor) - self.move_to_gpu_index_len = 0 + # 新请求索引为空，暂停恢复前也已统一释放并 reset，可直接追加命中前缀。 + self.token_index.append_token_indexes(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 @@ -690,15 +688,13 @@ def _hybrid_match_radix_cache(self): key, block_hashs=block_hashs, update_refs=True ) if share_node is not None: - self.move_to_gpu_index_len = 0 + # 与普通前缀匹配一致：新请求或已完成暂停释放的请求，其索引状态为空。 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 - g_infer_context.req_manager.req_to_token_indexs_cpu[self.req_idx, :ready_cache_len].copy_( - value_tensor - ) + self.token_index.append_token_indexes(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 @@ -712,9 +708,7 @@ def _hybrid_match_radix_cache(self): # 如果 大页本质是被禁用的，可以直接使用小页的匹配结果 self.shared_kv_node = share_node ready_cache_len = share_node.node_prefix_total_len - g_infer_context.req_manager.req_to_token_indexs_cpu[self.req_idx, :ready_cache_len].copy_( - value_tensor - ) + self.token_index.append_token_indexes(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 @@ -734,12 +728,8 @@ def _hybrid_match_radix_cache(self): 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_cpu[ - self.req_idx, :cur_big_page_tokens - ].copy_(value_tensor[:cur_big_page_tokens]) - g_infer_context.req_manager.req_to_token_indexs_cpu[ - self.req_idx, cur_big_page_tokens:shared_kv_len - ].copy_(tail_mems) + self.token_index.append_token_indexes(value_tensor[:cur_big_page_tokens]) + self.token_index.append_token_indexes(tail_mems) # 将 对应的 value_tensors 中的 kv 数据 拷贝到 tail_mems 中对应的数据去 radix_cache.mem_manager.operator.copy_mem_to_mem( @@ -769,9 +759,7 @@ def _hybrid_match_radix_cache(self): # 大页匹配 self.shared_kv_node = share_node ready_cache_len = share_node.node_prefix_total_len - g_infer_context.req_manager.req_to_token_indexs_cpu[ - self.req_idx, :ready_cache_len - ].copy_(value_tensor[:ready_cache_len]) + self.token_index.append_token_indexes(value_tensor[: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 diff --git a/lightllm/server/router/model_infer/infer_req_ext.py b/lightllm/server/router/model_infer/infer_req_ext.py index dc22556e18..cf96d57c51 100644 --- a/lightllm/server/router/model_infer/infer_req_ext.py +++ b/lightllm/server/router/model_infer/infer_req_ext.py @@ -148,12 +148,6 @@ class FinalTokenMetadataExt: def __init__(self, req: "InferReq") -> None: self._req = req - def _mem_indexes(self) -> torch.Tensor: - # 延迟导入，避免与 infer_batch 形成模块级循环依赖。 - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - return g_infer_context.req_manager.req_to_token_indexs_cpu[self._req.req_idx] - def collect_prompt_logprobs(self) -> Optional[Tuple[np.ndarray, np.ndarray]]: """从 PromptLogprobsCaptureManager 提取 ``prompt_logprobs>0`` 的 top-k。 @@ -171,7 +165,7 @@ def collect_prompt_logprobs(self) -> Optional[Tuple[np.ndarray, np.ndarray]]: return None mgr = PromptLogprobsCaptureManager.get_instance() - mem_indexes = self._mem_indexes()[: req.shm_req.input_len - 1] + mem_indexes = req.token_index.get_cpu_token_indexes(0, req.shm_req.input_len - 1) return mgr.extract(mem_indexes, topk) def collect_routed_experts(self) -> Optional[np.ndarray]: @@ -183,7 +177,7 @@ def collect_routed_experts(self) -> Optional[np.ndarray]: if capture_len <= 0: return None - mem_indexes = self._mem_indexes()[0:capture_len] + mem_indexes = req.token_index.get_cpu_token_indexes(0, capture_len) return MoeRouteInfoManager.get_instance().extract(mem_indexes) def dump(self) -> None: diff --git a/lightllm/server/router/model_infer/infer_req_ext1.py b/lightllm/server/router/model_infer/infer_req_ext1.py new file mode 100644 index 0000000000..a6f1cb1508 --- /dev/null +++ b/lightllm/server/router/model_infer/infer_req_ext1.py @@ -0,0 +1,215 @@ +"""InferReq 的 token 索引扩展，仅保存请求槽位和索引状态。""" + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Optional + +import torch + +if TYPE_CHECKING: + from lightllm.server.router.model_infer.infer_batch import InferReq + + +class ReqTokenIndex(ABC): + """单个请求的 token 索引，使用调用方指定的流维护；区间右端点不包含在内。 + + 调用方保证 token 区间有效，并在构造时传入统一的索引流；在线服务在请求构造入口传入当前流。 + GPU 索引读写固定在索引流，返回 GPU 视图前让调用流等待，Helper 不维护逐请求 CUDA 事件。 + """ + + def __init__(self, infer_req: "InferReq", index_stream: torch.cuda.Stream): + # 仅保存固定槽位编号和 GPU 行视图，避免反向持有 InferReq 形成循环引用。 + self.req_idx = infer_req.req_idx + self.gpu_indexes = infer_req.req_manager.req_to_token_indexs[self.req_idx] + # 索引流由外部选择，后续 GPU 索引操作不随调用流切换。 + self._index_stream = index_stream + # 当前连续记录的 token 索引长度，包含页内预留槽位； + # CPU 扩容时复制 [0, token_index_len)，它不是缓存容量，也不是已算好 KV 的长度。 + self.token_index_len = 0 + # [0, gpu_index_submitted_len) 的最新索引写入已提交，后续索引流操作可按序使用。 + # 前缀重建前由 reset 清零；它不是累计上传量，也不表示 GPU 传输已经完成。 + self.gpu_index_submitted_len = 0 + + @abstractmethod + def append_token_indexes(self, mem_indexes: torch.Tensor): + """将 token 索引追加到当前记录的末尾，追加长度由 mem_indexes.numel() 决定。 + + 非空输入必须是 CPU Tensor；CPU 缓存模式直接复制到自己的 pinned 缓存。 + pure_gpu 按需 pin 后异步上传；已 pinned 的输入在 H2D 完成前不得改写，但可释放引用。 + 前缀替换先 reset，再按顺序追加重建。 + """ + pass + + @abstractmethod + def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + """返回独立 CPU 快照，允许缓存或传输任务跨请求生命周期持有。""" + pass + + @abstractmethod + def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + """在索引流准备 GPU 表视图；跨请求槽位复用持有时需另取快照。""" + pass + + @abstractmethod + def get_upload_indexes(self) -> Optional[torch.Tensor]: + """借用 [gpu_index_submitted_len, token_index_len) 的 CPU 视图，不更新提交进度。 + + 调用方可直接批量上传；缓存只追加，reset/扩容换用新存储，上传期间已有区间不变。 + GPU 模式返回 None。 + """ + pass + + @abstractmethod + def get_free_token_indexes(self, start: int, end: int) -> torch.Tensor: + """返回仅用于释放的 token 索引，可能是 CPU 或 GPU Tensor。 + + pin_mem_cache 返回 CPU 快照；pure_gpu 返回借用的 GPU 请求表视图。 + + GPU 视图借用请求表，返回前让调用流等待索引流。重写请求槽位前，须在同一流 + 按序消费，或让索引流等待消费流。此接口不提供跨槽位复用的快照语义， + 缓存和元数据仍使用 CPU getter。 + """ + pass + + @abstractmethod + def reset(self): + """重置有效索引长度和 GPU 提交进度，释放 CPU 缓存引用。""" + pass + + +class PinMemCacheReqTokenIndex(ReqTokenIndex): + """按请求按需申请完整 token 级 pinned CPU 缓存，GPU 更新可批量提交。""" + + def __init__(self, infer_req: "InferReq", index_stream: torch.cuda.Stream): + super().__init__(infer_req, index_stream) + self.input_len = infer_req.shm_req.input_len + self._cpu_buffer: Optional[torch.Tensor] = None + + def append_token_indexes(self, mem_indexes: torch.Tensor): + """在 CPU 索引末尾追加，GPU 提交进度由上传流程维护。""" + append_len = mem_indexes.numel() + if append_len == 0: + return + assert mem_indexes.device.type == "cpu", "mem_indexes 必须是 CPU Tensor" + start = self.token_index_len + end = start + append_len + self._reserve(end) + self._cpu_buffer[start:end].copy_(mem_indexes) + self.token_index_len = end + + def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + if start == end: + return torch.empty(0, dtype=self.gpu_indexes.dtype, device="cpu", requires_grad=False) + return self._cpu_buffer[start:end].clone() + + def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + """在索引流提交全部待上传索引，让调用流承接依赖后返回 GPU 视图。""" + assert end <= self.token_index_len, "读取范围超出已记录的 token 索引" + current_stream = torch.cuda.current_stream(self._index_stream.device) + with torch.cuda.stream(self._index_stream): + if self.gpu_index_submitted_len < end: + source_indexes = self.get_upload_indexes() + upload_end = self.token_index_len + # 缓存只追加，reset 会丢弃旧 buffer，已有上传区间不会被 CPU 重写。 + # copy_ 向 PyTorch pinned allocator 登记流使用，即使 reset 释放引用， + # 旧内存也会等 H2D 完成后才允许回收复用，因此可以直接上传缓存视图。 + self.gpu_indexes[self.gpu_index_submitted_len : upload_end].copy_(source_indexes, non_blocking=True) + self.gpu_index_submitted_len = upload_end + # 没有新上传时，已有写入也可能尚未完成；调用流仍需等待索引流。 + if current_stream != self._index_stream: + current_stream.wait_stream(self._index_stream) + return self.gpu_indexes[start:end] + + def get_upload_indexes(self) -> torch.Tensor: + if self._cpu_buffer is None: + return torch.empty(0, dtype=self.gpu_indexes.dtype, device="cpu", pin_memory=True, requires_grad=False) + return self._cpu_buffer[self.gpu_index_submitted_len : self.token_index_len] + + def get_free_token_indexes(self, start: int, end: int) -> torch.Tensor: + """返回用于释放的 CPU token 索引快照。""" + return self.get_cpu_token_indexes(start, end) + + def reset(self): + # 丢弃旧缓存引用，下一次追加重新申请，避免 CPU 覆写仍被在途 H2D 读取的内存。 + # 旧存储的回收时机由 PyTorch pinned allocator 跟踪，不需要在这里同步 CUDA。 + self.token_index_len = self.gpu_index_submitted_len = 0 + self._cpu_buffer = None + + def _reserve(self, end: int): + """将 CPU 缓存扩至至少 end 个 token，只复制已有的有效索引。""" + old_capacity = 0 if self._cpu_buffer is None else self._cpu_buffer.numel() + if old_capacity >= end: + return + if old_capacity == 0: + # 首次按完整输入长度预留，同时覆盖本次写入需要的 token 区间。 + reserve_len = max(self.input_len, end) + else: + # 容量不足时在所需长度上额外预留 2K token，减少后续 decode 扩容。 + reserve_len = end + 2 * 1024 + # 按 1K token 向上对齐，最多申请到 GPU 请求表允许的长度。 + new_capacity = min(self.gpu_indexes.numel(), (reserve_len + 1023) // 1024 * 1024) + new_buffer = torch.empty( + new_capacity, dtype=self.gpu_indexes.dtype, device="cpu", pin_memory=True, requires_grad=False + ) + if self.token_index_len: + new_buffer[: self.token_index_len].copy_(self._cpu_buffer[: self.token_index_len]) + self._cpu_buffer = new_buffer + + +class PureGpuReqTokenIndex(ReqTokenIndex): + """直接维护 GPU 请求表，不持有 CPU 索引缓存；CPU 读取会等待 D2H 完成。""" + + def append_token_indexes(self, mem_indexes: torch.Tensor): + append_len = mem_indexes.numel() + if append_len == 0: + return + assert mem_indexes.device.type == "cpu", "mem_indexes 必须是 CPU Tensor" + # 普通 CPU 输入在这里转为独立 pinned storage；已有 pinned 输入直接复用。 + if not mem_indexes.is_pinned(): + mem_indexes = mem_indexes.pin_memory() + start = self.token_index_len + end = start + append_len + # copy_ 会向 PyTorch pinned allocator 登记流使用，H2D 完成前不会回收源内存。 + # 普通 CPU 输入已形成独立 pinned 副本；直接复用的 pinned 输入在传输完成前不能改写。 + current_stream = torch.cuda.current_stream(self._index_stream.device) + with torch.cuda.stream(self._index_stream): + self.gpu_indexes[start:end].copy_(mem_indexes, non_blocking=True) + if current_stream != self._index_stream: + current_stream.wait_stream(self._index_stream) + self.token_index_len = end + self.gpu_index_submitted_len = end + + def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + # 阻塞到 D2H 完成，GPU 索引拷回后得到独立 CPU 快照。 + with torch.cuda.stream(self._index_stream): + return self.gpu_indexes[start:end].detach().cpu() + + def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + # append 已在索引流提交写入，读取只校验范围，不推进提交进度。 + assert end <= self.gpu_index_submitted_len, "读取范围超出已提交的 GPU 索引" + current_stream = torch.cuda.current_stream(self._index_stream.device) + if current_stream != self._index_stream: + current_stream.wait_stream(self._index_stream) + return self.gpu_indexes[start:end] + + def get_upload_indexes(self) -> None: + return None + + def get_free_token_indexes(self, start: int, end: int) -> torch.Tensor: + """返回用于释放的 GPU token 索引视图，让调用流先等待索引写入。""" + # 复用 GPU getter 的范围校验和流等待，避免跨流释放读取尚未就绪的索引。 + # allocator.free 会按页取首索引并阻塞 D2H；借用视图仍须在槽位重写前消费。 + return self.get_gpu_token_indexes(start, end) + + def reset(self): + # 保留 GPU 行内容，统一释放流程还需在槽位重写前消费借用的索引视图。 + self.token_index_len = self.gpu_index_submitted_len = 0 + + +def create_req_token_index(infer_req: "InferReq", index_stream: torch.cuda.Stream) -> ReqTokenIndex: + """从请求参数选择实现；索引流由调用方显式指定，内部不选择默认流或当前流。""" + mode = infer_req.args.req_token_index_mode + if mode == "pin_mem_cache": + return PinMemCacheReqTokenIndex(infer_req, index_stream) + if mode == "pure_gpu": + return PureGpuReqTokenIndex(infer_req, index_stream) + raise ValueError(f"未知请求 token 索引模式：{mode}，可选 pin_mem_cache / pure_gpu") 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 4d1143dcd4..ee4badd471 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -657,34 +657,43 @@ def _alloc_req_kv_mem( 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_cpu[req_obj.req_idx, old_hold_kv_len:new_hold_kv_len].copy_( - mem_indexes - ) - req_obj.hold_kv_len = new_hold_kv_len + req_obj.token_index.append_token_indexes(mem_indexes) + req_obj.hold_kv_len += alloc_token_num return mem_indexes def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): - """Collect pending index tails, upload them in one batch, then advance request lengths.""" - req_manager = g_infer_context.req_manager - cpu_table, gpu_table = req_manager.req_to_token_indexs_cpu, req_manager.req_to_token_indexs - pending_copies = [] - src_tensors, dst_tensors = [], [] + """在本批请求共用的索引流上批量上传，调用流按需等待索引流。""" + # 第一阶段：空批直接返回，确认请求共用一个索引流，并记录当前调用流。 + if not run_reqs: + return + index_streams = {req.token_index._index_stream for req in run_reqs} + assert len(index_streams) == 1, "同一批请求必须使用同一个索引流" + index_stream = index_streams.pop() + current_stream = torch.cuda.current_stream(index_stream.device) + # 第二阶段：切出尚未提交的 CPU 源视图和 GPU 目标视图，记录本次提交终点。 + cpu_index_slices, gpu_index_slices = [], [] + pending_index_uploads = [] for req in run_reqs: - start, end = req.move_to_gpu_index_len, req.hold_kv_len - assert 0 <= start <= end <= cpu_table.shape[1] + token_index = req.token_index + start, end = token_index.gpu_index_submitted_len, token_index.token_index_len if start < end: - src_tensors.append(cpu_table[req.req_idx, start:end]) - dst_tensors.append(gpu_table[req.req_idx, start:end]) - pending_copies.append((req, end)) - if not pending_copies: - return - - cuda_memcpy_pinned_h2d_batch_async(dst_tensors, src_tensors) - for req, end in pending_copies: - req.move_to_gpu_index_len = end + source_indexes = token_index.get_upload_indexes() + assert source_indexes is not None, "GPU 索引提交进度与记录长度不一致，GPU 模式应在追加时完成提交" + cpu_index_slices.append(source_indexes) + gpu_index_slices.append(token_index.gpu_indexes[start:end]) + pending_index_uploads.append((token_index, end)) + # 第三阶段：在索引流上批量提交 H2D，默认流适配和 pinned 源保活由拷贝封装处理。 + if cpu_index_slices: + with torch.cuda.stream(index_stream): + cuda_memcpy_pinned_h2d_batch_async(gpu_index_slices, cpu_index_slices) + # 第四阶段：成功提交后再更新进度；提交完成不代表 GPU 传输已经完成。 + for token_index, end in pending_index_uploads: + token_index.gpu_index_submitted_len = end + # 第五阶段：调用流与索引流不同时建立等待依赖，保证后续 GPU 访问能读到更新。 + # 两种模式在没有新上传时，已提交的索引写入也可能尚未完成，仍需承接依赖。 + if current_stream != index_stream: + current_stream.wait_stream(index_stream) def _get_classed_reqs( self, @@ -711,7 +720,7 @@ def _get_classed_reqs( 4. prefill_reqs 需要进行prefill操作的请求 5. decode_reqs 需要进行decode操作的请求 - 返回前将两组请求尚未上传的 token 索引提交到推理流，并更新 move_to_gpu_index_len。 + 返回前按所选模式准备两组请求的 GPU token 索引，并更新 token_index.gpu_index_submitted_len。 """ # 定期对 radix cache 进行 merge，防止查询插入的操作效率下降 self._timer_merge_radix_tree() @@ -872,9 +881,7 @@ def _get_classed_reqs( decode_reqs = [] if prefill_reqs or decode_reqs: - # 使用已有推理流提交映射表更新，与前后推理操作保持顺序。 - with torch.cuda.stream(g_infer_context.get_overlap_stream()): - self._copy_req_token_indexes_to_gpu(prefill_reqs + decode_reqs) + self._copy_req_token_indexes_to_gpu(prefill_reqs + decode_reqs) return prefill_reqs, decode_reqs # 一些可以复用的通用功能函数 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 c5090f467f..1d8bf4b51e 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 @@ -129,6 +129,10 @@ def _diverse_copy( # 一些可以复用的通用功能函数 def _diverse_pre_post_handle(self, run_reqs: List[InferReq], is_chuncked_mode: bool) -> List[InferReqUpdatePack]: + # diverse 会把 master/slave 已有前缀重新指向共享 KV，属于旧索引重写。 + # 先在默认流等待上一轮模型读取结束，随后 CPU 缓存模式的批量上传、纯 GPU + # 模式的直接写入都沿默认流排序；普通请求仅追加新索引，不需要这一等待。 + torch.cuda.default_stream().wait_stream(g_infer_context.get_overlap_stream()) update_func_objs: List[InferReqUpdatePack] = [] # 通用状态预先填充 is_master_in_dp = self.is_master_in_dp @@ -180,13 +184,14 @@ def _diverse_pre_post_handle(self, run_reqs: List[InferReq], is_chuncked_mode: b def _master_req_to_radix_cache(self, master_req: InferReq): 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_cpu[master_req.req_idx][: master_req.cur_kv_len].clone() - prefix_len, new_shared_kv_node = self.radix_cache.insert(key, value) + # 一次性保存全部已持有索引。重建共享前缀后，页内预留的尾部仍由 master 持有， + # 不能随 reset 丢掉；独立 CPU 快照也避免缓存复用时覆盖尚未追加的尾部。 + token_indexes = master_req.token_index.get_cpu_token_indexes(0, master_req.hold_kv_len) + prefix_len, new_shared_kv_node = self.radix_cache.insert(key, token_indexes[: master_req.cur_kv_len]) 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_cpu[master_req.req_idx][old_prefix_len:prefix_len] - ) + # 释放在默认流同步消费；纯 GPU 模式可借用索引表视图，避免整段 token 索引 D2H。 + self.model.mem_manager.free(master_req.token_index.get_free_token_indexes(old_prefix_len, prefix_len)) # 将原有共享节点替换为新共享节点，新共享节点对应的长度为当前的cur_kv_len self.radix_cache.dec_node_ref_counter(master_req.shared_kv_node) @@ -198,8 +203,9 @@ 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_cpu[master_req.req_idx, : master_req.cur_kv_len].copy_(value) - master_req.move_to_gpu_index_len = 0 + master_req.token_index.reset() + master_req.token_index.append_token_indexes(value) + master_req.token_index.append_token_indexes(token_indexes[master_req.cur_kv_len :]) return def _copy_master_req_to_slave_req(self, slave_req: InferReq): @@ -211,12 +217,9 @@ 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_cpu[slave_req.req_idx, :kv_len].copy_( - self.model.req_manager.req_to_token_indexs_cpu[master_req.req_idx, :kv_len] - ) - slave_req.move_to_gpu_index_len = 0 - # TODO: 确认 CPU 镜像复制后是否需要同步当前流，以及对 diverse 并行流程的影响。 - # torch.cuda.current_stream().synchronize() + token_indexes = master_req.token_index.get_cpu_token_indexes(0, kv_len) + slave_req.token_index.reset() + slave_req.token_index.append_token_indexes(token_indexes) 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 3c4200a161..d7873584ff 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 @@ -23,9 +23,8 @@ def __init__(self, max_req_num: int, dp_size_in_node: int, backend): self.backend: DPChunkedPrefillBackend = backend self.max_req_num = max_req_num - self.copy_stream = torch.cuda.Stream() - # 0 代表 kv_len, 1 代表 radix_cache_len + # 每个请求、DP rank 对应两列：kv_len 和 req_idx。 self.shared_req_infos = ShmArray( name="dp_shared_req_infos", shape=(self.max_req_num, dp_size_in_node, 2), @@ -39,15 +38,11 @@ def fill_reqs_info(self, reqs: List[InferReq]): """ 填充请求的 kv 信息到共享内存中 """ - # 所有 TP rank 先发布 CPU 索引，确保其他 DP rank 读取 GPU 表时数据已经就绪。 - current_stream = torch.cuda.current_stream() - # 常规索引上传使用 overlap stream，本次上传需排在此前的索引更新之后。 - self.copy_stream.wait_stream(g_infer_context.get_overlap_stream()) - with torch.cuda.stream(self.copy_stream): - self.backend._copy_req_token_indexes_to_gpu(reqs) - current_stream.wait_stream(self.copy_stream) - # 同步 NCCL barrier 会等待当前流完成，保证其他 rank 读取到已上传的索引。 - # TODO: 分析传输流依赖及 NCCL barrier 等待对 DP shared 并行效率的影响。 + # 请求索引统一在默认流维护，原生 batch 的默认流适配由拷贝封装处理。 + # DP 发布复用统一入口，不另建索引流或逐请求事件。 + self.backend._copy_req_token_indexes_to_gpu(reqs) + # 同步 NCCL barrier 会等待默认流上排入的上传依赖完成；所有 TP rank + # 到达此屏障后，其他 DP rank 才能读取已发布的 GPU 索引和共享请求元数据。 dist.barrier(group=self.backend.node_nccl_group) if self.backend.is_master_in_dp: self.shared_req_infos.arr[0 : len(reqs), self.dp_rank_in_node, self._KV_LEN_INDEX] = [ @@ -132,7 +127,6 @@ def kv_trans(self, trans_tasks: List["TransTask"]): dp_size_in_node=self.backend.dp_size_in_node, rank_in_dp=self.backend.rank_in_dp, ) - self.backend.logger.info(f"dp_i {self.dp_rank_in_node} transfer kv tokens num: {len(mem_indexes_tensor)}") for trans_task in trans_tasks: trans_task.req.cur_kv_len += len(trans_task.mem_indexes) 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 51c95d134e..8aa9c37491 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 @@ -112,26 +112,17 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): # g_infer_context.get_overlap_stream().synchronize() mem_manager = self.backend.model.mem_manager - req_manager = self.backend.model.req_manager - page_indexes_cuda = torch.tensor(need_pages, dtype=torch.int32, device="cpu", pin_memory=True).cuda( non_blocking=True ) - # hybrid 页面加载必须按完整 page 处理，否则可能缺失恢复运行态所需的 checkpoint，所以 - # 这里需要进行pad操作，使操作的页面是完整的。 + # hybrid 加载区间必须覆盖完整缓存页，包含恢复运行态所需的 checkpoint。 _start = page_len_start_list[ready_page_num] _end = req.hold_kv_len assert 0 <= _start <= req.cur_kv_len <= _end, f"invalid pad range [{_start}, {_end}]" - # 借用已有 GPU 索引表，避免为加载临时分配索引显存；旧前缀和新分配的 - # 索引从 pinned CPU 镜像一次上传，copy_ 与 KV 加载在当前流上按序执行。 - # 风险：加载完成前须保持 CPU 索引不变、请求槽位有效，且不能有其他流 - # 冲突访问此区间；修改初始化或流调度时需重新检查依赖，不能只依赖写入相同值。 - # 此处不更新推理流的上传进度。 - mem_indexes_cuda = req_manager.req_to_token_indexs[req.req_idx, _start:_end] - mem_indexes_cuda.copy_( - req_manager.req_to_token_indexs_cpu[req.req_idx, _start:_end], non_blocking=True - ) + # 索引准备与 KV 加载都在默认流，按提交顺序即可保证索引就绪。 + # 纯 GPU 模式直接借用请求表；请求槽位在 KV 加载完成前保持有效。 + mem_indexes_cuda = req.token_index.get_gpu_token_indexes(_start, _end) assert ( len(mem_indexes_cuda) == page_len_list[len(page_list) - 1] - page_len_start_list[ready_page_num] @@ -283,16 +274,9 @@ def _start_kv_cache_offload_task( move_token_num = page_len_list[item_size - 1] assert req.cur_kv_len >= move_token_num - req_manager = self.backend.model.req_manager - # 借用已有 GPU 索引表作为 offload 输入，避免 .cuda() 额外分配临时索引显存。 - # 风险警告：此处跨流借用推理索引表，“写入相同值”不能替代同步保证。 - # 安全性依赖结束请求的延迟清理，确保此前的推理访问已结束；CPU 索引前缀在 - # 上传期间必须保持不变，卸载完成事件确认前不得释放 KV 或复用 req_idx。 - # copy_ 与 offload 在同一卸载流上按序执行；这里只借用表，不更新推理流的上传进度。 - # 后续若用于运行中请求、前缀重映射，或修改 abort/清理/流调度逻辑，必须重新 - # 检查跨流依赖，必要时增加事件等待或改用卸载流独立暂存区，请谨慎修改。 - token_indexes = req_manager.req_to_token_indexs[req.req_idx, :move_token_num] - token_indexes.copy_(req_manager.req_to_token_indexs_cpu[req.req_idx, :move_token_num], non_blocking=True) + # getter 在请求索引流准备数据，并让当前卸载流等待，无需调用方额外同步。 + token_indexes = req.token_index.get_gpu_token_indexes(0, move_token_num) + # 请求的 KV 和 req_idx 仍由卸载任务持有，直到下面的完成事件确认后才释放。 mem_manager = self.backend.model.mem_manager @@ -304,7 +288,7 @@ def _start_kv_cache_offload_task( req=req, ) - # 这个操作只是为了在offload 对应的cuda stream中，同步标记下对应的kv cache offload 操作已经完成， + # 多 DP rank 先在卸载流会合，再记录任务完成事件，供后续确认请求可释放。 if self.backend.dp_world_size > 1: dist.all_reduce(self.offload_sync_tensor, op=dist.ReduceOp.MAX, group=self.offload_sync_group) 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 51548ac16c..f0ef204c53 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 @@ -99,20 +99,12 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: f"(reason={req_obj.finish_status.get_finish_reason()}), kv transfer error" ) - # 提前释放有问题的 mem_index。cur_kv_len 只表示已经传输完成的 - # 逻辑长度，hold_kv_len 还包含最后一个页面中尚未使用的预留槽位。 + # 只有原共享前缀仍是可信 KV，失败尾部不能参与缓存插入。 + # 保留 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_cpu[ - req_obj.req_idx, old_prefix_len : req_obj.hold_kv_len - ].clone() - self.model.mem_manager.free(mem_indexes) - req_obj.cur_kv_len = old_prefix_len - req_obj.hold_kv_len = old_prefix_len - req_obj.move_to_gpu_index_len = min(req_obj.move_to_gpu_index_len, old_prefix_len) - if self.is_master_in_dp: - req_obj.shm_req.shm_cur_kv_len = req_obj.cur_kv_len + req_obj.cur_kv_len = old_prefix_len + if self.is_master_in_dp: + req_obj.shm_req.shm_cur_kv_len = req_obj.cur_kv_len ans_list.append(req_obj) return ans_list 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 4357db5359..451afe4fa9 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 @@ -122,9 +122,7 @@ 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_cpu[ - req_obj.req_idx, kv_start_index:kv_end_index - ].tolist() + mem_indexes = req_obj.token_index.get_cpu_token_indexes(kv_start_index, kv_end_index).tolist() req_idx = None elif page_kind == "att_state": mem_indexes = [] diff --git a/test/benchmark/benchmark_event_synchronize_gil.py b/test/benchmark/benchmark_event_synchronize_gil.py new file mode 100644 index 0000000000..86eee4b772 --- /dev/null +++ b/test/benchmark/benchmark_event_synchronize_gil.py @@ -0,0 +1,233 @@ +"""实测 CUDA Event 等待是否释放 GIL，以及服务保活线程等待时主线程能否执行 Python。 + +示例：CUDA_VISIBLE_DEVICES=6 PYTHONPATH=. python test/benchmark/benchmark_event_synchronize_gil.py +所有等待都在 cuda_wrapper 的同一个后台线程池执行，主线程只做 Python 计数和时间采样。 +PyDLL.usleep 是持有 GIL 的负对照，CDLL.usleep 是释放 GIL 的正对照。 +""" + +import argparse +import ctypes +from datetime import datetime, timezone +import json +import os +from pathlib import Path +import statistics +import sys +import threading +import time + + +METHODS = ( + "pydll_holds_gil", + "cdll_releases_gil", + "event_blocking_true", + "event_blocking_false", + "service_release_worker", +) +COUNTER_CHUNK = 4096 + + +def calibrate_delay(torch, stream, wait_ms): + """先预热，再按实际 GPU 时间估算约 wait_ms 的延时，避免依赖空闲时的 SM 频率。""" + reference_cycles = 50_000_000 + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(stream): + torch.cuda._sleep(reference_cycles) + stream.synchronize() + durations = [] + for _ in range(3): + with torch.cuda.stream(stream): + start.record(stream) + torch.cuda._sleep(reference_cycles) + end.record(stream) + end.synchronize() + durations.append(start.elapsed_time(end)) + return int(reference_cycles * wait_ms / statistics.median(durations)), durations + + +def timed_wait(wait, state): + state["thread"] = threading.current_thread().name + state["start_ns"] = time.perf_counter_ns() + try: + wait() + finally: + state["end_ns"] = time.perf_counter_ns() + state["finished"] = True + + +def count_python_work(state): + # 此循环不调用 torch、CUDA、sleep、锁或 Future；只有 Python 整数加法和时间采样。 + # 每 4096 次采样一次，既能验证持续推进，也避免保存过多时间戳。 + samples = [] + counter = 0 + while not state["finished"]: + for _ in range(COUNTER_CHUNK): + counter += 1 + samples.append((time.perf_counter_ns(), counter)) + return samples + + +def run_trial(torch, wrapper, stream, host_waits, cycles, method, repeat, wait_ms, margin_ms): + event = None + timing_start, timing_end = None, None + sources = None + pending_before_wait = None + if method in host_waits: + wait = lambda: host_waits[method](int(wait_ms * 1000)) + else: + # 被测 Event 与服务保持相同的 enable_timing=False 配置；计时用独立事件。 + event = torch.cuda.Event(blocking=method != "event_blocking_false") + timing_start = torch.cuda.Event(enable_timing=True) + timing_end = torch.cuda.Event(enable_timing=True) + if method == "service_release_worker": + sources = [torch.empty(16, dtype=torch.int32, device="cpu", pin_memory=True, requires_grad=False)] + with torch.cuda.stream(stream): + timing_start.record(stream) + torch.cuda._sleep(cycles) + event.record(stream) + timing_end.record(stream) + pending_before_wait = not event.query() + if method == "service_release_worker": + wait = lambda: wrapper.PinnedH2DBatchCopier._wait_and_release_pinned_sources(event, sources) + else: + wait = event.synchronize + + state = {"finished": False} + future = wrapper.PinnedH2DBatchCopier._get_release_pool().submit(timed_wait, wait, state) + samples = count_python_work(state) + future.result(timeout=10) + gpu_delay_ms = None + if event is not None: + stream.synchronize() + gpu_delay_ms = timing_start.elapsed_time(timing_end) + + # 严格截取等待内部区间，排除进入原生调用之前和返回 Python 之后的调度窗口。 + margin_ns = int(margin_ms * 1_000_000) + window_start = state["start_ns"] + margin_ns + window_end = state["end_ns"] - margin_ns + interior_samples = [(stamp, count) for stamp, count in samples if window_start < stamp < window_end] + increments = interior_samples[-1][1] - interior_samples[0][1] if len(interior_samples) > 1 else 0 + progress_span_ms = ( + (interior_samples[-1][0] - interior_samples[0][0]) / 1_000_000 if len(interior_samples) > 1 else 0.0 + ) + duration_ms = (state["end_ns"] - state["start_ns"]) / 1_000_000 + # 等待不足目标的 60% 或内部窗口不足 150ms 时不接受结果，避免过短等待误判 GIL 行为。 + valid_wait = duration_ms >= wait_ms * 0.6 and window_end - window_start >= 150_000_000 + expects_gil_release = method != "pydll_holds_gil" + if expects_gil_release: + # 正对照要求内部窗口持续推进至少 100ms；负对照要求内部完全没有 Python 采样。 + progress_matches = increments > 0 and progress_span_ms >= 100 + else: + progress_matches = len(interior_samples) == 0 + passed = valid_wait and progress_matches and pending_before_wait is not False + if sources is not None: + passed = passed and len(sources) == 0 + + return { + "method": method, + "repeat": repeat, + "wait_ms": duration_ms, + "gpu_delay_ms": gpu_delay_ms, + "event_pending_before_wait": pending_before_wait, + "interior_window_ms": max(0.0, (window_end - window_start) / 1_000_000), + "interior_samples": len(interior_samples), + "python_increments": increments, + "progress_span_ms": progress_span_ms, + "worker_thread": state["thread"], + "sources_cleared": len(sources) == 0 if sources is not None else None, + "valid_wait": valid_wait, + "passed": passed, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--wait-ms", type=float, default=500) + parser.add_argument("--margin-ms", type=float, default=50) + args = parser.parse_args() + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not visible or "," in visible or visible == "3": + parser.error("请通过 CUDA_VISIBLE_DEVICES 显式选择一张空闲 GPU，不能使用 3 号卡") + if args.repeats < 1 or args.wait_ms < 400 or not 0 < args.margin_ms < args.wait_ms / 4: + parser.error("重复数须大于 0，等待至少 400ms，边缘排除区间须介于 0 和等待时间的四分之一之间") + + import torch + from lightllm.common import cuda_wrapper + + torch.set_num_threads(1) + torch.cuda.set_device(0) + stream = torch.cuda.Stream() + # ctypes 的两种加载方式调用同一个 libc 函数，唯一关注的区别是是否在原生等待中释放 GIL。 + libraries = {"pydll_holds_gil": ctypes.PyDLL(None), "cdll_releases_gil": ctypes.CDLL(None)} + host_waits = {} + for method, library in libraries.items(): + function = library.usleep + function.argtypes = [ctypes.c_uint] + function.restype = ctypes.c_int + host_waits[method] = function + cuda_wrapper.PinnedH2DBatchCopier._get_release_pool().submit(lambda: None).result(timeout=10) + cycles, calibration_ms = calibrate_delay(torch, stream, args.wait_ms) + root = args.output_dir or Path("/tmp") / ( + "event_synchronize_gil_" + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + ) + root.mkdir(parents=True, exist_ok=False) + print(f"结果目录：{root}", flush=True) + rows = [] + try: + for repeat in range(args.repeats): + for method in METHODS: + rows.append( + run_trial( + torch, cuda_wrapper, stream, host_waits, cycles, method, repeat, args.wait_ms, args.margin_ms + ) + ) + print(f"已完成第 {repeat + 1}/{args.repeats} 轮", flush=True) + finally: + stream.synchronize() + cuda_wrapper.PinnedH2DBatchCopier._get_release_pool().submit(lambda: None).result(timeout=10) + summaries = [] + for method in METHODS: + cases = [row for row in rows if row["method"] == method] + summary = { + "method": method, + "passed": sum(row["passed"] for row in cases), + "trials": len(cases), + "wait_ms_median": statistics.median(row["wait_ms"] for row in cases), + "python_increments_min": min(row["python_increments"] for row in cases), + "python_increments_median": statistics.median(row["python_increments"] for row in cases), + "python_increments_max": max(row["python_increments"] for row in cases), + "progress_span_ms_min": min(row["progress_span_ms"] for row in cases), + } + summaries.append(summary) + print(json.dumps(summary, ensure_ascii=False), flush=True) + evidence = { + "environment": { + "python": sys.version.split()[0], + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cuda_visible_devices": visible, + "gpu": torch.cuda.get_device_name(0), + "python_switch_interval_s": sys.getswitchinterval(), + }, + "config": { + "repeats": args.repeats, + "target_wait_ms": args.wait_ms, + "margin_ms": args.margin_ms, + "counter_chunk": COUNTER_CHUNK, + "gpu_delay_cycles": cycles, + "calibration_ms": calibration_ms, + }, + "summary": summaries, + "trials": rows, + "passed": all(row["passed"] for row in rows), + } + (root / "result.json").write_text(json.dumps(evidence, ensure_ascii=False, indent=2) + "\n") + if not evidence["passed"]: + raise SystemExit("实验判据未全部满足，请查看 result.json 中的有效等待时长与内部窗口进度") + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/benchmark_pinned_empty_h2d.py b/test/benchmark/benchmark_pinned_empty_h2d.py new file mode 100644 index 0000000000..537c2f12e3 --- /dev/null +++ b/test/benchmark/benchmark_pinned_empty_h2d.py @@ -0,0 +1,248 @@ +"""实测空 H2D 能否像单元素 H2D 一样保护 raw batch 使用的 pinned storage。 + +示例：CUDA_VISIBLE_DEVICES=6 python test/benchmark/benchmark_pinned_empty_h2d.py +每种方法和规模在独立子进程运行，避免其他案例留下的 allocator 状态影响观察。 +只操作实验自身申请的内存，不启动模型、不修改服务的生命周期保护实现。 +""" + +import argparse +import csv +import json +import os +from pathlib import Path +import subprocess +import sys +from datetime import datetime, timezone + + +METHODS = ("none", "empty_copy", "empty_to", "one_element_copy", "empty_copy_keep_view") + + +def protect_source(target, source, method): + if method == "empty_copy": + target[:0].copy_(source[:0], non_blocking=True) + elif method == "empty_to": + # 与 empty_copy 一样丢弃临时视图和结果；只有 keep_view 对照延长源引用。 + source[:0].to(target.device, non_blocking=True) + elif method == "one_element_copy": + target[:1].copy_(source[:1], non_blocking=True) + elif method == "empty_copy_keep_view": + # 此对照有意保留 CPU 空视图：即使 numel 为 0，它仍持有整个原 storage。 + # 如果这组成功，不能据此认为空 H2D 向 allocator 登记了流。 + view = source[:0] + target[:0].copy_(view, non_blocking=True) + return view + return None + + +def submit_raw_batch(runtime, target, source, stream): + attrs = runtime.cudaMemcpyAttributes() + # 与服务相同的 Stream 模式：真正读取源数据的时间遵循上传流顺序， + # raw batch 不会替调用方向 PyTorch host allocator 登记源内存。 + attrs.srcAccessOrder = runtime.cudaMemcpySrcAccessOrder.cudaMemcpySrcAccessOrderStream + (error,) = runtime.cudaMemcpyBatchAsync( + [target.data_ptr()], + [source.data_ptr()], + [source.numel() * source.element_size()], + 1, + [attrs], + [0], + 1, + stream.cuda_stream, + ) + if error != runtime.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaMemcpyBatchAsync failed: {error}") + + +def run_trial(torch, runtime, stream, method, token_count, repeat, delay_cycles, attempts): + offset = 64 + allocation_size = token_count + offset * 2 + expected_value, overwritten_value = repeat + 100, -repeat - 1000 + owner = torch.full((allocation_size,), expected_value, dtype=torch.int32, pin_memory=True, requires_grad=False) + source = owner[offset : offset + token_count] + source_storage_ptr = owner.untyped_storage().data_ptr() + empty_view = source[:0] + empty_metadata = { + "numel": empty_view.numel(), + "data_ptr": empty_view.data_ptr(), + "shares_source_storage": empty_view.untyped_storage().data_ptr() == source_storage_ptr, + } + del empty_view + target = torch.empty(token_count, dtype=torch.int32, device="cuda", requires_grad=False) + stream.wait_stream(torch.cuda.default_stream()) + # 同一流先延时，再提交真实 raw H2D 和待验证的保护操作。 + # 这样 CPU 能在 GPU 尚未读取源数据时释放原 Tensor、触发 allocator 复用。 + with torch.cuda.stream(stream): + torch.cuda._sleep(delay_cycles) + submit_raw_batch(runtime, target, source, stream) + kept_view = protect_source(target, source, method) + del owner, source + + held_allocations = [] + reused_while_pending = False + reuse_seen = False + overwrite_finished_while_pending = False + pending_checks = 0 + try: + for _ in range(attempts): + if stream.query(): + break + replacement = torch.empty(allocation_size, dtype=torch.int32, pin_memory=True, requires_grad=False) + # 保留所有新分配块。若旧块被复用并改写，也让它活到上传流结束， + # 将风险局限为本实验的数据被覆盖，不让 DMA 访问已经归还驱动的内存。 + held_allocations.append(replacement) + still_pending = not stream.query() + pending_checks += int(still_pending) + if replacement.untyped_storage().data_ptr() == source_storage_ptr: + reuse_seen = True + reused_while_pending = still_pending + if still_pending: + replacement.fill_(overwritten_value) + overwrite_finished_while_pending = not stream.query() + break + finally: + # 包括异常路径也先等待本实验的上传流，再释放 replacement 和保留的空视图。 + stream.synchronize() + del kept_view + result = target.cpu() + matching_tokens = int(result.eq(expected_value).sum().item()) + overwritten_tokens = int(result.eq(overwritten_value).sum().item()) + + reused_after_completion = None + if not reuse_seen: + # 正向对照还应证明上传完成后原分配块可以再次使用，而非被永久保留。 + reused_after_completion = False + for _ in range(attempts): + replacement = torch.empty(allocation_size, dtype=torch.int32, pin_memory=True, requires_grad=False) + held_allocations.append(replacement) + if replacement.untyped_storage().data_ptr() == source_storage_ptr: + reused_after_completion = True + break + return { + "repeat": repeat, + "method": method, + "token_count": token_count, + "payload_bytes": token_count * 4, + "source_storage_ptr": source_storage_ptr, + "empty_view": empty_metadata, + "pending_allocation_checks": pending_checks, + "source_reused_while_pending": reused_while_pending, + "overwrite_finished_while_pending": overwrite_finished_while_pending, + "matching_tokens": matching_tokens, + "overwritten_tokens": overwritten_tokens, + "output_correct": matching_tokens == token_count, + "source_reused_after_completion": reused_after_completion, + } + + +def run_worker(args): + import torch + from cuda.bindings import runtime + + torch.set_num_threads(1) + torch.cuda.set_device(0) + stream = torch.cuda.Stream() + # 首次初始化必须在延时试验前完成，不能让导入/驱动初始化耗时掩盖复用窗口。 + source = torch.ones(256, dtype=torch.int32, pin_memory=True, requires_grad=False) + target = torch.empty(256, dtype=torch.int32, device="cuda", requires_grad=False) + with torch.cuda.stream(stream): + torch.cuda._sleep(1) + submit_raw_batch(runtime, target, source, stream) + protect_source(target, source, args.method) + stream.synchronize() + del source, target + rows = [ + run_trial(torch, runtime, stream, args.method, args.token_count, repeat, args.delay_cycles, args.attempts) + for repeat in range(args.repeats) + ] + print(json.dumps({"torch_version": torch.__version__, "cuda_version": torch.version.cuda, "trials": rows})) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--sizes", type=int, nargs="+", default=[8192, 262144], help="展开后的 int32 token 数") + parser.add_argument("--delay-cycles", type=int, default=200_000_000) + parser.add_argument("--attempts", type=int, default=64) + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--method", choices=METHODS, help=argparse.SUPPRESS) + parser.add_argument("--token-count", type=int, help=argparse.SUPPRESS) + args = parser.parse_args() + # 显式限制到一张卡，防止实验误用用户禁止的 GPU 3 或其他运行中的服务。 + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not visible or "," in visible or visible == "3": + parser.error("请通过 CUDA_VISIBLE_DEVICES 显式选择一张空闲 GPU，不能使用 3 号卡") + if args.repeats < 1 or min(args.sizes) < 1 or args.delay_cycles < 1 or args.attempts < 1: + parser.error("重复数、规模、延时和分配尝试次数必须大于 0") + if args.worker: + run_worker(args) + return + + root = args.output_dir or Path("/tmp") / ( + "pinned_empty_h2d_" + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + ) + root.mkdir(parents=True, exist_ok=False) + all_trials, summaries = [], [] + version = None + for token_count in args.sizes: + for method in METHODS: + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--worker", + "--method", + method, + "--token-count", + str(token_count), + "--repeats", + str(args.repeats), + "--delay-cycles", + str(args.delay_cycles), + "--attempts", + str(args.attempts), + ] + child = subprocess.run(command, capture_output=True, text=True, check=True, timeout=180) + (root / f"{token_count}_{method}.json").write_text(child.stdout) + if child.stderr: + (root / f"{token_count}_{method}.stderr").write_text(child.stderr) + data = json.loads(child.stdout) + version = {key: data[key] for key in ("torch_version", "cuda_version")} + rows = data["trials"] + all_trials.extend(rows) + summary = { + "token_count": token_count, + "method": method, + "repeats": len(rows), + "pending_window_trials": sum(row["pending_allocation_checks"] > 0 for row in rows), + "reused_while_pending": sum(row["source_reused_while_pending"] for row in rows), + "overwrite_finished_while_pending": sum(row["overwrite_finished_while_pending"] for row in rows), + "correct_outputs": sum(row["output_correct"] for row in rows), + "fully_overwritten_outputs": sum(row["overwritten_tokens"] == token_count for row in rows), + "reusable_after_completion": sum(row["source_reused_after_completion"] is True for row in rows), + } + summaries.append(summary) + print(json.dumps(summary, ensure_ascii=False), flush=True) + + evidence = { + **version, + "cuda_visible_devices": visible, + "delay_cycles": args.delay_cycles, + "allocation_attempts": args.attempts, + "source_slice_offset_tokens": 64, + "raw_copy": "cudaMemcpyBatchAsync with cudaMemcpySrcAccessOrderStream", + "isolated_subprocess_per_method_and_size": True, + "replacement_allocations_retained_until_stream_completion": True, + "summary": summaries, + "trials": all_trials, + } + (root / "result.json").write_text(json.dumps(evidence, ensure_ascii=False, indent=2)) + with (root / "summary.csv").open("w", newline="") as output: + writer = csv.DictWriter(output, fieldnames=list(summaries[0])) + writer.writeheader() + writer.writerows(summaries) + print(f"实验结果：{root}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 7205c34e53..ce19f3bc61 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -28,6 +28,7 @@ sys.path.append(str(REPO_ROOT)) from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput +from lightllm.server.router.model_infer.infer_req_ext1 import ReqTokenIndex, create_req_token_index from lightllm.common.basemodel.triton_kernel.gen_mtp_prefill_params import gen_mtp_new_input_ids from lightllm.models import get_draft_model_class, get_model from lightllm.server.api_cli import make_argument_parser @@ -158,6 +159,7 @@ def __init__( self.draft_models = draft_models self.token_source = token_source self._hold_kv_len_by_req: Dict[int, int] = {} + self._req_token_indexes: Dict[int, ReqTokenIndex] = {} def _case_iters(self, warmup: bool) -> int: return self.args.warmup_iters if warmup else self.args.bench_iters @@ -558,32 +560,43 @@ def _make_decode_input( def _ensure_req_kv_capacity(self, req_idx: torch.Tensor, target_kv_lens: torch.Tensor) -> None: """按请求预分配完整 KV 页，并将物理索引写入请求索引表。""" - req_ids = req_idx.detach().cpu().tolist() + req_indexes = req_idx.detach().cpu().tolist() target_lens = target_kv_lens.detach().cpu().tolist() - assert len(req_ids) == len(target_lens) + assert len(req_indexes) == len(target_lens) - # MTP decode 会为同一个请求展开多行，只需按其中最大的逻辑长度扩容一次。 - max_target_len_by_req = {} - for req_id, target_len in zip(req_ids, target_lens): - 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)) + # 按 ReqManager 槽位合并 MTP 展开的多行，只按最大逻辑 token 长度扩容一次。 + max_target_len_by_req_index = {} + for req_index, target_len in zip(req_indexes, target_lens): + req_index = int(req_index) + max_target_len_by_req_index[req_index] = max(max_target_len_by_req_index.get(req_index, 0), int(target_len)) mem_manager = self.model.req_manager.mem_manager page_size = mem_manager.page_size - cpu_table = self.model.req_manager.req_to_token_indexs_cpu - gpu_table = self.model.req_manager.req_to_token_indexs - for req_id, target_len in max_target_len_by_req.items(): + req_manager = self.model.req_manager + for req_index, target_len in max_target_len_by_req_index.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) + old_hold_kv_len = self._hold_kv_len_by_req.get(req_index, 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) - cpu_indexes = cpu_table[req_id, old_hold_kv_len:new_hold_kv_len] - cpu_indexes.copy_(mem_indexes) - gpu_table[req_id, old_hold_kv_len:new_hold_kv_len].copy_(cpu_indexes, non_blocking=True) - self._hold_kv_len_by_req[req_id] = new_hold_kv_len + if req_index not in self._req_token_indexes: + # 静态 benchmark 不进入真实 InferReq 的共享内存生命周期，轻量对象 + # 提供请求槽位、启动参数和输入长度，复用在线请求的索引工厂。 + benchmark_req = SimpleNamespace( + req_idx=req_index, + req_manager=req_manager, + args=self.args, + shm_req=SimpleNamespace(input_len=target_len), + ) + self._req_token_indexes[req_index] = create_req_token_index( + benchmark_req, index_stream=torch.cuda.current_stream() + ) + token_index = self._req_token_indexes[req_index] + token_index.append_token_indexes(mem_indexes) + token_index.get_gpu_token_indexes(old_hold_kv_len, new_hold_kv_len) + self._hold_kv_len_by_req[req_index] = new_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: @@ -736,6 +749,7 @@ def _reset_model_cache(self): self.model.mem_manager.free_all() self.model.req_manager.free_all() self._hold_kv_len_by_req.clear() + self._req_token_indexes.clear() torch.cuda.synchronize() torch.cuda.empty_cache() diff --git a/unit_tests/common/cpu_index_helpers.py b/unit_tests/common/cpu_index_helpers.py new file mode 100644 index 0000000000..138c0a7e77 --- /dev/null +++ b/unit_tests/common/cpu_index_helpers.py @@ -0,0 +1,70 @@ +"""在轻量生命周期测试中复用生产 ReqTokenIndex；矩阵仅作为测试数据源。""" +from types import SimpleNamespace + +import torch + + +def make_cpu_index_manager(token_table, page_size=1, **attributes): + from lightllm.common.req_manager.base import ReqManager + from lightllm.server.router.model_infer.infer_req_ext1 import PinMemCacheReqTokenIndex + + class CpuOnlyTokenIndex(PinMemCacheReqTokenIndex): + def _reserve(self, end): + # CPU 生命周期测试只模拟缓冲区分配，追加、快照和 reset 仍复用生产实现。 + # 按测试请求表容量申请普通内存，不依赖 CUDA 或 pinned allocator。 + if self._cpu_buffer is not None and self._cpu_buffer.numel() >= end: + return + buffer = torch.empty_like(self.gpu_indexes, device="cpu", requires_grad=False) + if self.token_index_len: + buffer[: self.token_index_len].copy_(self._cpu_buffer[: self.token_index_len]) + self._cpu_buffer = buffer + + class TestReqManager(ReqManager): + def get_req_token_index(self, req_idx): + if req_idx not in self._req_token_indexes: + req = SimpleNamespace( + req_idx=req_idx, + req_manager=self, + args=SimpleNamespace(req_token_index_mode="pin_mem_cache"), + shm_req=SimpleNamespace(input_len=self.token_indexes_cpu.shape[1]), + ) + if self.req_to_token_indexs.is_cuda: + helper = PinMemCacheReqTokenIndex( + req, index_stream=torch.cuda.default_stream(self.req_to_token_indexs.device) + ) + else: + # CPU 测试不调用 GPU 接口，无需构造索引流。 + helper = CpuOnlyTokenIndex(req, index_stream=None) + # 既有底层 CUDA 测试直接布置 source；生产实现按请求独立申请。 + helper._cpu_buffer = self.token_indexes_cpu[req_idx] + self._req_token_indexes[req_idx] = helper + return self._req_token_indexes[req_idx] + + def get_cpu_token_indexes(self, req_idx, start, end): + return self.get_req_token_index(req_idx).get_cpu_token_indexes(start, end) + + def append_token_indexes(self, req_idx, indexes): + self.get_req_token_index(req_idx).append_token_indexes(indexes) + + manager = TestReqManager.__new__(TestReqManager) + manager.page_size = page_size + manager.req_token_index_mode = "pin_mem_cache" + manager.token_indexes_cpu = token_table + manager.req_to_token_indexs = torch.full_like(token_table, -1) + manager._req_token_indexes = {} + for name, value in attributes.items(): + setattr(manager, name, value) + return manager + + +def bind_cpu_indexes(req, manager, gpu_index_submitted_len=0, token_index_len=None): + req.req_manager = manager + if not hasattr(req, "args"): + req.args = SimpleNamespace(req_token_index_mode=manager.req_token_index_mode) + req.token_index = manager.get_req_token_index(req.req_idx) + # 矩阵只提供容量和预置数据；追加位置取请求实际持有的长度，不能使用整行容量。 + if token_index_len is None: + token_index_len = req.hold_kv_len if hasattr(req, "hold_kv_len") else 0 + req.token_index.token_index_len = token_index_len + req.token_index.gpu_index_submitted_len = gpu_index_submitted_len + return req diff --git a/unit_tests/common/test_kv_cache_allocator.py b/unit_tests/common/test_kv_cache_allocator.py new file mode 100644 index 0000000000..989338d437 --- /dev/null +++ b/unit_tests/common/test_kv_cache_allocator.py @@ -0,0 +1,333 @@ +import random +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.kv_cache_mem_manager import allocator as impl +from lightllm.common.kv_cache_mem_manager import mem_manager as mem_impl + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires pinned CPU allocations") + + +class _SharedInt: + def __init__(self, name): + self.name = name + self.value = None + + def set_value(self, value): + self.value = value + + def get_value(self): + return self.value + + +@pytest.fixture +def make_allocator(monkeypatch): + monkeypatch.setattr(impl, "SharedInt", _SharedInt) + monkeypatch.setattr(impl, "get_current_rank_in_node", lambda: 0) + return impl.KvCacheAllocator + + +def assert_capacity(allocator, available_tokens): + assert allocator.can_use_mem_size == available_tokens + assert allocator.shared_can_use_token_num.get_value() == available_tokens + assert (allocator.mark_end - allocator.mark_start) * allocator.page_size == available_tokens + + +@pytest.mark.parametrize("page_size", [1, 3, 4, 64]) +def test_compressed_page_state_returns_complete_pinned_token_indexes(make_allocator, page_size): + allocator = make_allocator(5 * page_size, page_size) + assert allocator.mem_state.tolist() == [i * page_size for i in range(5)] + assert allocator.mem_state.dtype == torch.int32 and allocator.mem_state.is_pinned() + assert allocator.mem_state.numel() == allocator.size // page_size + assert allocator._return_buffer is None + assert_capacity(allocator, 5 * page_size) + + first = allocator.alloc(2 * page_size) + second = allocator.alloc(page_size) + assert first.tolist() == list(range(2 * page_size)) + assert second.tolist() == list(range(2 * page_size, 3 * page_size)) + assert first.dtype == torch.int32 and first.is_pinned() + assert first.untyped_storage().data_ptr() == allocator._return_buffer.untyped_storage().data_ptr() + assert first.untyped_storage().nbytes() == 1024 * 1024 * first.element_size() + assert allocator._return_buffer.numel() == 1024 * 1024 - 3 * page_size + assert second.data_ptr() == first.data_ptr() + first.numel() * first.element_size() + assert first.untyped_storage().data_ptr() != allocator.mem_state.untyped_storage().data_ptr() + assert_capacity(allocator, 2 * page_size) + + +@pytest.mark.parametrize("page_size", [1, 3, 4, 64]) +@pytest.mark.parametrize("index_kind", ["list", "cpu", "cuda"]) +def test_nonadjacent_pages_are_recycled_from_token_indexes(make_allocator, page_size, index_kind): + allocator = make_allocator(5 * page_size, page_size) + indexes = allocator.alloc(5 * page_size).clone() + # Recycle physical pages 3 and 1, in that order, using full token-level indexes. + released = torch.cat([indexes[3 * page_size : 4 * page_size], indexes[page_size : 2 * page_size]]) + expected = released.clone() + if index_kind == "list": + released = released.tolist() + elif index_kind == "cuda": + released = released.cuda() + allocator.free(released) + assert_capacity(allocator, 2 * page_size) + assert allocator.mem_state[allocator.mark_start :].tolist() == [3 * page_size, page_size] + + # free() owns the compressed indexes even if the caller reuses its input buffer. + if isinstance(released, torch.Tensor): + released.fill_(-1) + else: + released[:] = [-1] * len(released) + recycled = allocator.alloc(2 * page_size) + assert torch.equal(recycled, expected) + assert_capacity(allocator, 0) + with pytest.raises(AssertionError, match="error alloc state"): + allocator.alloc(page_size) + assert_capacity(allocator, 0) + + +@pytest.mark.parametrize("page_size", [1, 3, 64]) +def test_return_buffer_replacement_and_free_preserve_retained_results(make_allocator, page_size): + # A full slab for aligned page sizes; page_size=3 leaves one unusable element. + size = 1024 * 1024 // page_size * page_size + allocator = make_allocator(size, page_size) + pointers, retained = [], [] + for iteration in range(4): + indexes = allocator.alloc(size) + pointers.append(indexes.data_ptr()) + for old, snapshot in retained: + assert torch.equal(old, snapshot) + snapshot = indexes.clone() + retained.append((indexes, snapshot)) + allocator.free(indexes.view(-1, page_size).flip(0).reshape(-1)) + assert torch.equal(indexes, snapshot) + assert_capacity(allocator, size) + assert len(set(pointers)) == 4 + assert allocator._return_buffer.numel() == 1024 * 1024 - size + + +@pytest.mark.parametrize("page_size", [1, 3, 64]) +def test_large_requests_allocate_separately_and_preserve_small_buffer(make_allocator, page_size): + large_size = (1024 * 1024 // page_size + 1) * page_size + allocator = make_allocator(large_size + page_size, page_size) + large = allocator.alloc(large_size) + snapshot = large.clone() + assert allocator._return_buffer is None + assert large.is_pinned() and large.dtype == torch.int32 + assert large.untyped_storage().nbytes() == large_size * large.element_size() + assert large.tolist() == list(range(large_size)) + + small = allocator.alloc(page_size) + remaining = allocator._return_buffer + assert remaining.numel() == 1024 * 1024 - page_size + assert small.tolist() == list(range(large_size, large_size + page_size)) + allocator.free(large.view(-1, page_size).flip(0).reshape(-1)) + recycled = allocator.alloc(large_size) + assert recycled.data_ptr() != large.data_ptr() + assert allocator._return_buffer is remaining + assert torch.equal(recycled, snapshot.view(-1, page_size).flip(0).reshape(-1)) + assert torch.equal(large, snapshot) + assert_capacity(allocator, 0) + + +@pytest.mark.parametrize("page_size", [1, 3, 4, 64]) +def test_free_all_restores_page_stride_and_token_capacity(make_allocator, page_size): + allocator = make_allocator(6 * page_size, page_size) + indexes = allocator.alloc(6 * page_size) + allocator.free(indexes.view(-1, page_size).flip(0).reshape(-1)) + allocator.alloc(page_size) + remaining = allocator._return_buffer + allocator.free_all() + assert allocator.mem_state.tolist() == [i * page_size for i in range(6)] + assert allocator._return_buffer is remaining + assert_capacity(allocator, 6 * page_size) + assert allocator.alloc(6 * page_size).tolist() == list(range(6 * page_size)) + + +@pytest.mark.parametrize("page_size", [1, 3, 4, 64]) +def test_resize_and_empty_operations_preserve_page_size(make_allocator, page_size): + allocator = make_allocator(3 * page_size, page_size) + previous = allocator.alloc(page_size) + snapshot = previous.clone() + for page_count in [7, 2, 0, 4]: + allocator.resize(page_count * page_size) + assert allocator.mem_state.tolist() == [i * page_size for i in range(page_count)] + assert allocator._return_buffer is None + assert torch.equal(previous, snapshot) + assert_capacity(allocator, page_count * page_size) + assert allocator.alloc(0).numel() == 0 + assert allocator._return_buffer is None + allocator.free([]) + allocator.free(torch.empty(0, dtype=torch.int32)) + assert_capacity(allocator, page_count * page_size) + assert allocator.alloc(page_count * page_size).tolist() == list(range(page_count * page_size)) + assert_capacity(allocator, 0) + allocator.free_all() + assert_capacity(allocator, page_count * page_size) + + +@pytest.mark.parametrize("size,page_size", [(4, 0), (4, -1), (-4, 4), (5, 4)]) +def test_invalid_constructor_sizes_are_rejected(make_allocator, size, page_size): + with pytest.raises(AssertionError): + make_allocator(size, page_size) + + +def test_invalid_operations_leave_allocation_state_unchanged(make_allocator): + allocator = make_allocator(16, 4) + allocator.alloc(4) + for operation in [ + lambda: allocator.alloc(-4), + lambda: allocator.alloc(1), + lambda: allocator.alloc(16), + lambda: allocator.free([0, 1, 2]), + lambda: allocator.free(list(range(8))), + lambda: allocator.resize(17), + ]: + before = allocator.mem_state.clone() + remaining = allocator._return_buffer + with pytest.raises(AssertionError): + operation() + assert torch.equal(allocator.mem_state, before) + assert allocator._return_buffer is remaining + assert allocator.size == 16 + assert_capacity(allocator, 12) + + +@pytest.mark.parametrize("page_size", [1, 3, 4, 64]) +def test_randomized_alloc_free_never_overlaps_live_pages(make_allocator, page_size): + page_count = 23 + allocator = make_allocator(page_count * page_size, page_size) + rng = random.Random(7341 + page_size) + live_pages = set() + for _ in range(300): + available = page_count - len(live_pages) + if available and (not live_pages or rng.random() < 0.6): + count = rng.randint(1, min(available, 5)) + indexes = allocator.alloc(count * page_size).clone() + blocks = indexes.view(-1, page_size) + bases = blocks[:, 0].tolist() + assert len(set(bases)) == count + assert all(0 <= base < allocator.size and base % page_size == 0 for base in bases) + assert live_pages.isdisjoint(bases) + assert torch.equal(blocks, blocks[:, :1] + torch.arange(page_size, dtype=torch.int32)) + live_pages.update(bases) + else: + bases = rng.sample(sorted(live_pages), rng.randint(1, min(len(live_pages), 5))) + indexes = [base + offset for base in bases for offset in range(page_size)] + allocator.free(indexes if rng.random() < 0.5 else torch.tensor(indexes, dtype=torch.int32)) + live_pages.difference_update(bases) + assert_capacity(allocator, (page_count - len(live_pages)) * page_size) + allocator.free([base + offset for base in sorted(live_pages) for offset in range(page_size)]) + assert_capacity(allocator, allocator.size) + assert sorted(allocator.alloc(allocator.size).tolist()) == list(range(allocator.size)) + + +@pytest.mark.parametrize("page_size", [1, 3, 4, 64]) +def test_memory_manager_passes_page_size_and_resizes_token_capacity(make_allocator, monkeypatch, page_size): + monkeypatch.setattr(mem_impl, "get_env_start_args", lambda: SimpleNamespace(page_size=page_size)) + manager = mem_impl.MemoryManager(5 * page_size + page_size - 1, torch.float16, 1, 8, 1) + assert manager.size == 5 * page_size + assert manager.allocator.page_size == page_size + assert manager.allocator.mem_state.tolist() == [i * page_size for i in range(5)] + assert manager.HOLD_TOKEN_MEMINDEXES == tuple(range(5 * page_size, 6 * page_size)) + indexes = manager.alloc(2 * page_size) + manager.free(indexes) + assert_capacity(manager.allocator, manager.size) + + manager.resize_mem(3 * page_size + page_size - 1) + assert manager.size == 3 * page_size + assert manager.allocator.mem_state.tolist() == [i * page_size for i in range(3)] + assert manager.HOLD_TOKEN_MEMINDEXES == tuple(range(3 * page_size, 4 * page_size)) + indexes = manager.alloc(manager.size) + assert indexes.tolist() == list(range(manager.size)) + manager.free(indexes.cuda()) + assert_capacity(manager.allocator, manager.size) + + +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.parametrize("page_size", [1, 64]) +def test_request_free_indexes_recycle_pages_after_reset_batch_and_slot_reuse( + make_allocator, monkeypatch, mode, page_size +): + # 延迟导入，沿用在线推理初始化顺序，不增加另一套请求/allocator 夹具。 + from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 + from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext + from lightllm.server.router.model_infer.infer_req_ext1 import create_req_token_index + + allocator = make_allocator(6 * page_size, page_size) + allocated = allocator.alloc(allocator.size) + first = torch.cat([allocated[4 * page_size : 5 * page_size], allocated[page_size : 2 * page_size]]).pin_memory() + second = allocated[3 * page_size : 4 * page_size] + expected = torch.cat([first, second]) + table = torch.full((2, 2 * page_size), -1, dtype=torch.int32, device="cuda") + manager = SimpleNamespace(req_to_token_indexs=table) + context = InferenceContext.__new__(InferenceContext) + context.req_manager = manager + context.radix_cache = None + reqs = [] + + def unexpected_snapshot(*args): + pytest.fail("pure_gpu 的实际释放路径不能先拷回完整 CPU token 快照") + + with torch.cuda.stream(torch.cuda.default_stream()): + for row, indexes in enumerate((first, second)): + req = InferReq.__new__(InferReq) + req.req_idx = row + req.req_manager = manager + req.args = SimpleNamespace(req_token_index_mode=mode) + req.cur_kv_len = req.hold_kv_len = indexes.numel() + req.shm_req = SimpleNamespace(input_len=req.cur_kv_len, shm_cur_kv_len=req.cur_kv_len) + req.token_index = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + req.token_index.append_token_indexes(indexes) + if mode == "pure_gpu": + monkeypatch.setattr(req.token_index, "get_cpu_token_indexes", unexpected_snapshot) + reqs.append(req) + + # 混合空区间和非空区间；GPU 模式的空区间也必须留在 GPU，才能直接 cat。 + torch.cuda._sleep(10_000_000) + released = [reqs[0].token_index.get_free_token_indexes(0, 0)] + for req in reqs: + context.free_a_req_mem(released, req) + assert req.token_index.token_index_len == req.token_index.gpu_index_submitted_len == 0 + assert released[0].numel() == 0 + assert all(indexes.device.type == ("cuda" if mode == "pure_gpu" else "cpu") for indexes in released) + merged = torch.cat(released) + + # 默认流先合并旧视图，再重写请求槽位；allocator 只能回收已合并的旧索引。 + # 不在 free 前同步或 D2H，真实 free 必须等待默认流并按 page_size 抽取页首。 + table.fill_(-99) + allocator.free(merged) + assert_capacity(allocator, 3 * page_size) + assert allocator.mem_state[allocator.mark_start :].tolist() == [4 * page_size, page_size, 3 * page_size] + assert allocator.alloc(3 * page_size).tolist() == expected.tolist() + assert_capacity(allocator, 0) + + +@pytest.mark.parametrize("page_size", [1, 64]) +def test_pure_gpu_free_waits_for_index_stream_before_allocator_d2h(make_allocator, page_size): + from lightllm.server.router.model_infer.infer_req_ext1 import PureGpuReqTokenIndex + + allocator = make_allocator(4 * page_size, page_size) + indexes = allocator.alloc(2 * page_size) + table = torch.full((1, indexes.numel()), -1, dtype=torch.int32, device="cuda") + index_stream = torch.cuda.Stream() + index_stream.wait_stream(torch.cuda.current_stream()) + req = SimpleNamespace(req_idx=0, req_manager=SimpleNamespace(req_to_token_indexs=table)) + helper = PureGpuReqTokenIndex(req, index_stream=index_stream) + + with torch.cuda.stream(index_stream): + # 延迟索引写入；主线程不等待，默认流必须通过释放 getter 承接写入依赖。 + torch.cuda._sleep(100_000_000) + helper.append_token_indexes(indexes) + + free_indexes = helper.get_free_token_indexes(0, indexes.numel()) + assert free_indexes.data_ptr() == table[0].data_ptr() + # 不先同步索引流：真实 allocator.free 在当前流按页取首索引并阻塞 D2H。 + allocator.free(free_indexes) + helper.reset() + with torch.cuda.stream(index_stream): + table.fill_(-99) + assert_capacity(allocator, allocator.size) + assert allocator.alloc(indexes.numel()).tolist() == indexes.tolist() + index_stream.synchronize() diff --git a/unit_tests/common/test_pinned_memory_stream.py b/unit_tests/common/test_pinned_memory_stream.py new file mode 100644 index 0000000000..aaff9db033 --- /dev/null +++ b/unit_tests/common/test_pinned_memory_stream.py @@ -0,0 +1,353 @@ +"""验证 raw H2D 的源引用由完成事件和后台线程保活，不依赖服务启动。""" + +import gc +from concurrent.futures import ThreadPoolExecutor +import threading +from types import SimpleNamespace +import weakref + +import pytest +import torch + +from lightllm.common import cuda_wrapper + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and pinned CPU memory") + + +@pytest.fixture(autouse=True) +def release_futures(monkeypatch): + monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", True) + pool_factory = cuda_wrapper.ThreadPoolExecutor + pools = {} + futures = [] + + def track_pool(pool): + original_submit = pool.submit + pools[pool] = original_submit + + def record_submit(*args, **kwargs): + future = original_submit(*args, **kwargs) + futures.append(future) + return future + + monkeypatch.setattr(pool, "submit", record_submit) + + def create_pool(*args, **kwargs): + pool = pool_factory(*args, **kwargs) + track_pool(pool) + return pool + + # 观察已有或后续实际创建的线程池，夹具自身不触发延迟初始化。 + pool = cuda_wrapper.PinnedH2DBatchCopier._release_pool + if pool is not None: + track_pool(pool) + monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", create_pool) + yield futures + # 先收尾所有真实等待任务，再撤销本例 monkeypatch；生产池可被后续测试继续复用。 + for future in futures: + future.result(timeout=10) + for submit in pools.values(): + submit(lambda: None).result(timeout=10) + + +def test_empty_batch_and_fallback_do_not_create_release_pool(monkeypatch): + copier = cuda_wrapper.PinnedH2DBatchCopier + monkeypatch.setattr(copier, "_release_pool", None) + monkeypatch.setattr(copier, "_native_batch_available", False) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(cuda_stream=0)) + + def unexpected_pool(*args, **kwargs): + pytest.fail("空批次和 copy_ 回退不能创建保活线程池") + + monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", unexpected_pool) + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([], []) + source = torch.arange(8, dtype=torch.int32) + target = torch.empty_like(source) + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) + assert copier._release_pool is None + assert torch.equal(target, source) + + +def test_release_pool_is_initialized_once_for_concurrent_callers(monkeypatch): + copier = cuda_wrapper.PinnedH2DBatchCopier + monkeypatch.setattr(copier, "_release_pool", None) + entered = threading.Event() + allow_creation = threading.Event() + start = threading.Barrier(8) + shared_pool = object() + creations = [] + + def create_pool(**kwargs): + creations.append(kwargs) + entered.set() + assert allow_creation.wait(timeout=5) + return shared_pool + + def get_pool(): + start.wait(timeout=5) + return copier._get_release_pool() + + monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", create_pool) + with ThreadPoolExecutor(max_workers=8) as callers: + futures = [callers.submit(get_pool) for _ in range(8)] + try: + assert entered.wait(timeout=5) + finally: + allow_creation.set() + assert all(future.result(timeout=5) is shared_pool for future in futures) + assert creations == [{"max_workers": 1, "thread_name_prefix": "lightllm-pinned-release"}] + + +@requires_cuda +@pytest.mark.parametrize("native_batch", [False, True]) +def test_default_stream_copy_orders_reads_and_reuses_upload_stream(monkeypatch, native_batch): + runtime = pytest.importorskip("cuda.bindings.runtime") + copier = cuda_wrapper.PinnedH2DBatchCopier + monkeypatch.setattr(copier, "_native_batch_available", native_batch) + stream = torch.cuda.default_stream() + source = torch.arange(8192, dtype=torch.int32, pin_memory=True) + target = torch.empty(8192, dtype=torch.int32, device="cuda") + native_copy = runtime.cudaMemcpyBatchAsync + submitted_streams, observed = [], [] + + def record_batch(*args): + submitted_streams.append(args[-1]) + return native_copy(*args) + + monkeypatch.setattr(runtime, "cudaMemcpyBatchAsync", record_batch) + with torch.cuda.stream(stream): + for _ in range(2): + torch.cuda._sleep(20_000_000) + target.fill_(-1) + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) + assert torch.cuda.current_stream() == stream + observed.append(target.clone()) + for result in observed: + assert torch.equal(result.cpu(), source) + if native_batch: + assert len(submitted_streams) == 2 + assert submitted_streams[0] == submitted_streams[1] != 0 + else: + assert submitted_streams == [] + + +@requires_cuda +def test_raw_batch_records_one_completion_event_without_small_copy(monkeypatch, release_futures): + runtime = pytest.importorskip("cuda.bindings.runtime") + source = torch.arange(2048, dtype=torch.int32, pin_memory=True) + another_source = torch.arange(2048, dtype=torch.int32, pin_memory=True).add_(10000) + sources = [source[128:384], source[1000:1256], another_source[17:273]] + target = torch.full((1024,), -1, dtype=torch.int32, device="cuda") + targets = [target[2:258], target[300:556], target[600:856]] + expected = torch.full((1024,), -1, dtype=torch.int32) + expected[2:258], expected[300:556], expected[600:856] = sources + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.default_stream()) + calls = [] + native_copy, event_record = runtime.cudaMemcpyBatchAsync, torch.cuda.Event.record + + def record_native(*args): + calls.append(("batch", args[3], args[-1])) + return native_copy(*args) + + def record_event(event, on_stream=None): + calls.append(("event", 1, on_stream.cuda_stream)) + return event_record(event, on_stream) + + def unexpected_copy(*args, **kwargs): + pytest.fail("raw batch 保活不应额外提交小片段 copy_ 或进入逐项回退") + + def unexpected_sync(*args, **kwargs): + pytest.fail("正常 raw batch 不应在提交线程同步 CUDA 流") + + monkeypatch.setattr(runtime, "cudaMemcpyBatchAsync", record_native) + monkeypatch.setattr(torch.cuda.Event, "record", record_event) + monkeypatch.setattr(torch.Tensor, "copy_", unexpected_copy) + with torch.cuda.stream(stream), monkeypatch.context() as no_sync: + no_sync.setattr(torch.cuda.Stream, "synchronize", unexpected_sync) + no_sync.setattr(torch.cuda, "synchronize", unexpected_sync) + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async(targets, sources) + assert len(release_futures) == 1 + release_futures[0].result(timeout=10) + assert cuda_wrapper.PinnedH2DBatchCopier._native_batch_available + assert calls == [("batch", 3, stream.cuda_stream), ("event", 1, stream.cuda_stream)] + assert torch.equal(target.cpu(), expected) + + +@requires_cuda +def test_background_wait_holds_tensor_references_without_blocking_submitter(monkeypatch, release_futures): + pytest.importorskip("cuda.bindings.runtime") + source = torch.arange(8192, dtype=torch.int32, pin_memory=True) + source_views = [source[64:1088], source[1088:4160]] + expected = torch.cat(source_views) + source_ref = weakref.ref(source) + view_refs = [weakref.ref(view) for view in source_views] + target = torch.empty(4096, dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.default_stream()) + worker_wait_started = threading.Event() + allow_worker_finish = threading.Event() + main_thread = threading.get_ident() + event_synchronize = torch.cuda.Event.synchronize + + def controlled_wait(event): + assert threading.get_ident() != main_thread, "CUDA 完成事件只能由后台等待线程同步" + worker_wait_started.set() + assert allow_worker_finish.wait(timeout=10), "测试未允许后台等待结束" + event_synchronize(event) + + monkeypatch.setattr(torch.cuda.Event, "synchronize", controlled_wait) + try: + with torch.cuda.stream(stream): + torch.cuda._sleep(100_000_000) + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target[:1024], target[1024:]], source_views) + # gate 仍关闭时提交调用已返回，且后台任务正在等待，而非阻塞主线程。 + assert worker_wait_started.wait(timeout=5) + assert len(release_futures) == 1 and not release_futures[0].done() + source_views.clear() + del source_views, source + gc.collect() + assert source_ref() is not None and all(reference() is not None for reference in view_refs) + finally: + allow_worker_finish.set() + for future in release_futures: + assert future.result(timeout=10) is None + stream.synchronize() + gc.collect() + assert source_ref() is None and all(reference() is None for reference in view_refs) + assert torch.equal(target.cpu(), expected) + + +@requires_cuda +@pytest.mark.parametrize("native_batch", [True, False]) +@pytest.mark.parametrize("shared_storage", [True, False]) +def test_async_copy_keeps_freed_source_storage_out_of_reuse(monkeypatch, native_batch, shared_storage): + if native_batch: + pytest.importorskip("cuda.bindings.runtime") + monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", native_batch) + source = torch.arange(8192, dtype=torch.int32, pin_memory=True) + another_source = source if shared_storage else torch.arange(8192, dtype=torch.int32, pin_memory=True).add_(10000) + source_ptrs = {source.data_ptr(), another_source.data_ptr()} + source_views = [source[64:1088], another_source[1088:4160]] + expected = torch.cat(source_views) + target = torch.empty(4096, dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.default_stream()) + + with torch.cuda.stream(stream): + torch.cuda._sleep(200_000_000) + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target[:1024], target[1024:]], source_views) + assert cuda_wrapper.PinnedH2DBatchCopier._native_batch_available is native_batch + del source_views, source, another_source + + pending_checks = 0 + for _ in range(64): + still_pending = not stream.query() + replacement = torch.empty(8192, dtype=torch.int32, pin_memory=True) + if still_pending: + assert replacement.data_ptr() not in source_ptrs + pending_checks += 1 + replacement.fill_(-1) + del replacement + stream.synchronize() + assert pending_checks > 0, "延时传输应覆盖源释放和 pinned 内存复用检查" + assert torch.equal(target.cpu(), expected) + + +@requires_cuda +@pytest.mark.parametrize("failure", ["record", "pool", "submit"]) +def test_event_setup_failure_propagates_without_sync_or_fallback(monkeypatch, failure): + pytest.importorskip("cuda.bindings.runtime") + source = torch.arange(1024, dtype=torch.int32, pin_memory=True) + target = torch.empty_like(source, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.default_stream()) + synchronize = torch.cuda.Stream.synchronize + + def fail_setup(*args, **kwargs): + raise RuntimeError(f"test {failure} failure") + + def unexpected_sync(*args, **kwargs): + pytest.fail("保活接管异常应直接传播，不应同步收尾") + + def unexpected_copy(*args, **kwargs): + pytest.fail("raw batch 已提交后的保护失败不能被 copy_ 回退掩盖") + + if failure == "record": + monkeypatch.setattr(torch.cuda.Event, "record", fail_setup) + elif failure == "pool": + monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_release_pool", None) + monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", fail_setup) + else: + monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier._get_release_pool(), "submit", fail_setup) + monkeypatch.setattr(torch.cuda.Stream, "synchronize", unexpected_sync) + monkeypatch.setattr(torch.Tensor, "copy_", unexpected_copy) + try: + with torch.cuda.stream(stream): + torch.cuda._sleep(20_000_000) + with pytest.raises(RuntimeError, match=f"test {failure} failure"): + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) + finally: + # 测试捕获注入异常后仍要继续运行，自行持有源并等待在途传输结束。 + synchronize(stream) + assert cuda_wrapper.PinnedH2DBatchCopier._native_batch_available + assert torch.equal(target.cpu(), source) + + +@requires_cuda +def test_pytorch_fallback_needs_no_completion_event_or_pool_task(monkeypatch, release_futures): + monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", False) + source = torch.arange(1024, dtype=torch.int32, pin_memory=True) + target = torch.empty_like(source, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.default_stream()) + copy_sizes = [] + copy = torch.Tensor.copy_ + + def unexpected_event(*args, **kwargs): + pytest.fail("普通 PyTorch copy_ 已保护源 storage，不应建立额外完成事件") + + def record_copy(dst, src, *args, **kwargs): + copy_sizes.append(src.numel() * src.element_size()) + return copy(dst, src, *args, **kwargs) + + monkeypatch.setattr(torch.cuda, "Event", unexpected_event) + monkeypatch.setattr(torch.Tensor, "copy_", record_copy) + with torch.cuda.stream(stream): + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) + stream.synchronize() + assert not cuda_wrapper.PinnedH2DBatchCopier._native_batch_available + assert release_futures == [] + assert copy_sizes == [source.numel() * source.element_size()] + assert torch.equal(target.cpu(), source) + + +@requires_cuda +def test_source_release_pool_is_reused_across_batches(release_futures): + pytest.importorskip("cuda.bindings.runtime") + pool = cuda_wrapper.PinnedH2DBatchCopier._get_release_pool() + source = torch.arange(1024, dtype=torch.int32, pin_memory=True) + target = torch.empty((2, 1024), dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.default_stream()) + with torch.cuda.stream(stream): + for row in range(2): + cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target[row]], [source]) + assert cuda_wrapper.PinnedH2DBatchCopier._get_release_pool() is pool + assert len(release_futures) == 2 + for future in release_futures: + assert future.result(timeout=10) is None + assert torch.equal(target.cpu(), source.expand(2, -1)) + + +def test_background_wait_failure_propagates(): + source = torch.arange(8, dtype=torch.int32) + sources = [source] + + def failed_wait(): + raise RuntimeError("test background wait failure") + + with pytest.raises(RuntimeError, match="test background wait failure"): + cuda_wrapper.PinnedH2DBatchCopier._wait_and_release_pinned_sources( + SimpleNamespace(synchronize=failed_wait), sources + ) diff --git a/unit_tests/common/test_req_manager_cpu.py b/unit_tests/common/test_req_manager_cpu.py index 473624f95d..24a482f9af 100644 --- a/unit_tests/common/test_req_manager_cpu.py +++ b/unit_tests/common/test_req_manager_cpu.py @@ -11,25 +11,28 @@ from lightllm.common.req_manager.base import ReqManager, _ReqLinkedList from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext +from lightllm.server.router.model_infer.infer_req_ext1 import create_req_token_index from lightllm.server.router.model_infer import infer_batch from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend +from unit_tests.common.cpu_index_helpers import bind_cpu_indexes, make_cpu_index_manager @pytest.fixture(autouse=True) def reset_batch_copy_availability(monkeypatch): from lightllm.common import cuda_wrapper - monkeypatch.setattr(cuda_wrapper, "_cuda_memcpy_batch_available", True) + monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", True) + monkeypatch.setattr(infer_batch.g_infer_context, "req_manager", None) def _make_manager(device="cpu"): - manager = ReqManager.__new__(ReqManager) + manager = make_cpu_index_manager( + torch.full((3, 16), -1, dtype=torch.int32, pin_memory=torch.device(device).type == "cuda"), + req_to_token_indexs=torch.full((3, 16), -1, dtype=torch.int32, device=device), + ) + infer_batch.g_infer_context.req_manager = manager manager.max_request_num = 2 manager.req_list = _ReqLinkedList(2) - manager.req_to_token_indexs_cpu = torch.full( - (3, 16), -1, dtype=torch.int32, pin_memory=torch.device(device).type == "cuda" - ) - manager.req_to_token_indexs = torch.full((3, 16), -1, dtype=torch.int32, device=device) return manager @@ -66,31 +69,33 @@ def test_req_index_free_rejects_unallocated_and_duplicate_indexes(): def test_cpu_updates_own_allocator_values_and_defer_gpu_publication(): manager = _make_manager() indexes = torch.tensor([20, 21, 22, 23], dtype=torch.int32) - manager.req_to_token_indexs_cpu[0, 4:8].copy_(indexes) + manager.token_indexes_cpu[0, 4:8].copy_(indexes) indexes.fill_(99) - manager.req_to_token_indexs_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) + manager.token_indexes_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) - assert manager.req_to_token_indexs_cpu[0, :8].tolist() == [0, 1, 2, 3, 20, 21, 22, 23] + assert manager.token_indexes_cpu[0, :8].tolist() == [0, 1, 2, 3, 20, 21, 22, 23] assert manager.req_to_token_indexs.eq(-1).all() def test_releasing_request_memory_resets_uploaded_length(): manager = _make_manager() - manager.req_to_token_indexs_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) + manager.token_indexes_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) context = InferenceContext() context.req_manager = manager req = SimpleNamespace( req_idx=0, cur_kv_len=6, hold_kv_len=8, - move_to_gpu_index_len=8, shm_req=SimpleNamespace(shm_cur_kv_len=6), ) + bind_cpu_indexes(req, manager, gpu_index_submitted_len=8) released = [] context.free_a_req_mem(released, req) assert released[0].tolist() == list(range(8)) - assert req.cur_kv_len == req.hold_kv_len == req.move_to_gpu_index_len == 0 + assert req.cur_kv_len == req.hold_kv_len == req.token_index.gpu_index_submitted_len == 0 + assert req.token_index.token_index_len == 0 + assert req.token_index._cpu_buffer is None @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -99,13 +104,13 @@ def test_shared_export_keeps_cpu_mirror_local_and_pinned(monkeypatch): req_manager = _make_manager("cuda") mem_manager = impl.MemoryManager.__new__(impl.MemoryManager) - source = req_manager.req_to_token_indexs_cpu + source = req_manager.token_indexes_cpu source_ptr = source.data_ptr() exports = [] def serialize(manager): assert manager.req_to_token_indexs is req_manager.req_to_token_indexs - assert not hasattr(manager, "req_to_token_indexs_cpu") + assert not hasattr(manager, "token_indexes_cpu") exports.append(manager) return memoryview(b"test") @@ -132,40 +137,79 @@ def forbid_cpu_export(*args, **kwargs): mem_manager.write_to_shm(req_manager) assert len(exports) == 4 - assert req_manager.req_to_token_indexs_cpu is source + assert req_manager.token_indexes_cpu is source assert source.data_ptr() == source_ptr assert source.is_pinned() and not source.is_shared() -def test_prefix_match_invalidates_uploaded_length(monkeypatch): - manager = _make_manager() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.parametrize("hybrid", [False, True]) +@pytest.mark.parametrize("resuming", [False, True]) +def test_prefix_match_appends_to_initialized_or_released_index_state(monkeypatch, mode, hybrid, resuming): + manager = _make_manager("cuda") req = InferReq.__new__(InferReq) req.req_idx = 0 - req.args = SimpleNamespace(page_size=4) - req.cur_kv_len = 0 - req.hold_kv_len = req.move_to_gpu_index_len = 4 + req.req_manager = manager + req.args = SimpleNamespace( + page_size=4, + req_token_index_mode=mode, + linear_att_hash_page_size=2, + linear_att_page_block_num=2, + max_req_total_len=16, + ) + req.cur_kv_len = req.hold_kv_len = 0 + req.shared_kv_node = req.tail_small_page_buffer_id = None req.sampling_param = SimpleNamespace(disable_prompt_cache=False) - req.shm_req = SimpleNamespace(shm_prompt_ids=SimpleNamespace(arr=list(range(8)))) + req.shm_req = SimpleNamespace( + input_len=8, + shm_prompt_ids=SimpleNamespace(arr=list(range(8))), + hybrid_token_hash_list=SimpleNamespace(size=3, get_all=lambda: [1, 2, 3]), + ) + req.token_index = create_req_token_index(req, index_stream=torch.cuda.current_stream()) req.get_cur_total_len = lambda: 8 - node = SimpleNamespace(node_prefix_total_len=4) + if resuming: + req.token_index.append_token_indexes(torch.arange(8, dtype=torch.int32)) + req.token_index.get_gpu_token_indexes(0, 8) + req.cur_kv_len, req.hold_kv_len = 6, 8 + release_context = InferenceContext() + released = [] + release_context.free_a_req_mem(released, req) + assert released[0].cpu().tolist() == list(range(8)) + assert req.cur_kv_len == req.hold_kv_len == 0 + assert req.token_index.token_index_len == req.token_index.gpu_index_submitted_len == 0 + + node = SimpleNamespace(node_prefix_total_len=4, big_page_buffer_idx=0, is_big_page_node=lambda: True) indexes = torch.arange(32, 36, dtype=torch.int32) + manager.restore_big_page_state = lambda **kwargs: None context = SimpleNamespace( - is_hybrid_att_model=False, + is_hybrid_att_model=hybrid, req_manager=manager, - radix_cache=SimpleNamespace(match_prefix=lambda key, update_refs: (node, 4, indexes)), + radix_cache=SimpleNamespace(match_prefix=lambda key, **kwargs: (node, 4, indexes)), ) monkeypatch.setattr(infer_batch, "g_infer_context", context) - req._match_radix_cache() + def unexpected_reset(): + pytest.fail("前缀匹配入口的索引状态已为空，不应再次 reset") + + monkeypatch.setattr(req.token_index, "reset", unexpected_reset) + if hybrid: + req._hybrid_match_radix_cache() + else: + req._match_radix_cache() - assert req.move_to_gpu_index_len == 0 + assert req.token_index.token_index_len == 4 + assert req.token_index.gpu_index_submitted_len == (0 if mode == "pin_mem_cache" else 4) assert req.cur_kv_len == req.hold_kv_len == 4 - assert manager.req_to_token_indexs_cpu[0, :4].tolist() == [32, 33, 34, 35] + assert req.shared_kv_node is node + assert req.shm_req.prompt_cache_len == req.shm_req.shm_cur_kv_len == 4 + assert req.token_index.get_cpu_token_indexes(0, 4).tolist() == [32, 33, 34, 35] + assert req.token_index.get_gpu_token_indexes(0, 4).cpu().tolist() == [32, 33, 34, 35] def test_radix_insertion_owns_cpu_indexes_before_request_row_reuse(): manager = _make_manager() - manager.req_to_token_indexs_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) + manager.token_indexes_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) cached = [] def insert(key, indexes): @@ -183,10 +227,11 @@ def insert(key, indexes): shared_kv_node=None, get_input_token_ids=lambda: list(range(6)), ) + bind_cpu_indexes(req, manager) released = [] context._full_att_free_req(released, req) assert released[0].tolist() == [4, 5, 6, 7] - manager.req_to_token_indexs_cpu[0, :8].fill_(99) + manager.token_indexes_cpu[0, :8].fill_(99) assert cached[0].tolist() == [0, 1, 2, 3] assert manager.req_to_token_indexs.eq(-1).all() @@ -200,10 +245,10 @@ def test_async_batch_reads_source_after_prior_stream_writes_and_preserves_untouc stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): torch.cuda._sleep(2_000_000) - manager.req_to_token_indexs_cpu[0, 4:8].copy_(incoming_indexes, non_blocking=True) + manager.token_indexes_cpu[0, 4:8].copy_(incoming_indexes, non_blocking=True) cuda_memcpy_pinned_h2d_batch_async( [manager.req_to_token_indexs[0, 4:8]], - [manager.req_to_token_indexs_cpu[0, 4:8]], + [manager.token_indexes_cpu[0, 4:8]], ) stream.synchronize() @@ -211,11 +256,11 @@ def test_async_batch_reads_source_after_prior_stream_writes_and_preserves_untouc assert manager.req_to_token_indexs[0, :4].eq(-1).all() assert manager.req_to_token_indexs[0, 8:].eq(-1).all() assert manager.req_to_token_indexs[1:].eq(-1).all() - manager.req_to_token_indexs_cpu[0, 4:8].copy_(torch.tensor([80, 81, 82, 83], dtype=torch.int32)) + manager.token_indexes_cpu[0, 4:8].copy_(torch.tensor([80, 81, 82, 83], dtype=torch.int32)) with torch.cuda.stream(stream): cuda_memcpy_pinned_h2d_batch_async( [manager.req_to_token_indexs[0, 4:8]], - [manager.req_to_token_indexs_cpu[0, 4:8]], + [manager.token_indexes_cpu[0, 4:8]], ) stream.synchronize() assert manager.req_to_token_indexs[0, 4:8].cpu().tolist() == [80, 81, 82, 83] @@ -227,11 +272,11 @@ def test_native_batch_uses_stream_order_and_handles_disjoint_ranges(monkeypatch, from cuda.bindings import runtime as cuda_runtime manager = _make_manager("cuda") - manager.req_to_token_indexs_cpu.copy_(torch.arange(48, dtype=torch.int32).reshape(3, 16)) + manager.token_indexes_cpu.copy_(torch.arange(48, dtype=torch.int32).reshape(3, 16)) ranges = [(0, 4, 8), (1, 1, 5), (0, 12, 16)] - expected = torch.full_like(manager.req_to_token_indexs_cpu, -1) + expected = torch.full_like(manager.token_indexes_cpu, -1) for row, start, end in ranges: - expected[row, start:end].copy_(manager.req_to_token_indexs_cpu[row, start:end]) + expected[row, start:end].copy_(manager.token_indexes_cpu[row, start:end]) calls = [] native_copy = cuda_runtime.cudaMemcpyBatchAsync @@ -243,22 +288,12 @@ def record_native(*args): stream = torch.cuda.default_stream() if default_stream else torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - if default_stream: - from lightllm.common import cuda_wrapper as base - - with pytest.raises(AssertionError, match="非默认 CUDA 流"): - cuda_memcpy_pinned_h2d_batch_async( - [manager.req_to_token_indexs[row, start:end] for row, start, end in ranges], - [manager.req_to_token_indexs_cpu[row, start:end] for row, start, end in ranges], - ) - assert calls == [] - assert base._cuda_memcpy_batch_available - return torch.cuda._sleep(2_000_000) cuda_memcpy_pinned_h2d_batch_async( [manager.req_to_token_indexs[row, start:end] for row, start, end in ranges], - [manager.req_to_token_indexs_cpu[row, start:end] for row, start, end in ranges], + [manager.token_indexes_cpu[row, start:end] for row, start, end in ranges], ) + assert torch.cuda.current_stream() == stream result = manager.req_to_token_indexs.clone() stream.synchronize() @@ -285,11 +320,11 @@ def test_batch_fallback_uses_current_stream_without_synchronizing(monkeypatch, f from lightllm.common import cuda_wrapper as base manager = _make_manager("cuda") - manager.req_to_token_indexs_cpu.copy_(torch.arange(48, dtype=torch.int32).reshape(3, 16)) + manager.token_indexes_cpu.copy_(torch.arange(48, dtype=torch.int32).reshape(3, 16)) ranges = [(0, 4, 8), (1, 1, 5), (0, 12, 16)] - expected = torch.full_like(manager.req_to_token_indexs_cpu, -1) + expected = torch.full_like(manager.token_indexes_cpu, -1) for row, start, end in ranges: - expected[row, start:end].copy_(manager.req_to_token_indexs_cpu[row, start:end]) + expected[row, start:end].copy_(manager.token_indexes_cpu[row, start:end]) batch_attempts = [] if failure == "missing_binding": @@ -327,7 +362,7 @@ def record_copy(dst, src, *args, **kwargs): stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) dsts = [manager.req_to_token_indexs[row, start:end] for row, start, end in ranges] - srcs = [manager.req_to_token_indexs_cpu[row, start:end] for row, start, end in ranges] + srcs = [manager.token_indexes_cpu[row, start:end] for row, start, end in ranges] def unexpected_sync(*args, **kwargs): pytest.fail("copy submission must not synchronize") @@ -345,9 +380,9 @@ def unexpected_sync(*args, **kwargs): assert torch.equal(result.cpu(), expected) assert [size for size, _ in loop_calls] == [16] * 3 assert all(copy_stream == stream.cuda_stream for _, copy_stream in loop_calls) - assert not base._cuda_memcpy_batch_available + assert not base.PinnedH2DBatchCopier._native_batch_available attempts_after_failure = len(batch_attempts) - manager.req_to_token_indexs_cpu.fill_(99) + manager.token_indexes_cpu.fill_(99) with torch.cuda.stream(stream): cuda_memcpy_pinned_h2d_batch_async(dsts, srcs) stream.synchronize() @@ -396,7 +431,9 @@ def _make_scheduler(monkeypatch, manager, reqs, budget=16, mixed=False): pause_reqs=lambda reqs, is_master_in_dp: None, ) monkeypatch.setattr(base_backend, "g_infer_context", context) - manager.mem_manager = SimpleNamespace(alloc=lambda size: torch.arange(100, 100 + size, dtype=torch.int32)) + manager.mem_manager = SimpleNamespace( + alloc=lambda size: torch.arange(100, 100 + size, dtype=torch.int32, pin_memory=True) + ) backend = ModeBackend.__new__(ModeBackend) backend.args = args backend.support_overlap = False @@ -414,9 +451,9 @@ def _make_scheduler(monkeypatch, manager, reqs, budget=16, mixed=False): def record_copy(run_reqs): ranges = [ - (req.req_idx, req.move_to_gpu_index_len, req.hold_kv_len) + (req.req_idx, req.token_index.gpu_index_submitted_len, req.hold_kv_len) for req in run_reqs - if req.move_to_gpu_index_len < req.hold_kv_len + if req.token_index.gpu_index_submitted_len < req.hold_kv_len ] if ranges: copies.append(ranges) @@ -426,12 +463,12 @@ def record_copy(run_reqs): return backend, stream, copies -def _make_scheduled_req(row, total_len=12, moved=4): +def _make_scheduled_req(row, total_len=12, gpu_index_submitted_len=4, hold_kv_len=4): req = InferReq.__new__(InferReq) req.req_idx = req.req_id = row req.cur_kv_len = 3 - req.hold_kv_len = 4 - req.move_to_gpu_index_len = moved + req.hold_kv_len = hold_kv_len + bind_cpu_indexes(req, infer_batch.g_infer_context.req_manager, gpu_index_submitted_len=gpu_index_submitted_len) req.mtp_step = 0 req.filter_mark = req.wait_pause = req.paused = req.infer_aborted = False req.finish_status = infer_batch.FinishStatus() @@ -443,11 +480,11 @@ def _make_scheduled_req(row, total_len=12, moved=4): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_scheduler_submits_native_batch_on_inference_stream_from_default_stream(monkeypatch): +def test_scheduler_submits_native_batch_on_upload_stream_and_returns_to_default_stream(monkeypatch): from cuda.bindings import runtime as cuda_runtime manager = _make_manager("cuda") - req = _make_scheduled_req(0, moved=0) + req = _make_scheduled_req(0, gpu_index_submitted_len=0) backend, stream, _ = _make_scheduler(monkeypatch, manager, [req]) submitted_streams = [] native_copy = cuda_runtime.cudaMemcpyBatchAsync @@ -460,33 +497,34 @@ def record_batch(*args): with torch.cuda.stream(torch.cuda.default_stream()): backend._get_classed_reqs(req_ids=[0]) assert torch.cuda.current_stream().cuda_stream == 0 - stream.synchronize() + torch.cuda.default_stream().synchronize() - assert submitted_streams == [stream.cuda_stream] - assert req.move_to_gpu_index_len == req.hold_kv_len + assert len(submitted_streams) == 1 + assert submitted_streams[0] not in (0, stream.cuda_stream) + assert req.token_index.gpu_index_submitted_len == req.hold_kv_len assert torch.equal( - manager.req_to_token_indexs[0, : req.hold_kv_len].cpu(), manager.req_to_token_indexs_cpu[0, : req.hold_kv_len] + manager.req_to_token_indexs[0, : req.hold_kv_len].cpu(), manager.token_indexes_cpu[0, : req.hold_kv_len] ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_scheduler_uploads_only_new_tail_and_reuses_uploaded_page(monkeypatch): manager = _make_manager("cuda") - manager.req_to_token_indexs_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) - manager.req_to_token_indexs[0, :4].copy_(manager.req_to_token_indexs_cpu[0, :4]) + manager.token_indexes_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) + manager.req_to_token_indexs[0, :4].copy_(manager.token_indexes_cpu[0, :4]) req = _make_scheduled_req(0) backend, stream, copies = _make_scheduler(monkeypatch, manager, [req]) - with torch.cuda.stream(stream): + with torch.cuda.stream(torch.cuda.default_stream()): prefill, decode = backend._get_classed_reqs(req_ids=[0]) - stream.synchronize() + torch.cuda.default_stream().synchronize() assert prefill == [req] and decode == [] assert copies == [[(0, 4, 8)]] - assert req.move_to_gpu_index_len == req.hold_kv_len == 8 + assert req.token_index.gpu_index_submitted_len == req.hold_kv_len == 8 assert manager.req_to_token_indexs[0, :8].cpu().tolist() == [0, 1, 2, 3, 100, 101, 102, 103] req.cur_kv_len = 5 # The next chunk still fits into the reserved page. - with torch.cuda.stream(stream): + with torch.cuda.stream(torch.cuda.default_stream()): backend._get_classed_reqs(req_ids=[0]) assert copies == [[(0, 4, 8)]] @@ -494,53 +532,102 @@ def test_scheduler_uploads_only_new_tail_and_reuses_uploaded_page(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_scheduler_leaves_budget_skipped_request_unpublished(monkeypatch): manager = _make_manager("cuda") - manager.req_to_token_indexs_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) - req = _make_scheduled_req(0, moved=0) - req.hold_kv_len = 8 + manager.token_indexes_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) + req = _make_scheduled_req(0, gpu_index_submitted_len=0, hold_kv_len=8) backend, stream, copies = _make_scheduler(monkeypatch, manager, [req], budget=2) assert backend._get_classed_reqs(req_ids=[0]) == ([], []) - assert req.move_to_gpu_index_len == 0 + assert req.token_index.gpu_index_submitted_len == 0 assert copies == [] assert manager.req_to_token_indexs.eq(-1).all() backend.batch_max_tokens = 3 - with torch.cuda.stream(stream): + with torch.cuda.stream(torch.cuda.default_stream()): backend._get_classed_reqs(req_ids=[0]) - stream.synchronize() + torch.cuda.default_stream().synchronize() assert copies == [[(0, 0, 8)]] - assert req.move_to_gpu_index_len == 8 + assert req.token_index.gpu_index_submitted_len == 8 -def test_pd_transfer_failure_truncates_uploaded_length(monkeypatch): +@pytest.mark.parametrize("shared_prefix_len", [0, 4]) +@pytest.mark.parametrize("is_master_in_dp", [False, True]) +def test_pd_transfer_failure_defers_owned_tail_cleanup(monkeypatch, shared_prefix_len, is_master_in_dp): from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl manager = _make_manager() - manager.req_to_token_indexs_cpu[0, :8].copy_(torch.arange(8, dtype=torch.int32)) + manager.token_indexes_cpu[0, :12].copy_(torch.arange(12, dtype=torch.int32)) + shared_node = SimpleNamespace(node_prefix_total_len=shared_prefix_len) if shared_prefix_len else None + simulated_finishes = [] req = SimpleNamespace( req_idx=0, req_id=0, infer_aborted=False, - pd_task_num=1, + pd_task_num=2, pd_task_failed_num=1, pd_task_success_num=0, - finish_status=infer_batch.FinishStatus(infer_batch.FinishStatus.FINISHED_ERROR), - shared_kv_node=SimpleNamespace(node_prefix_total_len=4), - cur_kv_len=6, - hold_kv_len=8, - move_to_gpu_index_len=8, + finish_status=infer_batch.FinishStatus(), + shared_kv_node=shared_node, + cur_kv_len=10, + hold_kv_len=12, + cur_output_len=3, + cache_tiers=(infer_batch.CacheTier.GPU,), + shm_req=SimpleNamespace( + shm_cur_kv_len=10, + mark_simulated_finished=lambda status, output_len: simulated_finishes.append((status, output_len)), + ), + get_input_token_ids=lambda: list(range(12)), ) + bind_cpu_indexes(req, manager, gpu_index_submitted_len=12) + original_token_index_len = req.token_index.token_index_len + original_buffer = req.token_index._cpu_buffer monkeypatch.setattr(decode_impl, "g_infer_context", SimpleNamespace(requests_mapping={0: req})) freed = [] + manager.mem_manager = SimpleNamespace(free=lambda indexes: freed.extend(indexes.tolist())) backend = decode_impl.PDDecodeNode.__new__(decode_impl.PDDecodeNode) - backend.is_master_in_dp = False - backend.model = SimpleNamespace( - req_manager=manager, mem_manager=SimpleNamespace(free=lambda indexes: freed.extend(indexes.tolist())) - ) + backend.is_master_in_dp = is_master_in_dp + backend.model = SimpleNamespace(req_manager=manager, mem_manager=manager.mem_manager) + + # 任意传输任务仍未结束时，不能结束请求、回退逻辑长度或释放物理容量。 + assert backend._filter_not_ready_reqs([0]) == [] + assert not req.finish_status.is_finished() + assert req.cur_kv_len == req.shm_req.shm_cur_kv_len == 10 + assert freed == simulated_finishes == [] + + req.pd_task_success_num = 1 + for _ in range(2): + assert backend._filter_not_ready_reqs([0]) == [req] + assert req.finish_status.get_status() == infer_batch.FinishStatus.FINISHED_ERROR + assert req.cur_kv_len == shared_prefix_len + assert req.shm_req.shm_cur_kv_len == (shared_prefix_len if is_master_in_dp else 10) + assert req.hold_kv_len == req.token_index.gpu_index_submitted_len == 12 + assert req.token_index.token_index_len == original_token_index_len + assert req.token_index._cpu_buffer is original_buffer + assert freed == [] + assert simulated_finishes == ([(infer_batch.FinishStatus.FINISHED_ERROR, 3)] if is_master_in_dp else []) + + # 走真实统一释放路径：只保留可信共享前缀，失败的完整页和页内预留尾部都应回收。 + cached, released_refs = [], [] - assert backend._filter_not_ready_reqs([0]) == [req] - assert freed == [4, 5, 6, 7] - assert req.cur_kv_len == req.hold_kv_len == req.move_to_gpu_index_len == 4 + def insert(key, indexes): + cached.append((key, indexes)) + return len(key), shared_node + + context = InferenceContext() + context.args = SimpleNamespace(page_size=4) + context.req_manager = manager + context.radix_cache = SimpleNamespace(insert=insert, dec_node_ref_counter=released_refs.append) + released = [] + context.free_a_req_mem(released, req) + manager.free_token(torch.cat(released)) + + assert freed == list(range(shared_prefix_len, 12)) + assert len(cached) == 1 + assert cached[0][0].tolist() == cached[0][1].tolist() == list(range(shared_prefix_len)) + assert released_refs == ([shared_node] if shared_node is not None else []) + assert req.shared_kv_node is None + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == 0 + assert req.token_index.token_index_len == req.token_index.gpu_index_submitted_len == 0 + assert req.token_index._cpu_buffer is None @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -548,24 +635,24 @@ def test_pd_transfer_failure_truncates_uploaded_length(monkeypatch): def test_scheduler_publishes_prefill_and_decode_before_return(monkeypatch, mixed): manager = _make_manager("cuda") for row in (0, 1): - manager.req_to_token_indexs_cpu[row, :4].copy_(torch.arange(row * 4, row * 4 + 4, dtype=torch.int32)) - prefill_req = _make_scheduled_req(0, moved=0) - decode_req = _make_scheduled_req(1, total_len=4, moved=0) + manager.token_indexes_cpu[row, :4].copy_(torch.arange(row * 4, row * 4 + 4, dtype=torch.int32)) + prefill_req = _make_scheduled_req(0, gpu_index_submitted_len=0) + decode_req = _make_scheduled_req(1, total_len=4, gpu_index_submitted_len=0) backend, stream, copies = _make_scheduler(monkeypatch, manager, [prefill_req, decode_req], mixed=mixed) - with torch.cuda.stream(stream): + with torch.cuda.stream(torch.cuda.default_stream()): prefill, decode = backend._get_classed_reqs(req_ids=[0, 1]) - stream.synchronize() + torch.cuda.default_stream().synchronize() assert copies == [[(0, 0, 8), (1, 0, 4)]] - assert prefill_req.move_to_gpu_index_len == 8 - assert decode_req.move_to_gpu_index_len == 4 + assert prefill_req.token_index.gpu_index_submitted_len == 8 + assert decode_req.token_index.gpu_index_submitted_len == 4 assert prefill == ([prefill_req, decode_req] if mixed else [prefill_req]) assert decode == ([] if mixed else [decode_req]) for req in (prefill_req, decode_req): assert torch.equal( manager.req_to_token_indexs[req.req_idx, : req.hold_kv_len].cpu(), - manager.req_to_token_indexs_cpu[req.req_idx, : req.hold_kv_len], + manager.token_indexes_cpu[req.req_idx, : req.hold_kv_len], ) @@ -575,7 +662,10 @@ def test_request_upload_lengths_follow_fallback_copy_result(monkeypatch, loop_fa from cuda.bindings import runtime as cuda_runtime manager = _make_manager("cuda") - reqs = [_make_scheduled_req(0, moved=0), _make_scheduled_req(1, total_len=4, moved=0)] + reqs = [ + _make_scheduled_req(0, gpu_index_submitted_len=0), + _make_scheduled_req(1, total_len=4, gpu_index_submitted_len=0), + ] backend, stream, _ = _make_scheduler(monkeypatch, manager, reqs) monkeypatch.setattr( cuda_runtime, @@ -592,18 +682,20 @@ def failed_copy(dst, src, *args, **kwargs): return torch_copy(dst, src, *args, **kwargs) monkeypatch.setattr(torch.Tensor, "copy_", failed_copy) - with torch.cuda.stream(stream), pytest.raises(RuntimeError, match="individual copy failed"): + with torch.cuda.stream(torch.cuda.default_stream()), pytest.raises( + RuntimeError, match="individual copy failed" + ): backend._get_classed_reqs(req_ids=[0, 1]) - assert [req.move_to_gpu_index_len for req in reqs] == [0, 0] + assert [req.token_index.gpu_index_submitted_len for req in reqs] == [0, 0] else: - with torch.cuda.stream(stream): + with torch.cuda.stream(torch.cuda.default_stream()): backend._get_classed_reqs(req_ids=[0, 1]) - stream.synchronize() + torch.cuda.default_stream().synchronize() for req in reqs: - assert req.move_to_gpu_index_len == req.hold_kv_len + assert req.token_index.gpu_index_submitted_len == req.hold_kv_len assert torch.equal( manager.req_to_token_indexs[req.req_idx, : req.hold_kv_len].cpu(), - manager.req_to_token_indexs_cpu[req.req_idx, : req.hold_kv_len], + manager.token_indexes_cpu[req.req_idx, : req.hold_kv_len], ) @@ -612,16 +704,16 @@ def failed_copy(dst, src, *args, **kwargs): def test_copy_uses_current_stream_on_current_device(monkeypatch, batch_available): from lightllm.common import cuda_wrapper as base - monkeypatch.setattr(base, "_cuda_memcpy_batch_available", batch_available) + monkeypatch.setattr(base.PinnedH2DBatchCopier, "_native_batch_available", batch_available) manager = _make_manager("cuda:1") - manager.req_to_token_indexs_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) + manager.token_indexes_cpu[0, :4].copy_(torch.arange(4, dtype=torch.int32)) stream = torch.cuda.Stream(device=1) stream.wait_stream(torch.cuda.current_stream(1)) with torch.cuda.device(0): with torch.cuda.stream(stream): cuda_memcpy_pinned_h2d_batch_async( [manager.req_to_token_indexs[0, :4]], - [manager.req_to_token_indexs_cpu[0, :4]], + [manager.token_indexes_cpu[0, :4]], ) assert torch.cuda.current_device() == 1 assert torch.cuda.current_stream().cuda_stream == stream.cuda_stream diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index ea057a481d..1dcf637ff8 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -8,6 +8,7 @@ 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 +from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes class _FakeMemManager: @@ -28,15 +29,15 @@ 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_cpu = torch.full((3, 8), -1, dtype=torch.int32) + req_manager.page_size = 4 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_cpu[2].tolist() == [100, 101, 102, 103] * 2 - assert req_manager.req_to_token_indexs_cpu[:2].eq(-1).all() + 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(): @@ -49,11 +50,10 @@ def test_hold_request_indexes_require_all_requests_released(): def _make_context(monkeypatch): mem_manager = _FakeMemManager() - req_manager = ReqManager.__new__(ReqManager) + req_manager = make_cpu_index_manager(torch.full((2, 16), -1, dtype=torch.int32), page_size=4) req_manager.mem_manager = mem_manager req_manager.max_request_num = 2 - req_manager.req_to_token_indexs_cpu = torch.full((2, 16), -1, dtype=torch.int32) - req_manager.req_to_token_indexs = torch.full((2, 16), -1, dtype=torch.int32) + req_manager.page_size = 4 context = SimpleNamespace( args=SimpleNamespace(page_size=4), @@ -63,23 +63,27 @@ def _make_context(monkeypatch): ) monkeypatch.setattr(generic_pre_process, "g_infer_context", context) monkeypatch.setattr(base_backend, "g_infer_context", context) + monkeypatch.setattr(infer_batch, "g_infer_context", context) backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) backend.args = context.args def copy_indexes(run_reqs): for req in run_reqs: - start, end = req.move_to_gpu_index_len, req.hold_kv_len + start, end = req.token_index.gpu_index_submitted_len, req.hold_kv_len req_manager.req_to_token_indexs[req.req_idx, start:end].copy_( - req_manager.req_to_token_indexs_cpu[req.req_idx, start:end] + req.token_index.get_cpu_token_indexes(start, end) ) - req.move_to_gpu_index_len = end + req.token_index.gpu_index_submitted_len = end backend._copy_req_token_indexes_to_gpu = copy_indexes return context, backend def _make_req(req_idx): - req = SimpleNamespace(req_idx=req_idx, cur_kv_len=0, hold_kv_len=0, move_to_gpu_index_len=0) + req = InferReq.__new__(InferReq) + req.req_idx = req_idx + bind_cpu_indexes(req, infer_batch.g_infer_context.req_manager) + req.cur_kv_len = req.hold_kv_len = 0 req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) return req @@ -92,7 +96,7 @@ def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypat 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_cpu[0, :4].tolist() == [0, 1, 2, 3] + assert context.req_manager.get_cpu_token_indexes(0, 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 @@ -102,7 +106,7 @@ def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypat 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_cpu[0, :8].tolist() == list(range(8)) + assert context.req_manager.get_cpu_token_indexes(0, 0, 8).tolist() == list(range(8)) def test_reservation_fills_each_request_table_row(monkeypatch): @@ -113,8 +117,8 @@ def test_reservation_fills_each_request_table_row(monkeypatch): 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_cpu[0, :4].tolist() == [0, 1, 2, 3] - assert generic_pre_process.g_infer_context.req_manager.req_to_token_indexs_cpu[1, :4].tolist() == [4, 5, 6, 7] + assert generic_pre_process.g_infer_context.req_manager.get_cpu_token_indexes(0, 0, 4).tolist() == [0, 1, 2, 3] + assert generic_pre_process.g_infer_context.req_manager.get_cpu_token_indexes(1, 0, 4).tolist() == [4, 5, 6, 7] assert req0.hold_kv_len == req1.hold_kv_len == 4 @@ -125,9 +129,9 @@ def test_alloc_req_kv_mem_defers_gpu_table_update(monkeypatch): backend._alloc_req_kv_mem(req, alloc_token_num=4) backend._alloc_req_kv_mem(req, alloc_token_num=4) - assert context.req_manager.req_to_token_indexs_cpu[0, :8].tolist() == list(range(8)) + assert context.req_manager.get_cpu_token_indexes(0, 0, 8).tolist() == list(range(8)) assert context.req_manager.req_to_token_indexs.eq(-1).all() - assert req.move_to_gpu_index_len == 0 + assert req.token_index.gpu_index_submitted_len == 0 def test_need_token_num_distinguishes_compute_and_page_allocation(): @@ -199,8 +203,8 @@ def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( 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_cpu[req.req_idx, :hold_kv_len] = torch.arange( - start_index, start_index + hold_kv_len, dtype=torch.int32 + context.req_manager.append_token_indexes( + req.req_idx, 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 @@ -243,7 +247,7 @@ def test_decode_scheduler_reserves_cpu_request_table(monkeypatch): 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_cpu[0, :4] = torch.arange(4, dtype=torch.int32) + context.req_manager.append_token_indexes(0, torch.arange(4, dtype=torch.int32)) context.req_manager.mem_manager.next_index = 4 backend._filter_not_ready_reqs = lambda req_ids: [req] @@ -264,7 +268,8 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step ): context, backend = _make_context(monkeypatch) context.args.page_size = page_size - context.req_manager.req_to_token_indexs_cpu = torch.full((2, 64), -1, dtype=torch.int32) + context.req_manager.page_size = page_size + context.req_manager.token_indexes_cpu = torch.full((2, (64 + page_size - 1) // page_size), -1, dtype=torch.int32) 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 @@ -292,7 +297,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step 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_cpu[0, :page_size] = torch.arange(page_size, dtype=torch.int32) + context.req_manager.append_token_indexes(0, 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] @@ -317,7 +322,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): req.get_radix_cache_shared_len = lambda: 0 req.args = context.args - context.req_manager.req_to_token_indexs_cpu[0, :4] = torch.arange(4, dtype=torch.int32) + context.req_manager.append_token_indexes(0, 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 @@ -330,12 +335,14 @@ def test_decode_reserves_mtp_headroom(monkeypatch): 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_cpu[0, :12].tolist() == list(range(12)) + assert context.req_manager.get_cpu_token_indexes(0, 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 + context.req_manager.page_size = 1 + context.req_manager.token_indexes_cpu = torch.full((2, 16), -1, dtype=torch.int32) req = _make_req(0) req.cur_kv_len = 3 req.hold_kv_len = 3 @@ -345,7 +352,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): 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_cpu[0, :3] = torch.arange(3, dtype=torch.int32) + context.req_manager.append_token_indexes(0, 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 @@ -355,33 +362,39 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): assert req.hold_kv_len == 9 assert context.req_manager.mem_manager.alloc_sizes == [6] - assert context.req_manager.req_to_token_indexs_cpu[0, :9].tolist() == list(range(9)) + assert context.req_manager.get_cpu_token_indexes(0, 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_cpu=torch.arange(16, dtype=torch.int32)[None, :]) + infer_context.req_manager = make_cpu_index_manager( + token_table=torch.arange(16, dtype=torch.int32)[None, :], page_size=1 + ) req = SimpleNamespace( req_idx=0, cur_kv_len=3, hold_kv_len=12, shm_req=SimpleNamespace(shm_cur_kv_len=3), ) + bind_cpu_indexes(req, infer_context.req_manager) 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 + assert req.token_index.token_index_len == req.token_index.gpu_index_submitted_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_cpu=torch.arange(16, dtype=torch.int32)[None, :]) + infer_context.req_manager = make_cpu_index_manager( + token_table=torch.arange(16, dtype=torch.int32)[None, :], page_size=1 + ) req = SimpleNamespace( req_idx=0, cur_kv_len=5, @@ -391,6 +404,7 @@ def test_linear_attention_frees_reserved_page_tail(monkeypatch): hybrid_len_to_big_page_id={}, shared_kv_node=None, ) + bind_cpu_indexes(req, infer_context.req_manager) 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) @@ -404,7 +418,9 @@ def test_linear_attention_frees_preallocated_pages_before_first_forward(monkeypa 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_cpu=torch.arange(8, dtype=torch.int32)[None, :]) + infer_context.req_manager = make_cpu_index_manager( + token_table=torch.arange(8, dtype=torch.int32)[None, :], page_size=1 + ) req = SimpleNamespace( req_idx=0, cur_kv_len=0, @@ -414,6 +430,7 @@ def test_linear_attention_frees_preallocated_pages_before_first_forward(monkeypa hybrid_len_to_big_page_id={}, shared_kv_node=None, ) + bind_cpu_indexes(req, infer_context.req_manager) 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) @@ -433,7 +450,6 @@ def _make_paused_req(req_idx, target_kv_len, page_size, match_kv_len=0, chunk_kv cur_kv_len=0, hold_kv_len=0, paused=True, - move_to_gpu_index_len=0, match_call_count=0, shared_kv_node=None, shm_req=SimpleNamespace(is_paused=True, shm_cur_kv_len=0), @@ -459,9 +475,10 @@ def test_recover_paused_reqs_uses_page_allocation_need(monkeypatch): infer_context.backend = SimpleNamespace(disable_chunked_prefill=False) infer_context.radix_cache = None freed_indexes = [] - infer_context.req_manager = SimpleNamespace( - req_to_token_indexs_cpu=torch.arange(16, dtype=torch.int32).reshape(2, 8), + infer_context.req_manager = make_cpu_index_manager( + token_table=torch.arange(16, dtype=torch.int32).reshape(2, 8), free_token=lambda indexes: freed_indexes.extend(indexes.tolist()), + page_size=1, ) infer_context.get_can_alloc_token_num = lambda: 5 large_req = _make_paused_req(req_idx=0, target_kv_len=5, page_size=4) @@ -483,9 +500,8 @@ 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_cpu=torch.arange(8, dtype=torch.int32)[None, :], - free_token=lambda indexes: None, + infer_context.req_manager = make_cpu_index_manager( + token_table=torch.arange(8, dtype=torch.int32)[None, :], free_token=lambda indexes: None, page_size=1 ) 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) diff --git a/unit_tests/common/test_req_token_index.py b/unit_tests/common/test_req_token_index.py new file mode 100644 index 0000000000..16f0e1fc4d --- /dev/null +++ b/unit_tests/common/test_req_token_index.py @@ -0,0 +1,594 @@ +from types import SimpleNamespace +import weakref + +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 +from lightllm.common.req_manager import base +from lightllm.server.router.model_infer.infer_req_ext1 import ( + PinMemCacheReqTokenIndex, + PureGpuReqTokenIndex, + create_req_token_index, +) +from lightllm.server.router.model_infer.infer_batch import InferReq +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.core.objs.start_args_type import StartArgs + + +MODES = ["pin_mem_cache", "pure_gpu"] +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and pinned CPU memory") + + +def make_req(token_table, mode="pin_mem_cache", req_idx=0, input_len=0): + req = InferReq.__new__(InferReq) + req.req_idx = req_idx + req.req_manager = SimpleNamespace(req_to_token_indexs=token_table) + req.args = StartArgs() if mode is None else StartArgs(req_token_index_mode=mode) + req.shm_req = SimpleNamespace(input_len=input_len) + return req + + +@pytest.mark.parametrize( + ("mode", "expected_type"), + [ + (None, PinMemCacheReqTokenIndex), + ("pin_mem_cache", PinMemCacheReqTokenIndex), + ("pure_gpu", PureGpuReqTokenIndex), + ], +) +def test_factory_uses_request_mode_and_row_view_without_retaining_request(mode, expected_type): + table = torch.arange(24, dtype=torch.int32).view(3, 8) + req = make_req(table, mode, req_idx=1) + # 模式来自请求参数；manager 的相反配置不能影响该请求的选择。 + req.req_manager.req_token_index_mode = "pin_mem_cache" if mode == "pure_gpu" else "pure_gpu" + + # 构造只保存传入的流，不执行 CUDA 操作；用占位对象验证透传即可。 + index_stream = object() + helper = create_req_token_index(req, index_stream=index_stream) + + assert not hasattr(InferReq, "move_to_gpu_index_len") + assert isinstance(helper, expected_type) + assert helper._index_stream is index_stream + assert helper.req_idx == req.req_idx + assert helper.gpu_indexes.data_ptr() == table[req.req_idx].data_ptr() + expected = table.clone() + expected[req.req_idx, 3] = -99 + helper.gpu_indexes[3] = -99 + assert torch.equal(table, expected) + # 模拟真实持有关系；保留 Helper 时，请求也应能立即释放，无需循环 GC。 + req.token_index = helper + request_ref = weakref.ref(req) + del req + assert request_ref() is None + + +@pytest.mark.parametrize("mode", MODES) +def test_start_args_and_cli_select_mode(mode): + parser = make_argument_parser() + assert parser.parse_args([]).req_token_index_mode == "pin_mem_cache" + assert StartArgs().req_token_index_mode == "pin_mem_cache" + assert parser.parse_args(["--req_token_index_mode", mode]).req_token_index_mode == mode + assert StartArgs(req_token_index_mode=mode).req_token_index_mode == mode + with pytest.raises(SystemExit): + parser.parse_args(["--req_token_index_mode", "unknown"]) + with pytest.raises(ValueError, match="未知请求 token 索引模式"): + create_req_token_index(make_req(torch.empty((1, 8), dtype=torch.int32), "unknown"), index_stream=object()) + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("pinned", [False, True]) +def test_append_accepts_cpu_input_and_only_pins_when_gpu_upload_requires_it(monkeypatch, mode, pinned): + table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") + helper = create_req_token_index(make_req(table, mode), index_stream=torch.cuda.current_stream()) + indexes = torch.arange(4, dtype=torch.int32, pin_memory=pinned) + pin_memory = torch.Tensor.pin_memory + pinned_inputs = [] + + def record_pin_memory(tensor, *args, **kwargs): + pinned_inputs.append(tensor) + return pin_memory(tensor, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "pin_memory", record_pin_memory) + helper.append_token_indexes(indexes) + # 普通 CPU 输入已经复制到自有 storage，可在 append 返回后立即改写。 + if not pinned: + indexes.fill_(-99) + + assert len(pinned_inputs) == (1 if mode == "pure_gpu" and not pinned else 0) + assert helper.token_index_len == 4 + assert helper.get_cpu_token_indexes(0, 4).tolist() == [0, 1, 2, 3] + assert helper.get_gpu_token_indexes(0, 4).cpu().tolist() == [0, 1, 2, 3] + if mode == "pin_mem_cache": + assert helper._cpu_buffer.is_pinned() + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +def test_append_rejects_gpu_input_without_changing_index_state(mode): + table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") + helper = create_req_token_index(make_req(table, mode), index_stream=torch.cuda.current_stream()) + + with pytest.raises(AssertionError, match="mem_indexes 必须是 CPU Tensor"): + helper.append_token_indexes(torch.arange(4, dtype=torch.int32, device="cuda")) + + assert helper.token_index_len == helper.gpu_index_submitted_len == 0 + assert table.eq(-1).all() + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +def test_gpu_getter_rejects_unrecorded_indexes_without_advancing_progress(mode): + table = torch.full((1, 8), -7, dtype=torch.int32, device="cuda") + helper = create_req_token_index(make_req(table, mode), index_stream=torch.cuda.current_stream()) + helper.append_token_indexes(torch.tensor([10, 11], dtype=torch.int32, pin_memory=True)) + submitted_len = helper.gpu_index_submitted_len + + with pytest.raises(AssertionError, match="读取范围超出"): + helper.get_gpu_token_indexes(0, 4) + + assert helper.token_index_len == 2 + assert helper.gpu_index_submitted_len == submitted_len + assert helper.get_gpu_token_indexes(0, 2).cpu().tolist() == [10, 11] + + +@requires_cuda +def test_batch_upload_rejects_missing_gpu_submission_without_advancing_progress(monkeypatch): + from lightllm.server.router.model_infer.mode_backend import base_backend + + table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") + req = make_req(table, "pure_gpu") + req.token_index = create_req_token_index(req, index_stream=torch.cuda.current_stream()) + # 模拟不一致状态：记录了索引，但没有提交 GPU 写入；批量入口不能替它补写进度。 + req.token_index.token_index_len = 4 + monkeypatch.setattr(base_backend, "g_infer_context", SimpleNamespace(req_manager=req.req_manager)) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + + with pytest.raises(AssertionError, match="GPU 索引提交进度与记录长度不一致"): + backend._copy_req_token_indexes_to_gpu([req]) + + assert req.token_index.gpu_index_submitted_len == 0 + assert not hasattr(backend, "_index_upload_stream") + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +def test_backend_rejects_requests_with_different_index_streams(monkeypatch, mode): + from lightllm.server.router.model_infer.mode_backend import base_backend + + table = torch.full((2, 8), -1, dtype=torch.int32, device="cuda") + reqs = [] + for row in range(2): + index_stream = torch.cuda.Stream() + index_stream.wait_stream(torch.cuda.current_stream()) + req = make_req(table, mode, req_idx=row) + req.token_index = create_req_token_index(req, index_stream=index_stream) + req.token_index.append_token_indexes(torch.arange(4, dtype=torch.int32, pin_memory=True)) + reqs.append(req) + + def unexpected_upload(*args): + pytest.fail("索引流不一致时应在提交上传前拒绝该批请求") + + monkeypatch.setattr(base_backend, "cuda_memcpy_pinned_h2d_batch_async", unexpected_upload) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + with pytest.raises(AssertionError, match="同一批请求必须使用同一个索引流"): + backend._copy_req_token_indexes_to_gpu(reqs) + assert [req.token_index.gpu_index_submitted_len for req in reqs] == [0 if mode == "pin_mem_cache" else 4] * 2 + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("page_size", [1, 3, 4, 64]) +def test_token_units_allow_partial_page_appends_and_snapshot_ownership(mode, page_size): + row = torch.full((3 * page_size + 5,), -1, dtype=torch.int32, device="cuda") + req = make_req(row.unsqueeze(0), mode) + helper = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + req.token_index = helper + count = row.numel() + original = (torch.arange(count, dtype=torch.int32) * 3).pin_memory() + split = page_size + 1 + req.token_index.append_token_indexes(original[:split]) + req.token_index.append_token_indexes(torch.empty(0, dtype=torch.int32)) + assert helper.token_index_len == split + req.token_index.append_token_indexes(original[split:]) + assert helper.token_index_len == count + expected = torch.arange(count, dtype=torch.int32) * 3 + for start, end in [(0, count), (1, count - 1), (page_size, page_size + 2), (count, count)]: + assert torch.equal(req.token_index.get_cpu_token_indexes(start, end), expected[start:end]) + snapshot = req.token_index.get_cpu_token_indexes(0, count) + # pure_gpu 的阻塞 CPU getter 已等待 H2D 完成，此后才可主动改写上传源。 + original.fill_(99) + # GPU 读取区间只决定返回视图；首次读取同时提交内部记录的全部待上传索引。 + assert torch.equal(req.token_index.get_gpu_token_indexes(1, 3).cpu(), expected[1:3]) + assert helper.gpu_index_submitted_len == count + gpu = req.token_index.get_gpu_token_indexes(0, count) + assert gpu.data_ptr() == row.data_ptr() + assert torch.equal(gpu.cpu(), expected) + if mode == "pin_mem_cache": + cpu_buffer = helper._cpu_buffer + helper.reset() + assert helper.token_index_len == helper.gpu_index_submitted_len == 0 + if mode == "pin_mem_cache": + assert helper._cpu_buffer is None + else: + assert not hasattr(helper, "_cpu_buffer") + helper.append_token_indexes(torch.zeros(count, dtype=torch.int32, pin_memory=True)) + if mode == "pin_mem_cache": + assert helper._cpu_buffer is not cpu_buffer + assert helper.get_gpu_token_indexes(0, count).cpu().eq(0).all() + assert torch.equal(snapshot, torch.arange(count, dtype=torch.int32) * 3) + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +def test_free_getter_borrows_gpu_indexes_but_metadata_getter_keeps_cpu_snapshot(monkeypatch, mode): + table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") + req = make_req(table, mode) + req.token_index = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + req.token_index.append_token_indexes(torch.arange(6, dtype=torch.int32, pin_memory=True)) + snapshot = req.token_index.get_cpu_token_indexes(1, 5) + + def unexpected_snapshot(*args): + pytest.fail("pure_gpu 的释放接口不能调用普通 CPU 快照 getter") + + with monkeypatch.context() as no_snapshot: + if mode == "pure_gpu": + no_snapshot.setattr(req.token_index, "get_cpu_token_indexes", unexpected_snapshot) + free_indexes = req.token_index.get_free_token_indexes(1, 5) + empty = req.token_index.get_free_token_indexes(3, 3) + + expected_device = "cuda" if mode == "pure_gpu" else "cpu" + assert free_indexes.device.type == empty.device.type == expected_device + assert empty.numel() == 0 + if mode == "pure_gpu": + assert free_indexes.data_ptr() == table[0, 1:5].data_ptr() + updated = torch.arange(101, 105, dtype=torch.int32) + replacement = torch.arange(6, dtype=torch.int32, pin_memory=True) + replacement[1:5] = updated + req.token_index.reset() + req.token_index.append_token_indexes(replacement) + + assert snapshot.device.type == "cpu" + assert snapshot.tolist() == [1, 2, 3, 4] + assert free_indexes.cpu().tolist() == (updated.tolist() if mode == "pure_gpu" else [1, 2, 3, 4]) + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +def test_reset_rebuild_keeps_cpu_appends_and_gpu_submission_separate(mode): + helper = create_req_token_index( + make_req(torch.zeros((1, 8), dtype=torch.int32, device="cuda"), mode), index_stream=torch.cuda.default_stream() + ) + pending = helper.get_upload_indexes() + if mode == "pure_gpu": + assert pending is None + else: + assert pending.numel() == 0 + helper.append_token_indexes(torch.arange(4, dtype=torch.int32, pin_memory=True)) + pending = helper.get_upload_indexes() + if mode == "pure_gpu": + assert pending is None + else: + assert pending.tolist() == [0, 1, 2, 3] + assert helper.gpu_index_submitted_len == (0 if mode == "pin_mem_cache" else 4) + before = helper.get_cpu_token_indexes(0, 4) + helper.get_gpu_token_indexes(0, 4) + pending = helper.get_upload_indexes() + if mode == "pure_gpu": + assert pending is None + else: + assert pending.numel() == 0 + helper.reset() + helper.append_token_indexes(torch.tensor([0, 31], dtype=torch.int32, pin_memory=True)) + assert helper.token_index_len == 2 + assert helper.gpu_index_submitted_len == (0 if mode == "pin_mem_cache" else 2) + if mode == "pin_mem_cache": + assert helper.gpu_indexes[:4].cpu().tolist() == [0, 1, 2, 3] + helper.append_token_indexes(torch.tensor([32, 3], dtype=torch.int32, pin_memory=True)) + assert helper.token_index_len == 4 + assert helper.gpu_index_submitted_len == (0 if mode == "pin_mem_cache" else 4) + assert helper.get_gpu_token_indexes(0, 4).cpu().tolist() == [0, 31, 32, 3] + assert helper.gpu_index_submitted_len == 4 + assert before.tolist() == [0, 1, 2, 3] + + +@requires_cuda +@pytest.mark.parametrize( + ("input_len", "max_length", "initial_capacity", "next_end", "grown_capacity"), + [ + (1023, 16384, 1024, 1025, 4096), + (1024, 16384, 1024, 1025, 4096), + (1025, 16384, 2048, 2049, 5120), + (2048, 16384, 2048, 5000, 7168), + (1025, 3500, 2048, 3000, 3500), + ], +) +def test_pin_buffer_reserves_input_length_and_grows_with_headroom( + input_len, max_length, initial_capacity, next_end, grown_capacity +): + req = make_req(torch.empty((1, max_length), dtype=torch.int32, device="cuda"), input_len=input_len) + helper = PinMemCacheReqTokenIndex(req, index_stream=torch.cuda.default_stream()) + assert helper.input_len == input_len + assert helper._cpu_buffer is None + helper.append_token_indexes(torch.empty(0, dtype=torch.int32)) + assert helper._cpu_buffer is None + # 首次只追加一个索引，仍按完整输入长度预留。 + first_index = torch.zeros(1, dtype=torch.int32, pin_memory=True) + helper.append_token_indexes(first_index) + assert helper._cpu_buffer.numel() == initial_capacity + assert helper._cpu_buffer.is_pinned() and helper._cpu_buffer.data_ptr() != first_index.data_ptr() + helper.append_token_indexes(torch.arange(1, initial_capacity, dtype=torch.int32, pin_memory=True)) + old_buffer = helper._cpu_buffer + helper.append_token_indexes(torch.arange(initial_capacity, next_end, dtype=torch.int32, pin_memory=True)) + assert helper._cpu_buffer.numel() == grown_capacity and helper._cpu_buffer.is_pinned() + assert helper._cpu_buffer.data_ptr() != old_buffer.data_ptr() + snapshot = helper.get_cpu_token_indexes(0, next_end) + assert snapshot.tolist() == list(range(next_end)) + grown_buffer = helper._cpu_buffer + helper.reset() + assert helper.input_len == input_len and helper._cpu_buffer is None + assert helper.token_index_len == helper.gpu_index_submitted_len == 0 + helper.append_token_indexes(torch.tensor([99], dtype=torch.int32, pin_memory=True)) + assert helper._cpu_buffer is not grown_buffer + assert helper._cpu_buffer.numel() == initial_capacity + assert helper.token_index_len == 1 + assert helper.get_cpu_token_indexes(0, helper.token_index_len).tolist() == [99] + assert snapshot.tolist() == list(range(next_end)) + + +@requires_cuda +@pytest.mark.parametrize("use_default_stream", [False, True]) +def test_pin_buffer_reset_during_upload_keeps_old_data_until_h2d_completes(use_default_stream): + count = 8192 + row = torch.full((count,), -1, dtype=torch.int32, device="cuda") + index_stream = torch.cuda.default_stream() if use_default_stream else torch.cuda.Stream() + index_stream.wait_stream(torch.cuda.current_stream()) + helper = PinMemCacheReqTokenIndex(make_req(row.unsqueeze(0), input_len=count), index_stream=index_stream) + original = torch.arange(count, dtype=torch.int32, pin_memory=True) + replacement = torch.full((count,), -99, dtype=torch.int32, pin_memory=True) + helper.append_token_indexes(original) + source_ref = weakref.ref(helper._cpu_buffer) + + with torch.cuda.stream(index_stream): + torch.cuda._sleep(200_000_000) + helper.get_gpu_token_indexes(0, count) + assert not index_stream.query(), "延迟 H2D 应覆盖 reset 和重新追加的过程" + helper.reset() + assert helper._cpu_buffer is None and source_ref() is None + helper.append_token_indexes(replacement) + + # reset 已释放旧 Tensor 的 Python 引用；新缓存写入不能污染仍在传输的旧数据。 + index_stream.synchronize() + assert torch.equal(row.cpu(), original) + assert torch.equal(helper.get_gpu_token_indexes(0, count).cpu(), replacement) + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("pre_upload", [False, True]) +@pytest.mark.parametrize("use_default_stream", [False, True]) +def test_gpu_getter_uses_injected_stream_and_orders_calling_stream(monkeypatch, mode, pre_upload, use_default_stream): + row = torch.full((16,), -1, dtype=torch.int32, device="cuda") + index_stream = torch.cuda.default_stream() if use_default_stream else torch.cuda.Stream() + index_stream.wait_stream(torch.cuda.current_stream()) + reader = torch.cuda.Stream() + with torch.cuda.stream(reader): + # 构造时的当前流与传入的索引流不同，Helper 必须始终使用外部指定的流。 + helper = create_req_token_index(make_req(row.unsqueeze(0), mode), index_stream=index_stream) + assert helper._index_stream is index_stream + assert torch.cuda.current_stream() == reader + copy = torch.Tensor.copy_ + cpu = torch.Tensor.cpu + uploads = [] + + def record_copy(dst, src, *args, **kwargs): + if dst.is_cuda and src.device.type == "cpu": + assert torch.cuda.current_stream() == index_stream + uploads.append(src.numel()) + return copy(dst, src, *args, **kwargs) + + def check_cpu_read(src, *args, **kwargs): + if src.is_cuda: + assert torch.cuda.current_stream() == index_stream + return cpu(src, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "copy_", record_copy) + monkeypatch.setattr(torch.Tensor, "cpu", check_cpu_read) + with torch.cuda.stream(index_stream): + torch.cuda._sleep(30_000_000) + with torch.cuda.stream(reader): + helper.append_token_indexes(torch.arange(8, dtype=torch.int32, pin_memory=True)) + assert torch.cuda.current_stream() == reader + if pre_upload: + # 索引流已提交上传，但尚未完成；另一个流读取时仍必须承接依赖。 + helper.get_gpu_token_indexes(0, 8) + with torch.cuda.stream(reader): + indexes = helper.get_gpu_token_indexes(0, 8) + assert torch.cuda.current_stream() == reader + actual = indexes.clone() + assert helper.get_cpu_token_indexes(0, 8).tolist() == list(range(8)) + assert torch.cuda.current_stream() == reader + reader.synchronize() + assert uploads == [8] + with torch.cuda.stream(index_stream): + assert actual.cpu().tolist() == list(range(8)) + + +@requires_cuda +def test_pure_gpu_slot_reuse_keeps_old_and_new_helper_writes_in_order(): + row = torch.full((8,), -1, dtype=torch.int32, device="cuda") + torch.cuda._sleep(30_000_000) + old = create_req_token_index(make_req(row.unsqueeze(0), "pure_gpu"), index_stream=torch.cuda.default_stream()) + old.append_token_indexes(torch.full((8,), 11, dtype=torch.int32, pin_memory=True)) + # 源 Tensor 引用已释放；先保留旧写入结果，验证后续 pinned 分配不会覆盖在途 H2D。 + first_write = row.clone() + new = create_req_token_index(make_req(row.unsqueeze(0), "pure_gpu"), index_stream=torch.cuda.default_stream()) + new.append_token_indexes(torch.full((8,), 22, dtype=torch.int32, pin_memory=True)) + assert new.get_cpu_token_indexes(0, 8).tolist() == [22] * 8 + assert first_write.cpu().tolist() == [11] * 8 + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("page_size", [1, 4, 64]) +def test_manager_owns_only_gpu_table_and_initializes_hold_page(monkeypatch, mode, page_size): + from lightllm.common.req_manager import req_sampling_params + + monkeypatch.setattr( + base, "get_env_start_args", lambda: SimpleNamespace(page_size=page_size, req_token_index_mode=mode) + ) + monkeypatch.setattr(req_sampling_params, "ReqSamplingParamsManager", lambda count: None) + mem_manager = SimpleNamespace(page_size=page_size, HOLD_TOKEN_MEMINDEXES=tuple(range(page_size * 8, page_size * 9))) + manager = base.ReqManager(2, page_size * 2 + 1, mem_manager) + assert not hasattr(manager, "req_to_page_indexs_cpu") + assert not hasattr(manager, "get_cpu_token_indexes") + assert manager.req_to_token_indexs.shape == (3, 3 * page_size) + assert ( + manager.req_to_token_indexs[manager.HOLD_REQUEST_ID].cpu().tolist() + == list(mem_manager.HOLD_TOKEN_MEMINDEXES) * 3 + ) + + +@requires_cuda +@pytest.mark.parametrize("batch_available", [False, True]) +@pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("index_stream_kind", ["default", "shared"]) +def test_backend_uploads_pinned_views_without_request_events_or_keepalive_list( + monkeypatch, mode, batch_available, index_stream_kind +): + from lightllm.common import cuda_wrapper + from lightllm.server.router.model_infer.mode_backend import base_backend + + table = torch.full((2, 16), -1, dtype=torch.int32, device="cuda") + expected = (torch.arange(16, dtype=torch.int32).view(2, 8) + 32).pin_memory() + index_streams = [torch.cuda.default_stream()] * 2 + if index_stream_kind == "shared": + index_streams = [torch.cuda.Stream()] * 2 + for index_stream in set(index_streams): + index_stream.wait_stream(torch.cuda.default_stream()) + reqs = [] + for row, index_stream in enumerate(index_streams): + req = make_req(table, mode, req_idx=row) + req.token_index = create_req_token_index(req, index_stream=index_stream) + with torch.cuda.stream(index_stream): + # 批量上传须先等待旧写入；pure_gpu 在本流追加，调用流尚未承接依赖。 + torch.cuda._sleep(60_000_000) + table[row, :8].fill_(-2) + req.token_index.append_token_indexes(expected[row]) + req.hold_kv_len = 8 + reqs.append(req) + monkeypatch.setattr(base_backend, "g_infer_context", SimpleNamespace(req_manager=reqs[0].req_manager)) + monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", batch_available) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + stream = index_streams[0] + stream.wait_stream(torch.cuda.default_stream()) + with torch.cuda.stream(stream): + torch.cuda._sleep(30_000_000) + events = [] + batch_calls = [] + copy_batch = base_backend.cuda_memcpy_pinned_h2d_batch_async + event_type = torch.cuda.Event + + def record_batch(dst_tensors, src_tensors): + # 直接借用各请求缓存的待上传区间，不能再创建聚合快照或改成逐请求上传。 + assert [tensor.data_ptr() for tensor in src_tensors] == [ + req.token_index.get_upload_indexes().data_ptr() for req in reqs + ] + batch_calls.append(len(src_tensors)) + return copy_batch(dst_tensors, src_tensors) + + def record_event(*args, **kwargs): + event = event_type(*args, **kwargs) + events.append(event) + return event + + with monkeypatch.context() as batch_events: + batch_events.setattr(torch.cuda, "Event", record_event) + batch_events.setattr(base_backend, "cuda_memcpy_pinned_h2d_batch_async", record_batch) + backend._copy_req_token_indexes_to_gpu(reqs) + observed = [] + for req, index_stream in zip(reqs, index_streams): + with torch.cuda.stream(index_stream): + # 不等待后台线程：已标记提交的索引必须能在请求自己的索引流立即使用。 + observed.append(req.token_index.get_gpu_token_indexes(0, 8).clone()) + if mode == "pin_mem_cache": + for req in reqs: + req.token_index.reset() + req.token_index.append_token_indexes(torch.zeros(8, dtype=torch.int32, pin_memory=True)) + # 两个请求只建立一个 batch 保活事件；pure_gpu 与 PyTorch 回退无需该事件。 + assert len(events) == (1 if mode == "pin_mem_cache" and batch_available else 0) + assert batch_calls == ([2] if mode == "pin_mem_cache" else []) + pool = cuda_wrapper.PinnedH2DBatchCopier._release_pool + if pool is not None: + pool.submit(lambda: None).result(timeout=10) + assert [req.token_index.gpu_index_submitted_len for req in reqs] == [0 if mode == "pin_mem_cache" else 8] * 2 + assert not hasattr(backend, "_pending_cpu_index_uploads") + assert not hasattr(backend, "_index_upload_stream") + assert all(not hasattr(req.token_index, "_gpu_write_event") for req in reqs) + assert torch.equal(table[:, :8].cpu(), expected) + for actual, expected_row, index_stream in zip(observed, expected, index_streams): + with torch.cuda.stream(index_stream): + assert torch.equal(actual.cpu(), expected_row) + backend._copy_req_token_indexes_to_gpu([]) + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +def test_backend_waits_for_submitted_indexes_without_pending_cpu_uploads(monkeypatch, mode): + from lightllm.server.router.model_infer.mode_backend import base_backend + + table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") + index_stream = torch.cuda.default_stream() + reader_stream = torch.cuda.Stream() + reader_stream.wait_stream(index_stream) + req = make_req(table, mode) + req.token_index = create_req_token_index(req, index_stream=index_stream) + with torch.cuda.stream(index_stream): + torch.cuda._sleep(200_000_000) + req.token_index.append_token_indexes(torch.arange(8, dtype=torch.int32, pin_memory=True)) + req.token_index.get_gpu_token_indexes(0, 8) + assert req.token_index.gpu_index_submitted_len == req.token_index.token_index_len == 8 + assert not index_stream.query(), "索引应已提交但尚未完成，且没有新的 CPU 上传区间" + + def unexpected_upload(*args): + pytest.fail("没有待上传区间时不应再次提交拷贝") + + monkeypatch.setattr(base_backend, "cuda_memcpy_pinned_h2d_batch_async", unexpected_upload) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + with torch.cuda.stream(reader_stream): + backend._copy_req_token_indexes_to_gpu([req]) + # 直接读取请求表，避免 GPU getter 自带的等待掩盖 batch 入口遗漏的依赖。 + observed = table.clone() + reader_stream.synchronize() + assert observed.cpu().tolist() == [list(range(8))] + + +@requires_cuda +@pytest.mark.parametrize("mode", MODES) +def test_dp_publication_waits_for_the_selected_index_mode(monkeypatch, mode): + from lightllm.server.router.model_infer.mode_backend import base_backend + from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans + + table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") + torch.cuda._sleep(30_000_000) + table.fill_(-2) + req = make_req(table, mode) + req.token_index = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + req.token_index.append_token_indexes(torch.arange(8, dtype=torch.int32, pin_memory=True)) + req.cur_kv_len = req.hold_kv_len = 8 + context = SimpleNamespace(req_manager=req.req_manager) + monkeypatch.setattr(base_backend, "g_infer_context", context) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + backend.node_nccl_group = object() + backend.is_master_in_dp = False + + def barrier(group): + torch.cuda.current_stream().synchronize() + assert table.cpu().tolist() == [list(range(8))] + + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", barrier) + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = backend + module.fill_reqs_info([req]) 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 index 2f58371e78..fe8e3547a0 100644 --- 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 @@ -1,11 +1,14 @@ from types import SimpleNamespace import torch +import pytest # 沿用生产代码的模块加载顺序，避免独立运行本测试文件时形成循环导入。 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 +from lightllm.server.router.model_infer.mode_backend.diverse_backend import impl as diverse_impl +from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes class _FakeRadixCache: @@ -36,8 +39,8 @@ def test_diverse_copy_skips_slaves_waiting_for_overlap_cleanup(): def test_copy_master_kv_updates_slave_hold_kv_len(): - req_to_token_indexs_cpu = torch.full((2, 8), -1, dtype=torch.int32) - req_to_token_indexs_cpu[0, :4] = torch.tensor([10, 11, 12, 13], dtype=torch.int32) + token_indexes_cpu = torch.full((2, 8), -1, dtype=torch.int32) + token_indexes_cpu[0, :4] = torch.tensor([10, 11, 12, 13], dtype=torch.int32) old_shared_node = object() new_shared_node = object() master_req = SimpleNamespace( @@ -50,7 +53,6 @@ def test_copy_master_kv_updates_slave_hold_kv_len(): req_idx=1, cur_kv_len=2, hold_kv_len=2, - move_to_gpu_index_len=2, cur_output_len=0, shared_kv_node=old_shared_node, related_master_req=master_req, @@ -59,20 +61,99 @@ def test_copy_master_kv_updates_slave_hold_kv_len(): ) backend = DiversehBackend.__new__(DiversehBackend) backend.model = SimpleNamespace( - req_manager=SimpleNamespace( - req_to_token_indexs_cpu=req_to_token_indexs_cpu, - ), + req_manager=make_cpu_index_manager(token_table=token_indexes_cpu, page_size=1), ) backend.radix_cache = _FakeRadixCache() backend.is_master_in_dp = False + bind_cpu_indexes(master_req, backend.model.req_manager) + bind_cpu_indexes(slave_req, backend.model.req_manager, gpu_index_submitted_len=2) backend._copy_master_req_to_slave_req(slave_req) - assert req_to_token_indexs_cpu[1, :4].tolist() == [10, 11, 12, 13] + assert slave_req.token_index.get_cpu_token_indexes(0, 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 slave_req.move_to_gpu_index_len == 0 + assert slave_req.token_index.token_index_len == 4 + assert slave_req.token_index.gpu_index_submitted_len == 0 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] + + +def test_master_shared_prefix_rebuild_preserves_reserved_token_indexes(): + token_indexes_cpu = torch.tensor([[10, 11, 12, 13, 20, 21, 22, 23]], dtype=torch.int32) + shared_indexes = torch.tensor([100, 101, 102, 103], dtype=torch.int32) + new_shared_node = SimpleNamespace(node_prefix_total_len=4) + master_req = SimpleNamespace( + req_idx=0, + cur_kv_len=4, + hold_kv_len=8, + shared_kv_node=None, + get_input_token_ids=lambda: [1, 2, 3, 4], + args=SimpleNamespace(page_size=4), + ) + freed_indexes = [] + + class RadixCache(_FakeRadixCache): + def insert(self, key, value): + assert key.tolist() == [1, 2, 3, 4] + assert value.tolist() == [10, 11, 12, 13] + return 4, new_shared_node + + def match_prefix(self, key, update_refs): + assert not update_refs + return new_shared_node, 4, shared_indexes + + manager = make_cpu_index_manager(token_table=token_indexes_cpu, page_size=4) + bind_cpu_indexes(master_req, manager, gpu_index_submitted_len=8) + original_buffer = master_req.token_index._cpu_buffer + backend = DiversehBackend.__new__(DiversehBackend) + backend.model = SimpleNamespace(mem_manager=SimpleNamespace(free=lambda indexes: freed_indexes.append(indexes))) + backend.radix_cache = RadixCache() + + backend._master_req_to_radix_cache(master_req) + + assert master_req.token_index.get_cpu_token_indexes(0, 8).tolist() == [100, 101, 102, 103, 20, 21, 22, 23] + assert master_req.token_index._cpu_buffer is not original_buffer + assert master_req.token_index.token_index_len == master_req.hold_kv_len == 8 + assert master_req.token_index.gpu_index_submitted_len == 0 + assert InferReq._kv_cache_alloc_need(master_req, target_kv_len=8) == 0 + assert InferReq._kv_cache_alloc_need(master_req, target_kv_len=9) == 4 + assert [indexes.tolist() for indexes in freed_indexes] == [[10, 11, 12, 13]] + assert master_req.shared_kv_node is new_shared_node + + +@pytest.mark.parametrize("is_chuncked_mode", [False, True]) +def test_diverse_waits_for_model_before_rewriting_shared_prefix(monkeypatch, is_chuncked_mode): + order = [] + overlap_stream = object() + + class DefaultStream: + def wait_stream(self, source): + assert source is overlap_stream + order.append("wait_model") + + def synchronize(self): + order.append("synchronize") + + default_stream = DefaultStream() + monkeypatch.setattr(torch.cuda, "default_stream", lambda: default_stream) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: default_stream) + monkeypatch.setattr(diverse_impl.g_infer_context, "get_overlap_stream", lambda: overlap_stream) + master = SimpleNamespace( + is_master_req=lambda: True, + get_chuncked_input_token_len=lambda: 4, + get_cur_total_len=lambda: 4, + cur_output_len=0, + slave_reqs=[object()], + ) + backend = DiversehBackend.__new__(DiversehBackend) + backend.is_master_in_dp = False + backend._master_req_to_radix_cache = lambda master_req: order.append("rewrite_indexes") + + packs = backend._diverse_pre_post_handle([master], is_chuncked_mode) + + assert order == ["wait_model", "rewrite_indexes", "synchronize"] + assert packs[0].req_obj is master + assert packs[0].output_len == 1 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 2fe5cdd73e..b3d504956f 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 @@ -11,7 +11,9 @@ ) from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as multi_level_kv_cache_impl -from lightllm.server.router.model_infer.infer_batch import InferenceContext +from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext +from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes +from lightllm.server.router.model_infer.infer_req_ext1 import create_req_token_index def test_cache_tiers_reassignment_is_rejected(): @@ -42,10 +44,11 @@ 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_cpu=torch.tensor([[10, 11, 12, 13, 14, 15, 16]])) + context.req_manager = make_cpu_index_manager(token_table=torch.tensor([[10, 11, 12, 13, 14, 15, 16]]), page_size=1) 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) + bind_cpu_indexes(req, context.req_manager) free_token_indexes = [] context._free_req_mem_without_radix_insert(free_token_indexes, req) @@ -59,14 +62,18 @@ def test_legacy_cache_tiers_still_insert_gpu_radix_cache(): context = InferenceContext() context.radix_cache = object() context.is_hybrid_att_model = False + context.req_manager = make_cpu_index_manager(token_table=torch.full((1, 4), -1, dtype=torch.int32)) inserted_reqs = [] context._full_att_free_req = lambda free_token_index, req: inserted_reqs.append(req) req = SimpleNamespace( + req_idx=0, cache_tiers=(CacheTier.GPU, CacheTier.CPU, CacheTier.DISK), cur_kv_len=3, shm_req=SimpleNamespace(shm_cur_kv_len=3), ) + bind_cpu_indexes(req, context.req_manager) + context.free_a_req_mem([], req) assert inserted_reqs == [req] @@ -136,7 +143,7 @@ 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_cpu=torch.tensor([[10, 11, 12]])) + context.req_manager = make_cpu_index_manager(token_table=torch.tensor([[10, 11, 12]]), page_size=1) 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,6 +156,7 @@ 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}, ) + bind_cpu_indexes(req, context.req_manager) free_token_indexes = [] context._free_req_mem_without_radix_insert(free_token_indexes, req) @@ -160,7 +168,9 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): assert req.hybrid_len_to_big_page_id == {} -def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch): +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires pinned CPU buffers") +def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch, mode): table = torch.full((1, 136), -1, dtype=torch.int32) gpu_table = torch.full_like(table, -2) table[0, :4] = torch.arange(4, dtype=torch.int32) @@ -171,13 +181,14 @@ def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch): def alloc(need_size): nonlocal next_index alloc_sizes.append(need_size) - mem_indexes = torch.arange(next_index, next_index + need_size, dtype=torch.int32) + mem_indexes = torch.arange(next_index, next_index + need_size, dtype=torch.int32, pin_memory=True) next_index += need_size return mem_indexes def alloc_req_kv_mem(req, alloc_token_num): mem_indexes = alloc(alloc_token_num) new_hold_kv_len = req.hold_kv_len + alloc_token_num + req.token_index.append_token_indexes(mem_indexes) table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes req.hold_kv_len = new_hold_kv_len return mem_indexes @@ -192,7 +203,7 @@ def load(mem_indexes, **kwargs): radix_cache=None, model=SimpleNamespace( mem_manager=SimpleNamespace(operator=operator), - req_manager=SimpleNamespace(req_to_token_indexs_cpu=table, req_to_token_indexs=gpu_table), + req_manager=make_cpu_index_manager(token_table=table, req_to_token_indexs=gpu_table, page_size=4), ), _alloc_req_kv_mem=alloc_req_kv_mem, ) @@ -203,7 +214,7 @@ def load(mem_indexes, **kwargs): module.need_sync_compute_stream = lambda: False module.cpu_cache_client = SimpleNamespace() context = SimpleNamespace( - req_manager=SimpleNamespace(req_to_token_indexs_cpu=table, mem_manager=SimpleNamespace(alloc=alloc)), + req_manager=make_cpu_index_manager(token_table=table, mem_manager=SimpleNamespace(alloc=alloc), page_size=1), get_can_alloc_token_num=lambda: 200, ) req = SimpleNamespace( @@ -218,6 +229,12 @@ def load(mem_indexes, **kwargs): ), sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), ) + bind_cpu_indexes(req, backend.model.req_manager) + req.args.req_token_index_mode = mode + req.token_index = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + req.token_index.append_token_indexes(table[0, :4].pin_memory()) + backend.model.req_manager._req_token_indexes[0] = req.token_index + context.req_manager = backend.model.req_manager monkeypatch.setattr(multi_level_kv_cache_impl, "g_infer_context", context) monkeypatch.setattr(multi_level_kv_cache_impl.dist, "barrier", lambda group: None) uploaded_indexes = [] @@ -227,10 +244,10 @@ def upload(tensor, **kwargs): return tensor def unexpected_copy(*args, **kwargs): - pytest.fail("CPU mirror upload must not re-pin indexes or concatenate GPU tensors") + pytest.fail("CPU page-index upload must not concatenate GPU tensors") monkeypatch.setattr(torch.Tensor, "cuda", upload) - monkeypatch.setattr(torch.Tensor, "pin_memory", unexpected_copy) + monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: tensor) monkeypatch.setattr(torch, "cat", unexpected_copy) module.load_cpu_cache_to_reqs([req]) @@ -243,3 +260,82 @@ def unexpected_copy(*args, **kwargs): assert table[0, :132].tolist() == list(range(132)) assert gpu_table[0, :132].tolist() == list(range(132)) assert gpu_table[0, 132:].eq(-2).all() + + +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.parametrize("use_default_index_stream", [False, True]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cpu_cache_offload_getter_orders_transfer_stream(monkeypatch, mode, use_default_index_stream): + default_stream = torch.cuda.default_stream() + index_stream = default_stream if use_default_index_stream else torch.cuda.Stream() + index_stream.wait_stream(default_stream) + transfer_stream = torch.cuda.Stream() + gpu_indexes = torch.full((64,), -1, dtype=torch.int32, device="cuda") + req = SimpleNamespace( + req_idx=0, + req_manager=SimpleNamespace(req_to_token_indexs=gpu_indexes.unsqueeze(0)), + args=SimpleNamespace(req_token_index_mode=mode), + cache_tiers=(CacheTier.CPU,), + cur_kv_len=64, + shm_req=SimpleNamespace( + input_len=64, + token_hash_list=SimpleNamespace(get_all=lambda: [123]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [64]), + ), + ) + helper = create_req_token_index(req, index_stream=index_stream) + req.token_index = helper + expected = torch.arange(64, dtype=torch.int32, pin_memory=True) + # 延迟索引流上的写入，验证 getter 在卸载流调用时也会承接依赖。 + with torch.cuda.stream(index_stream): + torch.cuda._sleep(20_000_000) + helper.append_token_indexes(expected) + order = [] + observed_indexes = [] + original_get_gpu_token_indexes = helper.get_gpu_token_indexes + + def prepare_indexes(start, end): + assert torch.cuda.current_stream() == transfer_stream + order.append("prepare") + return original_get_gpu_token_indexes(start, end) + + original_wait_stream = torch.cuda.Stream.wait_stream + + def wait_stream(waiting_stream, source_stream): + assert waiting_stream == transfer_stream and source_stream == index_stream + order.append("transfer_waits_index") + original_wait_stream(waiting_stream, source_stream) + + def offload(mem_indexes, **kwargs): + assert torch.cuda.current_stream() == transfer_stream + order.append("offload") + observed_indexes.append(mem_indexes.clone()) + + def broadcast(data, **kwargs): + data[0] = {"item_size": 1, "page_list": [3], "ready_list": [False]} + + req.token_index.get_gpu_token_indexes = prepare_indexes + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.backend = SimpleNamespace( + is_master_in_dp=False, + dp_world_size=1, + model=SimpleNamespace( + mem_manager=SimpleNamespace(operator=SimpleNamespace(offload_gpu_kv_to_cpu_cache=offload)) + ), + ) + module.gloo_group = object() + module.cpu_cache_client = object() + module.page_index_buffer = torch.empty(8, dtype=torch.int32, device="cuda") + module.page_ready_buffer = torch.empty(8, dtype=torch.bool, device="cuda") + monkeypatch.setattr(torch.cuda.Stream, "wait_stream", wait_stream) + monkeypatch.setattr(multi_level_kv_cache_impl.dist, "broadcast_object_list", broadcast) + + task = module._start_kv_cache_offload_task(req, transfer_stream) + task.sync_event.synchronize() + + assert order == ["prepare", "transfer_waits_index", "offload"] + assert torch.cuda.current_stream() == default_stream + assert torch.equal(observed_indexes[0].cpu(), expected) + assert task.req_obj is req + assert task.move_token_num == 64 + assert req.cpu_cache_task_status == InferReq._CpuCacheTaskStatus.RUNNING 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 index 6af1b1b854..1719180a8e 100644 --- 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 @@ -12,6 +12,7 @@ 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 +from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes def _bind_alloc_need(req, page_size): @@ -41,7 +42,7 @@ def test_pd_decode_reserves_model_pages_but_transfers_only_logical_kv(monkeypatc 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_cpu=table)) + backend.model = SimpleNamespace(req_manager=make_cpu_index_manager(token_table=table, page_size=1)) backend.is_master_in_dp = False backend._alloc_req_kv_mem = lambda req, size: _reserve_into_table(table, req, alloc_sizes, next_index, size) @@ -109,28 +110,25 @@ def test_dp_cache_fetch_publishes_gpu_indexes_before_barrier_returns(monkeypatch torch.arange(32, out=cpu_table.view(-1)) gpu_table = torch.full((2, 16), -1, dtype=torch.int32, device="cuda") gpu_table[1, :2].copy_(cpu_table[1, :2], non_blocking=True) - req_manager = SimpleNamespace(req_to_token_indexs_cpu=cpu_table, req_to_token_indexs=gpu_table) + req_manager = make_cpu_index_manager(token_table=cpu_table, req_to_token_indexs=gpu_table, page_size=4) reqs = [ - SimpleNamespace(req_idx=1, cur_kv_len=4, hold_kv_len=4, move_to_gpu_index_len=2), - SimpleNamespace(req_idx=0, cur_kv_len=4, hold_kv_len=8, move_to_gpu_index_len=0), + SimpleNamespace(req_idx=1, cur_kv_len=4, hold_kv_len=4), + SimpleNamespace(req_idx=0, cur_kv_len=4, hold_kv_len=8), ] - stream = torch.cuda.Stream() - overlap_stream = torch.cuda.Stream() - current_stream = torch.cuda.current_stream() - # 测试刚创建的 GPU 表先完成初始化；生产表在模型初始化期间已经就绪。 - overlap_stream.wait_stream(current_stream) - # 模拟 overlap 流上尚未完成的旧索引写入，验证 DP 上传会排在其后。 - with torch.cuda.stream(overlap_stream): - torch.cuda._sleep(20_000_000) - gpu_table[0, :8].fill_(-2) - context = SimpleNamespace(req_manager=req_manager, get_overlap_stream=lambda: overlap_stream) + for req, gpu_index_submitted_len in zip(reqs, (2, 0)): + bind_cpu_indexes(req, req_manager, gpu_index_submitted_len=gpu_index_submitted_len) + default_stream = torch.cuda.default_stream() + # 索引旧写入统一在默认流；上传流必须等待它，随后默认流再等待上传完成。 + torch.cuda._sleep(20_000_000) + gpu_table[0, :8].fill_(-2) + context = SimpleNamespace(req_manager=req_manager) monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) monkeypatch.setattr(base_backend, "g_infer_context", context) backend = SimpleNamespace(node_nccl_group=object(), is_master_in_dp=is_master_in_dp) order = [] def upload(run_reqs): - assert torch.cuda.current_stream() == stream + assert torch.cuda.current_stream() == default_stream order.append("upload") base_backend.ModeBackend._copy_req_token_indexes_to_gpu(backend, run_reqs) @@ -141,19 +139,19 @@ def forbid_synchronize(*args): pytest.fail("DP shared must rely on NCCL barrier rather than explicitly synchronizing a stream") def wait_stream(waiting_stream, source_stream): - if waiting_stream == stream: - assert source_stream == overlap_stream - order.append("copy_waits_overlap") + if waiting_stream != default_stream: + assert source_stream == default_stream + order.append("upload_waits_default") else: - assert waiting_stream == current_stream and source_stream == stream - order.append("current_waits_copy") + assert source_stream != default_stream + order.append("default_waits_upload") original_wait_stream(waiting_stream, source_stream) def barrier(group): assert group is backend.node_nccl_group - assert order == ["copy_waits_overlap", "upload", "current_waits_copy"] - # 模拟同步 NCCL barrier 在 C++ 层等待当前流完成。 - original_synchronize(current_stream) + assert order == ["upload", "upload_waits_default", "default_waits_upload"] + # 模拟同步 NCCL barrier 在 C++ 层等待默认流完成。 + original_synchronize(default_stream) assert torch.equal(gpu_table[1, :4].cpu(), cpu_table[1, :4]) assert torch.equal(gpu_table[0, :8].cpu(), cpu_table[0, :8]) assert gpu_table[0, 8:].eq(-1).all() @@ -166,14 +164,13 @@ def barrier(group): monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", barrier) module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) module.backend = backend - module.copy_stream = stream module.dp_rank_in_node = 0 module.shared_req_infos = SimpleNamespace(arr=np.full((2, 2, 2), -1, dtype=np.int64)) module.fill_reqs_info(reqs) - assert order == ["copy_waits_overlap", "upload", "current_waits_copy", "barrier"] - assert [req.move_to_gpu_index_len for req in reqs] == [4, 8] + assert order == ["upload", "upload_waits_default", "default_waits_upload", "barrier"] + assert [req.token_index.gpu_index_submitted_len for req in reqs] == [4, 8] if is_master_in_dp: assert module.shared_req_infos.arr[:, 0].tolist() == [[4, 1], [4, 0]] else: @@ -209,7 +206,7 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat 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_cpu=table), + req_manager=make_cpu_index_manager(token_table=table, page_size=1), get_can_alloc_token_num=lambda: 100, ) monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py index b22f9a7403..4c4aa6c46a 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py @@ -10,6 +10,7 @@ from lightllm.server.router.model_infer import infer_batch from lightllm.server.router.model_infer.infer_batch import InferReq, InferSamplingParams from lightllm.server.router.model_infer.mode_backend import base_backend +from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes def _make_req_with_cached_prefix(monkeypatch, prompt_logprobs, disable_prompt_cache=False, hybrid=False): @@ -22,7 +23,7 @@ def _make_req_with_cached_prefix(monkeypatch, prompt_logprobs, disable_prompt_ca req.args = SimpleNamespace( page_size=2, linear_att_hash_page_size=2, linear_att_page_block_num=2, max_req_total_len=16 ) - req.cur_kv_len = req.hold_kv_len = req.move_to_gpu_index_len = 0 + req.cur_kv_len = req.hold_kv_len = 0 req.shared_kv_node = req.tail_small_page_buffer_id = None req.shm_req = SimpleNamespace( sample_params=sample_params, @@ -35,11 +36,13 @@ def _make_req_with_cached_prefix(monkeypatch, prompt_logprobs, disable_prompt_ca req.get_cur_total_len = lambda: 6 node = SimpleNamespace(node_prefix_total_len=4, big_page_buffer_idx=0, is_big_page_node=lambda: True) match_prefix = Mock(return_value=(node, 4, torch.arange(4, dtype=torch.int32))) - manager = SimpleNamespace( - req_to_token_indexs_cpu=torch.zeros((1, 8), dtype=torch.int32), + manager = make_cpu_index_manager( + token_table=torch.zeros((1, 8), dtype=torch.int32), restore_big_page_state=Mock(), init_hybrid_attention_state=Mock(), + page_size=2, ) + bind_cpu_indexes(req, manager) monkeypatch.setattr( infer_batch, "g_infer_context", @@ -100,3 +103,27 @@ def test_cached_prompt_still_captures_logprobs_for_every_prompt_position(monkeyp captured = capture_manager.capture.call_args.kwargs assert captured["mem_indexes"].tolist() == list(range(5)) assert captured["top_token_ids"].shape == captured["top_logprobs"].shape == (5, prompt_logprobs) + + +@pytest.mark.parametrize("kind", ["prompt_logprobs", "routed_experts"]) +def test_final_metadata_extracts_owned_token_indexes_through_req_helper(monkeypatch, kind): + from lightllm.server.router.model_infer import infer_req_ext + + req, _ = _make_req_with_cached_prefix(monkeypatch, prompt_logprobs=3) + req.cur_kv_len = 6 + req.cur_output_len = 0 + req.token_index.append_token_indexes(torch.tensor([8, 9, 2, 3, 14, 15], dtype=torch.int32)) + manager = SimpleNamespace(extract=Mock(return_value="metadata")) + metadata = infer_req_ext.FinalTokenMetadataExt(req) + if kind == "prompt_logprobs": + monkeypatch.setattr(infer_req_ext.PromptLogprobsCaptureManager, "get_instance", lambda: manager) + result = metadata.collect_prompt_logprobs() + else: + monkeypatch.setattr(infer_req_ext.MoeRouteInfoManager, "get_instance", lambda: manager) + result = metadata.collect_routed_experts() + assert result == "metadata" + indexes = manager.extract.call_args.args[0] + assert indexes.tolist() == [8, 9, 2, 3, 14] + req.token_index.reset() + req.token_index.append_token_indexes(torch.arange(6, dtype=torch.int32)) + assert indexes.tolist() == [8, 9, 2, 3, 14] diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py index 99e99fb0f1..eba292cda4 100644 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ b/unit_tests/test_static_benchmark_paged_kv.py @@ -3,6 +3,7 @@ from types import SimpleNamespace import numpy as np +import pytest import torch @@ -20,6 +21,7 @@ decode_profile_batch_divisor, resolve_batch_max_prefill_cases, ) +from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes class FakeMemManager: @@ -31,21 +33,23 @@ def __init__(self, page_size: int): 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) + indexes = torch.arange(self.next_index, self.next_index + size, dtype=torch.int32, pin_memory=True) self.next_index += size return indexes -def make_executor(monkeypatch, page_size: int = 4): +def make_executor(monkeypatch, page_size: int = 4, mode="pin_mem_cache"): mem_manager = FakeMemManager(page_size=page_size) - req_manager = SimpleNamespace( + req_manager = make_cpu_index_manager( mem_manager=mem_manager, - req_to_token_indexs_cpu=torch.full((4, 32), -1, dtype=torch.int32), + token_table=torch.full((4, 32), -1, dtype=torch.int32), req_to_token_indexs=torch.full((4, 32), -1, dtype=torch.int32), + page_size=page_size, ) model = SimpleNamespace(mem_manager=mem_manager, req_manager=req_manager) - args = SimpleNamespace(mtp_mode=None, mtp_step=0) + args = SimpleNamespace(mtp_mode=None, mtp_step=0, req_token_index_mode=mode) executor = StaticBenchmarkExecutor(args=args, model=model, draft_models=[], token_source=None) + req_manager._req_token_indexes = executor._req_token_indexes return executor, mem_manager, req_manager @@ -55,8 +59,10 @@ def assert_complete_physical_pages(indexes: torch.Tensor, page_size: int): assert torch.all(page // page_size == page[0] // page_size) -def test_prefill_allocates_and_reuses_complete_pages_per_request(monkeypatch): - executor, mem_manager, req_manager = make_executor(monkeypatch, page_size=4) +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires pinned CPU buffers") +def test_prefill_allocates_and_reuses_complete_pages_per_request(monkeypatch, mode): + executor, mem_manager, req_manager = make_executor(monkeypatch, page_size=4, mode=mode) req_idx = torch.tensor([0, 1], dtype=torch.int32) model_input = executor._make_prefill_input( @@ -66,11 +72,16 @@ def test_prefill_allocates_and_reuses_complete_pages_per_request(monkeypatch): ) assert mem_manager.alloc_sizes == [4, 4] - assert torch.equal(req_manager.req_to_token_indexs, req_manager.req_to_token_indexs_cpu) + if mode == "pin_mem_cache": + assert executor._req_token_indexes[0]._cpu_buffer.is_pinned() + else: + assert not hasattr(executor._req_token_indexes[0], "_cpu_buffer") + for row in (0, 1): + assert torch.equal(req_manager.req_to_token_indexs[row, :4], req_manager.get_cpu_token_indexes(row, 0, 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_cpu[0, :4], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs_cpu[1, :4], page_size=4) + assert_complete_physical_pages(req_manager.get_cpu_token_indexes(0, 0, 4), page_size=4) + assert_complete_physical_pages(req_manager.get_cpu_token_indexes(1, 0, 4), page_size=4) # 两个请求都还在已有页内继续 decode，不应产生新的分配。 executor._make_decode_input( @@ -94,12 +105,14 @@ def test_prefill_allocates_and_reuses_complete_pages_per_request(monkeypatch): ) 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_cpu[0, :8], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs_cpu[1, :8], page_size=4) + assert_complete_physical_pages(req_manager.get_cpu_token_indexes(0, 0, 8), page_size=4) + assert_complete_physical_pages(req_manager.get_cpu_token_indexes(1, 0, 8), page_size=4) -def test_mtp_expanded_rows_allocate_once_per_request(monkeypatch): - executor, mem_manager, req_manager = make_executor(monkeypatch, page_size=4) +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires pinned CPU buffers") +def test_mtp_expanded_rows_allocate_once_per_request(monkeypatch, mode): + executor, mem_manager, req_manager = make_executor(monkeypatch, page_size=4, mode=mode) executor.args.mtp_mode = "eagle_with_att" executor.args.mtp_step = 3 @@ -116,8 +129,8 @@ def test_mtp_expanded_rows_allocate_once_per_request(monkeypatch): 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_cpu[0, :8], page_size=4) - assert_complete_physical_pages(req_manager.req_to_token_indexs_cpu[1, :8], page_size=4) + assert_complete_physical_pages(req_manager.get_cpu_token_indexes(0, 0, 8), page_size=4) + assert_complete_physical_pages(req_manager.get_cpu_token_indexes(1, 0, 8), page_size=4) def test_profile_capacity_uses_page_aligned_per_request_cost(): From 0af0bca94ce3cf30f3c120edd337e7de0c26aae8 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 03:05:56 +0000 Subject: [PATCH 11/22] refactor: remove redundant request index mode state Keep request index mode selection in the helper factory and remove duplicate ReqManager state, validation and logging. Use English diagnostic messages across the affected runtime, benchmark and test files while preserving Chinese comments. Update tests and fixtures accordingly. Validation: 225 GPU tests passed; 72 CPU tests passed and 153 CUDA-dependent tests skipped. Black, Flake8 and diff checks passed. --- lightllm/common/req_manager/base.py | 10 +----- .../router/model_infer/infer_req_ext1.py | 10 +++--- .../model_infer/mode_backend/base_backend.py | 7 +++-- .../benchmark_event_synchronize_gil.py | 14 ++++++--- test/benchmark/benchmark_pinned_empty_h2d.py | 6 ++-- unit_tests/common/cpu_index_helpers.py | 3 +- unit_tests/common/test_kv_cache_allocator.py | 2 +- .../common/test_pinned_memory_stream.py | 18 +++++------ unit_tests/common/test_req_manager_cpu.py | 2 +- unit_tests/common/test_req_token_index.py | 31 ++++++++----------- 10 files changed, 48 insertions(+), 55 deletions(-) diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index f334c24a00..49ace35c07 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -49,20 +49,12 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana # 的那个batch size 进行运行，所有 padding 的请求都会使用预留的这个请求管理 id 进行处理 # 这样让 DP 的实现更为简化一些。 self.req_list = _ReqLinkedList(max_request_num) - args = get_env_start_args() - page_size = args.page_size + page_size = get_env_start_args().page_size self.page_size = page_size aligned_sequence_length = (max_sequence_length + page_size - 1) // page_size * page_size self.req_to_token_indexs = torch.zeros( (max_request_num + 1, aligned_sequence_length), dtype=torch.int32, device="cuda" ) - self.req_token_index_mode = args.req_token_index_mode - if self.req_token_index_mode == "pin_mem_cache": - logger.info("请求 token 索引模式 pin_mem_cache：按请求申请 pinned CPU 缓存，索引更新批量上传 GPU") - elif self.req_token_index_mode == "pure_gpu": - logger.info("请求 token 索引模式 pure_gpu：直接维护 GPU 表，CPU 读取索引时执行 D2H 同步") - else: - raise ValueError(f"未知请求 token 索引模式：{self.req_token_index_mode}") self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num) self.max_request_num = max_request_num self.HOLD_REQUEST_ID = max_request_num diff --git a/lightllm/server/router/model_infer/infer_req_ext1.py b/lightllm/server/router/model_infer/infer_req_ext1.py index a6f1cb1508..0dfd4ce061 100644 --- a/lightllm/server/router/model_infer/infer_req_ext1.py +++ b/lightllm/server/router/model_infer/infer_req_ext1.py @@ -89,7 +89,7 @@ def append_token_indexes(self, mem_indexes: torch.Tensor): append_len = mem_indexes.numel() if append_len == 0: return - assert mem_indexes.device.type == "cpu", "mem_indexes 必须是 CPU Tensor" + assert mem_indexes.device.type == "cpu", "mem_indexes must be a CPU Tensor" start = self.token_index_len end = start + append_len self._reserve(end) @@ -103,7 +103,7 @@ def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: """在索引流提交全部待上传索引，让调用流承接依赖后返回 GPU 视图。""" - assert end <= self.token_index_len, "读取范围超出已记录的 token 索引" + assert end <= self.token_index_len, "Read range exceeds recorded token indexes" current_stream = torch.cuda.current_stream(self._index_stream.device) with torch.cuda.stream(self._index_stream): if self.gpu_index_submitted_len < end: @@ -162,7 +162,7 @@ def append_token_indexes(self, mem_indexes: torch.Tensor): append_len = mem_indexes.numel() if append_len == 0: return - assert mem_indexes.device.type == "cpu", "mem_indexes 必须是 CPU Tensor" + assert mem_indexes.device.type == "cpu", "mem_indexes must be a CPU Tensor" # 普通 CPU 输入在这里转为独立 pinned storage；已有 pinned 输入直接复用。 if not mem_indexes.is_pinned(): mem_indexes = mem_indexes.pin_memory() @@ -185,7 +185,7 @@ def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: # append 已在索引流提交写入，读取只校验范围，不推进提交进度。 - assert end <= self.gpu_index_submitted_len, "读取范围超出已提交的 GPU 索引" + assert end <= self.gpu_index_submitted_len, "Read range exceeds submitted GPU indexes" current_stream = torch.cuda.current_stream(self._index_stream.device) if current_stream != self._index_stream: current_stream.wait_stream(self._index_stream) @@ -212,4 +212,4 @@ def create_req_token_index(infer_req: "InferReq", index_stream: torch.cuda.Strea return PinMemCacheReqTokenIndex(infer_req, index_stream) if mode == "pure_gpu": return PureGpuReqTokenIndex(infer_req, index_stream) - raise ValueError(f"未知请求 token 索引模式：{mode}，可选 pin_mem_cache / pure_gpu") + raise ValueError(f"Unknown request token index mode: {mode}; expected pin_mem_cache or pure_gpu") 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 ee4badd471..f48dd8a435 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -668,7 +668,7 @@ def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): if not run_reqs: return index_streams = {req.token_index._index_stream for req in run_reqs} - assert len(index_streams) == 1, "同一批请求必须使用同一个索引流" + assert len(index_streams) == 1, "All requests in a batch must use the same index stream" index_stream = index_streams.pop() current_stream = torch.cuda.current_stream(index_stream.device) # 第二阶段：切出尚未提交的 CPU 源视图和 GPU 目标视图，记录本次提交终点。 @@ -679,7 +679,10 @@ def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): start, end = token_index.gpu_index_submitted_len, token_index.token_index_len if start < end: source_indexes = token_index.get_upload_indexes() - assert source_indexes is not None, "GPU 索引提交进度与记录长度不一致，GPU 模式应在追加时完成提交" + assert source_indexes is not None, ( + "GPU index submission progress does not match recorded length; " + "pure_gpu must submit indexes when appending" + ) cpu_index_slices.append(source_indexes) gpu_index_slices.append(token_index.gpu_indexes[start:end]) pending_index_uploads.append((token_index, end)) diff --git a/test/benchmark/benchmark_event_synchronize_gil.py b/test/benchmark/benchmark_event_synchronize_gil.py index 86eee4b772..669bc7eb3a 100644 --- a/test/benchmark/benchmark_event_synchronize_gil.py +++ b/test/benchmark/benchmark_event_synchronize_gil.py @@ -150,9 +150,11 @@ def main(): args = parser.parse_args() visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") if not visible or "," in visible or visible == "3": - parser.error("请通过 CUDA_VISIBLE_DEVICES 显式选择一张空闲 GPU，不能使用 3 号卡") + parser.error("Set CUDA_VISIBLE_DEVICES to one idle GPU; GPU 3 must not be used") if args.repeats < 1 or args.wait_ms < 400 or not 0 < args.margin_ms < args.wait_ms / 4: - parser.error("重复数须大于 0，等待至少 400ms，边缘排除区间须介于 0 和等待时间的四分之一之间") + parser.error( + "Repeats must be positive, wait must be at least 400 ms, and margin must be between 0 and wait / 4" + ) import torch from lightllm.common import cuda_wrapper @@ -174,7 +176,7 @@ def main(): "event_synchronize_gil_" + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") ) root.mkdir(parents=True, exist_ok=False) - print(f"结果目录：{root}", flush=True) + print(f"Results directory: {root}", flush=True) rows = [] try: for repeat in range(args.repeats): @@ -184,7 +186,7 @@ def main(): torch, cuda_wrapper, stream, host_waits, cycles, method, repeat, args.wait_ms, args.margin_ms ) ) - print(f"已完成第 {repeat + 1}/{args.repeats} 轮", flush=True) + print(f"Completed round {repeat + 1}/{args.repeats}", flush=True) finally: stream.synchronize() cuda_wrapper.PinnedH2DBatchCopier._get_release_pool().submit(lambda: None).result(timeout=10) @@ -226,7 +228,9 @@ def main(): } (root / "result.json").write_text(json.dumps(evidence, ensure_ascii=False, indent=2) + "\n") if not evidence["passed"]: - raise SystemExit("实验判据未全部满足，请查看 result.json 中的有效等待时长与内部窗口进度") + raise SystemExit( + "Experiment criteria were not fully met; check wait durations and window progress in result.json" + ) if __name__ == "__main__": diff --git a/test/benchmark/benchmark_pinned_empty_h2d.py b/test/benchmark/benchmark_pinned_empty_h2d.py index 537c2f12e3..f0cfc9f892 100644 --- a/test/benchmark/benchmark_pinned_empty_h2d.py +++ b/test/benchmark/benchmark_pinned_empty_h2d.py @@ -172,9 +172,9 @@ def main(): # 显式限制到一张卡，防止实验误用用户禁止的 GPU 3 或其他运行中的服务。 visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") if not visible or "," in visible or visible == "3": - parser.error("请通过 CUDA_VISIBLE_DEVICES 显式选择一张空闲 GPU，不能使用 3 号卡") + parser.error("Set CUDA_VISIBLE_DEVICES to one idle GPU; GPU 3 must not be used") if args.repeats < 1 or min(args.sizes) < 1 or args.delay_cycles < 1 or args.attempts < 1: - parser.error("重复数、规模、延时和分配尝试次数必须大于 0") + parser.error("Repeats, sizes, delay cycles and allocation attempts must be positive") if args.worker: run_worker(args) return @@ -241,7 +241,7 @@ def main(): writer = csv.DictWriter(output, fieldnames=list(summaries[0])) writer.writeheader() writer.writerows(summaries) - print(f"实验结果：{root}", flush=True) + print(f"Experiment results: {root}", flush=True) if __name__ == "__main__": diff --git a/unit_tests/common/cpu_index_helpers.py b/unit_tests/common/cpu_index_helpers.py index 138c0a7e77..36ac5f0b60 100644 --- a/unit_tests/common/cpu_index_helpers.py +++ b/unit_tests/common/cpu_index_helpers.py @@ -48,7 +48,6 @@ def append_token_indexes(self, req_idx, indexes): manager = TestReqManager.__new__(TestReqManager) manager.page_size = page_size - manager.req_token_index_mode = "pin_mem_cache" manager.token_indexes_cpu = token_table manager.req_to_token_indexs = torch.full_like(token_table, -1) manager._req_token_indexes = {} @@ -60,7 +59,7 @@ def append_token_indexes(self, req_idx, indexes): def bind_cpu_indexes(req, manager, gpu_index_submitted_len=0, token_index_len=None): req.req_manager = manager if not hasattr(req, "args"): - req.args = SimpleNamespace(req_token_index_mode=manager.req_token_index_mode) + req.args = SimpleNamespace(req_token_index_mode="pin_mem_cache") req.token_index = manager.get_req_token_index(req.req_idx) # 矩阵只提供容量和预置数据；追加位置取请求实际持有的长度，不能使用整行容量。 if token_index_len is None: diff --git a/unit_tests/common/test_kv_cache_allocator.py b/unit_tests/common/test_kv_cache_allocator.py index 989338d437..a083428403 100644 --- a/unit_tests/common/test_kv_cache_allocator.py +++ b/unit_tests/common/test_kv_cache_allocator.py @@ -268,7 +268,7 @@ def test_request_free_indexes_recycle_pages_after_reset_batch_and_slot_reuse( reqs = [] def unexpected_snapshot(*args): - pytest.fail("pure_gpu 的实际释放路径不能先拷回完整 CPU token 快照") + pytest.fail("The pure_gpu release path must not copy a full token index snapshot to CPU first") with torch.cuda.stream(torch.cuda.default_stream()): for row, indexes in enumerate((first, second)): diff --git a/unit_tests/common/test_pinned_memory_stream.py b/unit_tests/common/test_pinned_memory_stream.py index aaff9db033..f87ee97b51 100644 --- a/unit_tests/common/test_pinned_memory_stream.py +++ b/unit_tests/common/test_pinned_memory_stream.py @@ -57,7 +57,7 @@ def test_empty_batch_and_fallback_do_not_create_release_pool(monkeypatch): monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(cuda_stream=0)) def unexpected_pool(*args, **kwargs): - pytest.fail("空批次和 copy_ 回退不能创建保活线程池") + pytest.fail("Empty batches and copy_ fallback must not create the source release pool") monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", unexpected_pool) cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([], []) @@ -155,10 +155,10 @@ def record_event(event, on_stream=None): return event_record(event, on_stream) def unexpected_copy(*args, **kwargs): - pytest.fail("raw batch 保活不应额外提交小片段 copy_ 或进入逐项回退") + pytest.fail("Raw batch source protection must not submit a small copy_ or enter per-tensor fallback") def unexpected_sync(*args, **kwargs): - pytest.fail("正常 raw batch 不应在提交线程同步 CUDA 流") + pytest.fail("Normal raw batches must not synchronize a CUDA stream on the submitting thread") monkeypatch.setattr(runtime, "cudaMemcpyBatchAsync", record_native) monkeypatch.setattr(torch.cuda.Event, "record", record_event) @@ -191,9 +191,9 @@ def test_background_wait_holds_tensor_references_without_blocking_submitter(monk event_synchronize = torch.cuda.Event.synchronize def controlled_wait(event): - assert threading.get_ident() != main_thread, "CUDA 完成事件只能由后台等待线程同步" + assert threading.get_ident() != main_thread, "Only the background worker may synchronize completion events" worker_wait_started.set() - assert allow_worker_finish.wait(timeout=10), "测试未允许后台等待结束" + assert allow_worker_finish.wait(timeout=10), "The test did not allow the background wait to finish" event_synchronize(event) monkeypatch.setattr(torch.cuda.Event, "synchronize", controlled_wait) @@ -250,7 +250,7 @@ def test_async_copy_keeps_freed_source_storage_out_of_reuse(monkeypatch, native_ replacement.fill_(-1) del replacement stream.synchronize() - assert pending_checks > 0, "延时传输应覆盖源释放和 pinned 内存复用检查" + assert pending_checks > 0, "Delayed transfers must cover source release and pinned memory reuse checks" assert torch.equal(target.cpu(), expected) @@ -268,10 +268,10 @@ def fail_setup(*args, **kwargs): raise RuntimeError(f"test {failure} failure") def unexpected_sync(*args, **kwargs): - pytest.fail("保活接管异常应直接传播，不应同步收尾") + pytest.fail("Source protection setup errors must propagate without synchronous cleanup") def unexpected_copy(*args, **kwargs): - pytest.fail("raw batch 已提交后的保护失败不能被 copy_ 回退掩盖") + pytest.fail("copy_ fallback must not hide source protection failure after raw batch submission") if failure == "record": monkeypatch.setattr(torch.cuda.Event, "record", fail_setup) @@ -305,7 +305,7 @@ def test_pytorch_fallback_needs_no_completion_event_or_pool_task(monkeypatch, re copy = torch.Tensor.copy_ def unexpected_event(*args, **kwargs): - pytest.fail("普通 PyTorch copy_ 已保护源 storage，不应建立额外完成事件") + pytest.fail("PyTorch copy_ already protects source storage and must not create an extra completion event") def record_copy(dst, src, *args, **kwargs): copy_sizes.append(src.numel() * src.element_size()) diff --git a/unit_tests/common/test_req_manager_cpu.py b/unit_tests/common/test_req_manager_cpu.py index 24a482f9af..c36d3f7eee 100644 --- a/unit_tests/common/test_req_manager_cpu.py +++ b/unit_tests/common/test_req_manager_cpu.py @@ -190,7 +190,7 @@ def test_prefix_match_appends_to_initialized_or_released_index_state(monkeypatch monkeypatch.setattr(infer_batch, "g_infer_context", context) def unexpected_reset(): - pytest.fail("前缀匹配入口的索引状态已为空，不应再次 reset") + pytest.fail("Prefix matching starts with empty indexes and must not reset them again") monkeypatch.setattr(req.token_index, "reset", unexpected_reset) if hybrid: diff --git a/unit_tests/common/test_req_token_index.py b/unit_tests/common/test_req_token_index.py index 16f0e1fc4d..9569a0e238 100644 --- a/unit_tests/common/test_req_token_index.py +++ b/unit_tests/common/test_req_token_index.py @@ -40,9 +40,6 @@ def make_req(token_table, mode="pin_mem_cache", req_idx=0, input_len=0): def test_factory_uses_request_mode_and_row_view_without_retaining_request(mode, expected_type): table = torch.arange(24, dtype=torch.int32).view(3, 8) req = make_req(table, mode, req_idx=1) - # 模式来自请求参数；manager 的相反配置不能影响该请求的选择。 - req.req_manager.req_token_index_mode = "pin_mem_cache" if mode == "pure_gpu" else "pure_gpu" - # 构造只保存传入的流，不执行 CUDA 操作；用占位对象验证透传即可。 index_stream = object() helper = create_req_token_index(req, index_stream=index_stream) @@ -72,7 +69,7 @@ def test_start_args_and_cli_select_mode(mode): assert StartArgs(req_token_index_mode=mode).req_token_index_mode == mode with pytest.raises(SystemExit): parser.parse_args(["--req_token_index_mode", "unknown"]) - with pytest.raises(ValueError, match="未知请求 token 索引模式"): + with pytest.raises(ValueError, match="Unknown request token index mode"): create_req_token_index(make_req(torch.empty((1, 8), dtype=torch.int32), "unknown"), index_stream=object()) @@ -110,7 +107,7 @@ def test_append_rejects_gpu_input_without_changing_index_state(mode): table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") helper = create_req_token_index(make_req(table, mode), index_stream=torch.cuda.current_stream()) - with pytest.raises(AssertionError, match="mem_indexes 必须是 CPU Tensor"): + with pytest.raises(AssertionError, match="mem_indexes must be a CPU Tensor"): helper.append_token_indexes(torch.arange(4, dtype=torch.int32, device="cuda")) assert helper.token_index_len == helper.gpu_index_submitted_len == 0 @@ -125,7 +122,7 @@ def test_gpu_getter_rejects_unrecorded_indexes_without_advancing_progress(mode): helper.append_token_indexes(torch.tensor([10, 11], dtype=torch.int32, pin_memory=True)) submitted_len = helper.gpu_index_submitted_len - with pytest.raises(AssertionError, match="读取范围超出"): + with pytest.raises(AssertionError, match="Read range exceeds"): helper.get_gpu_token_indexes(0, 4) assert helper.token_index_len == 2 @@ -145,7 +142,7 @@ def test_batch_upload_rejects_missing_gpu_submission_without_advancing_progress( monkeypatch.setattr(base_backend, "g_infer_context", SimpleNamespace(req_manager=req.req_manager)) backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) - with pytest.raises(AssertionError, match="GPU 索引提交进度与记录长度不一致"): + with pytest.raises(AssertionError, match="GPU index submission progress does not match recorded length"): backend._copy_req_token_indexes_to_gpu([req]) assert req.token_index.gpu_index_submitted_len == 0 @@ -168,11 +165,11 @@ def test_backend_rejects_requests_with_different_index_streams(monkeypatch, mode reqs.append(req) def unexpected_upload(*args): - pytest.fail("索引流不一致时应在提交上传前拒绝该批请求") + pytest.fail("A batch with mismatched index streams must be rejected before upload") monkeypatch.setattr(base_backend, "cuda_memcpy_pinned_h2d_batch_async", unexpected_upload) backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) - with pytest.raises(AssertionError, match="同一批请求必须使用同一个索引流"): + with pytest.raises(AssertionError, match="All requests in a batch must use the same index stream"): backend._copy_req_token_indexes_to_gpu(reqs) assert [req.token_index.gpu_index_submitted_len for req in reqs] == [0 if mode == "pin_mem_cache" else 4] * 2 @@ -230,7 +227,7 @@ def test_free_getter_borrows_gpu_indexes_but_metadata_getter_keeps_cpu_snapshot( snapshot = req.token_index.get_cpu_token_indexes(1, 5) def unexpected_snapshot(*args): - pytest.fail("pure_gpu 的释放接口不能调用普通 CPU 快照 getter") + pytest.fail("The pure_gpu free getter must not call the CPU snapshot getter") with monkeypatch.context() as no_snapshot: if mode == "pure_gpu": @@ -353,7 +350,7 @@ def test_pin_buffer_reset_during_upload_keeps_old_data_until_h2d_completes(use_d with torch.cuda.stream(index_stream): torch.cuda._sleep(200_000_000) helper.get_gpu_token_indexes(0, count) - assert not index_stream.query(), "延迟 H2D 应覆盖 reset 和重新追加的过程" + assert not index_stream.query(), "Delayed H2D must remain pending through reset and append" helper.reset() assert helper._cpu_buffer is None and source_ref() is None helper.append_token_indexes(replacement) @@ -430,17 +427,15 @@ def test_pure_gpu_slot_reuse_keeps_old_and_new_helper_writes_in_order(): @requires_cuda -@pytest.mark.parametrize("mode", MODES) @pytest.mark.parametrize("page_size", [1, 4, 64]) -def test_manager_owns_only_gpu_table_and_initializes_hold_page(monkeypatch, mode, page_size): +def test_manager_owns_only_gpu_table_and_initializes_hold_page(monkeypatch, page_size): from lightllm.common.req_manager import req_sampling_params - monkeypatch.setattr( - base, "get_env_start_args", lambda: SimpleNamespace(page_size=page_size, req_token_index_mode=mode) - ) + monkeypatch.setattr(base, "get_env_start_args", lambda: SimpleNamespace(page_size=page_size)) monkeypatch.setattr(req_sampling_params, "ReqSamplingParamsManager", lambda count: None) mem_manager = SimpleNamespace(page_size=page_size, HOLD_TOKEN_MEMINDEXES=tuple(range(page_size * 8, page_size * 9))) manager = base.ReqManager(2, page_size * 2 + 1, mem_manager) + assert not hasattr(manager, "req_token_index_mode") assert not hasattr(manager, "req_to_page_indexs_cpu") assert not hasattr(manager, "get_cpu_token_indexes") assert manager.req_to_token_indexs.shape == (3, 3 * page_size) @@ -549,10 +544,10 @@ def test_backend_waits_for_submitted_indexes_without_pending_cpu_uploads(monkeyp req.token_index.append_token_indexes(torch.arange(8, dtype=torch.int32, pin_memory=True)) req.token_index.get_gpu_token_indexes(0, 8) assert req.token_index.gpu_index_submitted_len == req.token_index.token_index_len == 8 - assert not index_stream.query(), "索引应已提交但尚未完成，且没有新的 CPU 上传区间" + assert not index_stream.query(), "Submitted indexes must still be pending with no new CPU upload range" def unexpected_upload(*args): - pytest.fail("没有待上传区间时不应再次提交拷贝") + pytest.fail("No copy should be submitted when there are no pending upload ranges") monkeypatch.setattr(base_backend, "cuda_memcpy_pinned_h2d_batch_async", unexpected_upload) backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) From 6566f1ffc4a5101eeccb3af6faad02d146742fe4 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 05:15:29 +0000 Subject: [PATCH 12/22] Reserve 2K output tokens in initial pinned index buffer --- docs/kv_cache_page_size.md | 4 ++-- lightllm/server/router/model_infer/infer_req_ext1.py | 4 ++-- unit_tests/common/test_req_token_index.py | 12 ++++++------ 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index 04186270d4..630ac834cf 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -52,8 +52,8 @@ Helper 保存该流，后续不随调用点的当前流切换： | `pure_gpu` | `PureGpuReqTokenIndex` 直接维护 GPU 请求表，无常驻 CPU 索引缓存 | CPU 读取需要 D2H 并等待结果，频繁的 CPU 索引操作可能更慢 | 两种模式均保留相同的 GPU 请求表，显存占用基本相同。追加后的所需长度为 `end = token_index_len + mem_indexes.numel()`。 -CPU 缓存首次按 `max(input_len, end)` -向上对齐到 1K token 申请；后续容量不足时，在本次所需长度 `end` 上多预留 2K token，再向上对齐到 1K， +CPU 缓存首次按 `max(input_len + 2K, end)` +向上对齐到 1K token 申请，为输出预留空间；后续容量不足时，在本次所需长度 `end` 上多预留 2K token，再向上对齐到 1K， 容量上限为 GPU 请求表行长度。缓冲区在首次非空写入时按需申请，重置时保留输入长度，丢弃已申请的缓冲区引用； 由新建的 `torch.empty(..., pin_memory=True, requires_grad=False)` 独立持有，不借用 KV allocator 的返回缓冲区。索引缓存按 token 保存，不再按页压缩或展开。 diff --git a/lightllm/server/router/model_infer/infer_req_ext1.py b/lightllm/server/router/model_infer/infer_req_ext1.py index 0dfd4ce061..3fa0e0fb92 100644 --- a/lightllm/server/router/model_infer/infer_req_ext1.py +++ b/lightllm/server/router/model_infer/infer_req_ext1.py @@ -140,8 +140,8 @@ def _reserve(self, end: int): if old_capacity >= end: return if old_capacity == 0: - # 首次按完整输入长度预留，同时覆盖本次写入需要的 token 区间。 - reserve_len = max(self.input_len, end) + # 首次按完整输入长度额外预留 2K 输出 token，同时覆盖本次写入区间。 + reserve_len = max(self.input_len + 2 * 1024, end) else: # 容量不足时在所需长度上额外预留 2K token，减少后续 decode 扩容。 reserve_len = end + 2 * 1024 diff --git a/unit_tests/common/test_req_token_index.py b/unit_tests/common/test_req_token_index.py index 9569a0e238..e8eee3542d 100644 --- a/unit_tests/common/test_req_token_index.py +++ b/unit_tests/common/test_req_token_index.py @@ -294,11 +294,11 @@ def test_reset_rebuild_keeps_cpu_appends_and_gpu_submission_separate(mode): @pytest.mark.parametrize( ("input_len", "max_length", "initial_capacity", "next_end", "grown_capacity"), [ - (1023, 16384, 1024, 1025, 4096), - (1024, 16384, 1024, 1025, 4096), - (1025, 16384, 2048, 2049, 5120), - (2048, 16384, 2048, 5000, 7168), - (1025, 3500, 2048, 3000, 3500), + (1023, 16384, 3072, 3073, 6144), + (1024, 16384, 3072, 3073, 6144), + (1025, 16384, 4096, 4097, 7168), + (2048, 16384, 4096, 5000, 7168), + (1025, 5500, 4096, 5000, 5500), ], ) def test_pin_buffer_reserves_input_length_and_grows_with_headroom( @@ -310,7 +310,7 @@ def test_pin_buffer_reserves_input_length_and_grows_with_headroom( assert helper._cpu_buffer is None helper.append_token_indexes(torch.empty(0, dtype=torch.int32)) assert helper._cpu_buffer is None - # 首次只追加一个索引，仍按完整输入长度预留。 + # 首次只追加一个索引，仍按完整输入长度额外预留 2K 输出 token。 first_index = torch.zeros(1, dtype=torch.int32, pin_memory=True) helper.append_token_indexes(first_index) assert helper._cpu_buffer.numel() == initial_capacity From ae95b2daf64ac237a1565ea7d948aa4483ea3bef Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 06:44:00 +0000 Subject: [PATCH 13/22] Add page-compressed pinned request token index cache --- docs/kv_cache_page_size.md | 40 ++- lightllm/server/api_cli.py | 8 +- lightllm/server/core/objs/start_args_type.py | 4 +- .../router/model_infer/infer_req_ext1.py | 103 ++++++- .../model_infer/mode_backend/base_backend.py | 4 +- .../benchmark_req_token_index_alloc.py | 204 +++++++++++++ unit_tests/common/test_kv_cache_allocator.py | 7 +- unit_tests/common/test_req_manager_cpu.py | 8 +- unit_tests/common/test_req_token_index.py | 281 ++++++++++++++---- .../mode_backend/test_multi_level_kv_cache.py | 10 +- 10 files changed, 571 insertions(+), 98 deletions(-) create mode 100644 test/benchmark/benchmark_req_token_index_alloc.py diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index 630ac834cf..1187655fd8 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -38,7 +38,8 @@ KV 存储和请求表，但资源的申请、缓存和释放以完整物理页 `ReqManager` 只持有 GPU `req_to_token_indexs`，不再分配全局 CPU 页表。 每个 `InferReq` 通过成员 `token_index` 持有一个 `ReqTokenIndex`，实现位于 `server/router/model_infer/infer_req_ext1.py`。 索引操作直接使用 `req.token_index.append_token_indexes(...)` 等接口，`InferReq` 不再提供同名转发方法。 -基类仅在构造函数中初始化公共对象和状态，其余方法只定义抽象接口；两种子类各自实现读写、上传和重置逻辑。 +基类仅在构造函数中初始化公共对象和状态，其余方法只定义抽象接口； +各子类实现或复用读写、上传和重置逻辑。 扩展构造函数接收 `infer_req` 和必填的 `index_stream`，仅提取并保存固定的 `req_idx`、GPU 请求表行视图及索引流， 避免反向持有请求对象形成循环引用；索引流由外部选择，Helper 内部不查询默认流或捕获构造时的当前流。 CPU 缓存实现同时记录输入长度 `input_len`，helper 在共享请求初始化后、前缀匹配前构造。 @@ -49,20 +50,34 @@ Helper 保存该流，后续不随调用点的当前流切换： | `--req_token_index_mode` | 存储与优点 | 代价 | | --- | --- | --- | | `pin_mem_cache`（默认） | `PinMemCacheReqTokenIndex` 按请求直接申请 pinned CPU token 缓存；CPU 读取无需 D2H，GPU 更新可批量提交 | 占用 pinned CPU 内存，需要维护上传进度 | +| `pin_mem_page_cache` | `PinMemPageCacheReqTokenIndex` 每个 KV 页仅保存一个 pinned CPU 首 token 索引，常驻索引项数约缩小为原来的 `1 / page_size` | CPU 读取和 GPU 上传时需要展开 token 索引，适用于 `page_size > 1` | | `pure_gpu` | `PureGpuReqTokenIndex` 直接维护 GPU 请求表，无常驻 CPU 索引缓存 | CPU 读取需要 D2H 并等待结果，频繁的 CPU 索引操作可能更慢 | -两种模式均保留相同的 GPU 请求表，显存占用基本相同。追加后的所需长度为 `end = token_index_len + mem_indexes.numel()`。 -CPU 缓存首次按 `max(input_len + 2K, end)` +三种模式均保留相同的 GPU 请求表，显存占用基本相同。追加后的所需长度为 `end = token_index_len + mem_indexes.numel()`。 +完整 token 模式的 CPU 缓存首次按 `max(input_len + 2K, end)` 向上对齐到 1K token 申请，为输出预留空间；后续容量不足时，在本次所需长度 `end` 上多预留 2K token，再向上对齐到 1K， -容量上限为 GPU 请求表行长度。缓冲区在首次非空写入时按需申请，重置时保留输入长度，丢弃已申请的缓冲区引用； -由新建的 `torch.empty(..., pin_memory=True, requires_grad=False)` -独立持有，不借用 KV allocator 的返回缓冲区。索引缓存按 token 保存，不再按页压缩或展开。 +容量上限为 GPU 请求表行长度。两种 CPU 缓存模式均直接使用 `torch.empty(..., pin_memory=True, requires_grad=False)` +为每个请求独立申请常驻存储；完整模式保存每个 token，页模式在首次追加时按 GPU 请求表的最大长度 +一次性申请 `ceil(max_sequence_length / page_size)` 个页首索引，后续追加无需扩容或复制旧缓存。 +页模式继承完整模式的 GPU 上传、释放快照和 reset 流程，覆盖 CPU 索引的存储和展开。 +压缩依赖已有的页内连续索引约束；追加完整页，输入从下标 0 按 page_size 取页首。 +页模式展开 CPU 快照、释放索引或待上传索引时，断言 start 和 end 均为页边界，直接返回完整展开结果。 +缓冲区在首次非空写入时按需申请，重置时保留输入长度， +丢弃已申请的缓冲区引用，由 PyTorch 在引用和异步传输依赖结束后回收。 +页模式上传时先将待提交区间展开到临时 pinned token 索引，使用 `PinMemPageCacheReqTokenIndex._alloc_pinned_tensor` +按 1M token（4 MiB）申请 int32 slab。该接口及共享状态仅属于页缓存类，保存完整 buffer 和下一个空闲位置 `_pinned_offset`， +每次只创建返回值这一个切片。 +offset 只前进，不回绕；空间不足时换用新 slab 并从零开始，请求 reset 不重置共享 offset。 +超出 slab 容量的申请独立分配，空申请不改变 slab 或 offset。旧 slab 的 storage 在所有切片引用和异步传输依赖结束后 +由 PyTorch 回收。 +纯 GPU 模式对普通 CPU 输入使用 `mem_indexes.pin_memory()` 转为独立 pinned 存储，已有 pinned 输入继续复用。 +三种模式均不借用 KV allocator 的返回缓冲区作为请求常驻缓存。 统一接口为 `append_token_indexes(mem_indexes)`、`get_cpu_token_indexes(start, end)`、 `get_free_token_indexes(start, end)` 和 `get_gpu_token_indexes(start, end)`，均使用 token 单位， -追加接口将完整输入接到已记录索引的尾部，非空输入在两种模式下均断言为 CPU Tensor；空输入直接返回。 +追加接口将完整输入接到已记录索引的尾部，非空输入在三种模式下均断言为 CPU Tensor；空输入直接返回。 CPU 缓存模式直接复制到自身的 pinned 缓冲区，无需先 pin 输入；GPU 模式仅在输入未 pinned 时内部转换，已有 pinned 输入直接复用。 -前缀匹配和索引重建路径直接传入 CPU 索引，不再自行 pin。读取接口支持页内任意区间；提交给 allocator 的释放区间仍须覆盖完整页。 +前缀匹配和索引重建路径直接传入 CPU 索引，不再自行 pin。GPU 视图读取支持页内任意区间；页模式 CPU 展开和 allocator 释放使用完整页区间。 CPU getter 返回独立快照，缓存、PD 元数据和异步任务可跨请求释放持有；GPU getter 返回当前请求的 GPU 表视图，不应跨请求槽位复用持有。Helper 保存外部传入的索引流，GPU 索引读写固定在该流， 调用方保证 token 区间有效；GPU getter 返回前让调用流等待索引流，没有新增上传时也需承接已有写入的依赖。 @@ -80,8 +95,8 @@ allocator 仍维护 CPU 空闲页表，因此最终按 `page_size` 取页首并 `token_index_len` 是当前连续记录的 token 索引长度，包含页内预留槽位；每次追加从该位置开始， 并按输入元素个数增长。CPU 缓存扩容时保留 -`[0, token_index_len)` 的已有映射。它与 CPU 缓存容量 `_cpu_buffer.numel()`、已算好 KV 的长度 -`cur_kv_len` 不同。 +`[0, token_index_len)` 的已有映射。它与 CPU 缓存容量、已算好 KV 的长度 `cur_kv_len` 不同； +完整模式的 CPU token 容量为 `_cpu_buffer.numel()`，页模式为 `_cpu_buffer.numel() * page_size`。 `gpu_index_submitted_len` 是最新 GPU 索引写入已提交的前缀右端点，后续索引流操作可按序使用 `[0, gpu_index_submitted_len)`。它表示提交进度，不表示 GPU 传输已完成，也不是累计上传量； CPU 缓存模式的追加接口只更新 CPU 索引和 `token_index_len`，不修改 GPU 提交进度。 @@ -89,12 +104,13 @@ CPU 缓存模式的追加接口只更新 CPU 索引和 `token_index_len`，不 首次前缀匹配时索引为空；暂停恢复前已由统一释放流程重置，两处匹配入口直接追加命中前缀，不重复 `reset()`。 释放或暂停请求时也由索引扩展的 `reset()` 统一清零两个长度。 CPU 模式重置时将 `_cpu_buffer` 设为 `None`，下一次追加重新申请，随后从零开始提交 GPU 更新。 -单独的 GPU getter 直接上传 CPU 缓存视图：普通追加不会改写已有区间，扩容和重置使用新存储， +单独的 GPU getter 上传完整缓存视图或页缓存展开得到的暂存区：普通追加不会改写已有区间，扩容和重置使用新存储， PyTorch `copy_(non_blocking=True)` 会登记 pinned 源的流使用，旧存储在 H2D 完成前不会被回收复用。 因此不再为单个请求上传额外申请 pinned 快照或执行一次 CPU 拷贝，也不需要在 reset 中同步 CUDA。 `get_upload_indexes()` 不接受区间参数，CPU 模式返回内部计算的 -`[gpu_index_submitted_len, token_index_len)` 视图，不改变提交进度；无待上传内容时返回空 tensor， +`[gpu_index_submitted_len, token_index_len)` 的 pinned token 索引，完整模式返回缓存视图，页模式按需展开， +不改变提交进度；无待上传内容时返回空 tensor， 纯 GPU 模式返回 `None`。CPU 缓存模式的 GPU getter 校验 `end <= token_index_len` 后按需提交完整待上传区间， 纯 GPU 模式校验 `end <= gpu_index_submitted_len`；读取均不允许推进超出已有索引的提交进度。 批量入口遇到存在未提交区间、但没有 CPU 上传源的情况会直接断言，不再替 GPU 模式补写提交进度。 diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 1332156db2..4ad655f451 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -500,13 +500,15 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--req_token_index_mode", type=str, - choices=["pin_mem_cache", "pure_gpu"], + choices=["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"], default="pin_mem_cache", help="""Storage mode for request token indexes. pin_mem_cache (default): allocates a pinned CPU cache per request, allowing CPU reads without D2H transfers and batched GPU uploads, at the cost of pinned - memory usage and upload progress tracking. pure_gpu: maintains the GPU request table directly without + memory usage and upload progress tracking. pin_mem_page_cache: stores one pinned CPU index per KV page, + reducing persistent index memory but expanding indexes for token reads and GPU uploads; page_size=1 + provides no compression. pure_gpu: maintains the GPU request table directly without a persistent CPU index cache; CPU reads require D2H transfers and may incur synchronization waits. - Both modes retain the same GPU request table and have approximately the same GPU memory usage.""", + All modes retain the same GPU request table and have approximately the same GPU memory usage.""", ) parser.add_argument( "--vit_att_backend", diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 967b83c479..5af19daa02 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -176,7 +176,9 @@ class StartArgs: default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "flashinfer"]} ) page_size: int = field(default=1) - req_token_index_mode: str = field(default="pin_mem_cache", metadata={"choices": ["pin_mem_cache", "pure_gpu"]}) + req_token_index_mode: str = field( + default="pin_mem_cache", metadata={"choices": ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]} + ) vit_att_backend: List[str] = field( default_factory=lambda: ["auto"], metadata={"choices": ["auto", "triton", "fa3", "sdpa", "xformers"]} ) diff --git a/lightllm/server/router/model_infer/infer_req_ext1.py b/lightllm/server/router/model_infer/infer_req_ext1.py index 3fa0e0fb92..7f384bd81a 100644 --- a/lightllm/server/router/model_infer/infer_req_ext1.py +++ b/lightllm/server/router/model_infer/infer_req_ext1.py @@ -51,10 +51,10 @@ def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: @abstractmethod def get_upload_indexes(self) -> Optional[torch.Tensor]: - """借用 [gpu_index_submitted_len, token_index_len) 的 CPU 视图，不更新提交进度。 + """返回 [gpu_index_submitted_len, token_index_len) 的 pinned CPU token 索引，不更新提交进度。 - 调用方可直接批量上传；缓存只追加，reset/扩容换用新存储，上传期间已有区间不变。 - GPU 模式返回 None。 + 完整 CPU 缓存借用内部视图，页缓存按需展开到独立暂存区，调用方均可直接批量上传。 + 缓存只追加，reset/扩容换用新存储，上传期间已有区间不变；GPU 模式返回 None。 """ pass @@ -62,7 +62,7 @@ def get_upload_indexes(self) -> Optional[torch.Tensor]: def get_free_token_indexes(self, start: int, end: int) -> torch.Tensor: """返回仅用于释放的 token 索引，可能是 CPU 或 GPU Tensor。 - pin_mem_cache 返回 CPU 快照；pure_gpu 返回借用的 GPU 请求表视图。 + 两种 CPU 缓存模式返回 CPU 快照；pure_gpu 返回借用的 GPU 请求表视图。 GPU 视图借用请求表，返回前让调用流等待索引流。重写请求槽位前，须在同一流 按序消费，或让索引流等待消费流。此接口不提供跨槽位复用的快照语义， @@ -155,6 +155,95 @@ def _reserve(self, end: int): self._cpu_buffer = new_buffer +class PinMemPageCacheReqTokenIndex(PinMemCacheReqTokenIndex): + """每个逻辑页仅保存首 token 的 pinned CPU 索引，对外仍使用 token 单位。""" + + # 同一进程使用统一的 page_size，所有请求共享首次初始化的页内偏移。 + _page_offsets: Optional[torch.Tensor] = None + # 同一推理进程的页缓存实例共享 slab 和空闲 offset，由推理线程串行申请。 + _pinned_buffer: Optional[torch.Tensor] = None + _pinned_offset = 0 + + def __init__(self, infer_req: "InferReq", index_stream: torch.cuda.Stream): + super().__init__(infer_req, index_stream) + self.page_size = infer_req.req_manager.page_size + if PinMemPageCacheReqTokenIndex._page_offsets is None: + PinMemPageCacheReqTokenIndex._page_offsets = torch.arange( + self.page_size, dtype=self.gpu_indexes.dtype, device="cpu", requires_grad=False + ) + + def append_token_indexes(self, mem_indexes: torch.Tensor): + """从页边界追加 token 索引；调用方保证同一逻辑页内索引连续。""" + append_len = mem_indexes.numel() + if append_len == 0: + return + assert mem_indexes.device.type == "cpu", "mem_indexes must be a CPU Tensor" + start = self.token_index_len + end = start + append_len + self._reserve(end) + # 输入从页首开始，每隔 page_size 个 token 保存一个页首索引。 + first_page = start // self.page_size + last_page = (end + self.page_size - 1) // self.page_size + self._cpu_buffer[first_page:last_page].copy_(mem_indexes[:: self.page_size]) + self.token_index_len = end + + def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + return self._expand_token_indexes(start, end, pin_memory=False) + + def get_upload_indexes(self) -> torch.Tensor: + return self._expand_token_indexes(self.gpu_index_submitted_len, self.token_index_len, pin_memory=True) + + def _reserve(self, end: int): + """首次按 GPU 请求表的最大长度申请全部页索引，后续追加直接复用。""" + if self._cpu_buffer is not None: + return + page_capacity = (self.gpu_indexes.numel() + self.page_size - 1) // self.page_size + self._cpu_buffer = torch.empty( + page_capacity, dtype=self.gpu_indexes.dtype, device="cpu", pin_memory=True, requires_grad=False + ) + + def _expand_token_indexes(self, start: int, end: int, pin_memory: bool) -> torch.Tensor: + """展开页对齐区间内的 token 索引，返回值不借用压缩缓存。""" + assert start % self.page_size == 0, "start must be page-aligned" + assert end % self.page_size == 0, "end must be page-aligned" + if start == end: + return torch.empty( + 0, dtype=self.gpu_indexes.dtype, device="cpu", pin_memory=pin_memory, requires_grad=False + ) + first_page = start // self.page_size + last_page = end // self.page_size + token_num = end - start + if pin_memory: + # 上传使用一次性的 slab 区间，后续追加、扩容或 reset 均不会覆写在途 H2D 的源。 + token_indexes = self._alloc_pinned_tensor(token_num) + else: + token_indexes = torch.empty(token_num, dtype=self.gpu_indexes.dtype, device="cpu", requires_grad=False) + page_starts = self._cpu_buffer[first_page:last_page] + # 将每个 page 的首索引加上页内偏移，直接展开到 token 级别的目标缓冲区。 + torch.add(page_starts[:, None], self._page_offsets, out=token_indexes.view(-1, self.page_size)) + return token_indexes + + @staticmethod + def _alloc_pinned_tensor(need_size: int) -> torch.Tensor: + """申请 int32 token 索引的 pinned 区间，slab 容量为 1M token（4 MiB）。""" + buffer_token_num = 1024 * 1024 + if need_size == 0 or need_size > buffer_token_num: + return torch.empty(need_size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True) + buffer = PinMemPageCacheReqTokenIndex._pinned_buffer + start = PinMemPageCacheReqTokenIndex._pinned_offset + end = start + need_size + if buffer is None or end > buffer_token_num: + buffer = torch.empty( + buffer_token_num, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True + ) + PinMemPageCacheReqTokenIndex._pinned_buffer = buffer + start = 0 + end = need_size + # 每次只创建返回值这一个视图；offset 只前进，换新 slab 时才归零，不覆写在途 H2D 的源。 + PinMemPageCacheReqTokenIndex._pinned_offset = end + return buffer[start:end] + + class PureGpuReqTokenIndex(ReqTokenIndex): """直接维护 GPU 请求表，不持有 CPU 索引缓存；CPU 读取会等待 D2H 完成。""" @@ -210,6 +299,10 @@ def create_req_token_index(infer_req: "InferReq", index_stream: torch.cuda.Strea mode = infer_req.args.req_token_index_mode if mode == "pin_mem_cache": return PinMemCacheReqTokenIndex(infer_req, index_stream) + if mode == "pin_mem_page_cache": + return PinMemPageCacheReqTokenIndex(infer_req, index_stream) if mode == "pure_gpu": return PureGpuReqTokenIndex(infer_req, index_stream) - raise ValueError(f"Unknown request token index mode: {mode}; expected pin_mem_cache or pure_gpu") + raise ValueError( + f"Unknown request token index mode: {mode}; expected pin_mem_cache, pin_mem_page_cache or pure_gpu" + ) 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 f48dd8a435..aa1a11094e 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -671,7 +671,7 @@ def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): assert len(index_streams) == 1, "All requests in a batch must use the same index stream" index_stream = index_streams.pop() current_stream = torch.cuda.current_stream(index_stream.device) - # 第二阶段：切出尚未提交的 CPU 源视图和 GPU 目标视图，记录本次提交终点。 + # 第二阶段：准备尚未提交的 CPU token 索引和 GPU 目标视图，记录本次提交终点。 cpu_index_slices, gpu_index_slices = [], [] pending_index_uploads = [] for req in run_reqs: @@ -694,7 +694,7 @@ def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): for token_index, end in pending_index_uploads: token_index.gpu_index_submitted_len = end # 第五阶段：调用流与索引流不同时建立等待依赖，保证后续 GPU 访问能读到更新。 - # 两种模式在没有新上传时，已提交的索引写入也可能尚未完成，仍需承接依赖。 + # 各模式在没有新上传时，已提交的索引写入也可能尚未完成，仍需承接依赖。 if current_stream != index_stream: current_stream.wait_stream(index_stream) diff --git a/test/benchmark/benchmark_req_token_index_alloc.py b/test/benchmark/benchmark_req_token_index_alloc.py new file mode 100644 index 0000000000..55aa202017 --- /dev/null +++ b/test/benchmark/benchmark_req_token_index_alloc.py @@ -0,0 +1,204 @@ +"""Measure page-cache token index pinned allocation, excluding data writes and H2D. + +Example: + CUDA_VISIBLE_DEVICES=4 PYTHONPATH=. python test/benchmark/benchmark_req_token_index_alloc.py \ + --cpu 72 --output-dir /tmp/req_token_index_alloc +""" + +import argparse +import csv +import gc +import hashlib +import json +import math +import os +from pathlib import Path +import platform +import statistics +import time + +import torch + +from lightllm.server.router.model_infer.infer_req_ext1 import PinMemPageCacheReqTokenIndex + + +SLAB_TOKENS = 1024 * 1024 +CLOCK = time.perf_counter_ns + + +def allocate_direct(token_num): + return torch.empty(token_num, dtype=torch.int32, device="cpu", pin_memory=True, requires_grad=False) + + +def summarize(samples): + ordered = sorted(samples) + return { + "count": len(samples), + "min_us": min(samples), + "max_us": max(samples), + "mean_us": statistics.mean(samples), + "median_us": statistics.median(samples), + "p95_us": ordered[int(0.95 * (len(ordered) - 1))], + } + + +def measure_cached(mode, token_num, repeats, backing): + allocate = allocate_direct if mode == "direct_cached" else PinMemPageCacheReqTokenIndex._alloc_pinned_tensor + samples = [] + for _ in range(repeats): + if mode == "slab_slice" and 0 < token_num <= SLAB_TOKENS: + # 仅测切片路径：当前没有存活的返回值或异步传输，允许在计时外重置测试 backing。 + # 生产接口始终向后切分；slab_cached 模式保留其真实换块行为。 + if ( + PinMemPageCacheReqTokenIndex._pinned_buffer is None + or PinMemPageCacheReqTokenIndex._pinned_offset + token_num > SLAB_TOKENS + ): + PinMemPageCacheReqTokenIndex._pinned_buffer = backing + PinMemPageCacheReqTokenIndex._pinned_offset = 0 + start = CLOCK() + tensor = allocate(token_num) + samples.append((CLOCK() - start) / 1000) + del tensor + return samples + + +def measure_growth(mode, token_num, live_count): + # 仅清理 HOST 空闲缓存；用于构造缓存未命中，不用于生产服务。 + PinMemPageCacheReqTokenIndex._pinned_buffer = None + PinMemPageCacheReqTokenIndex._pinned_offset = 0 + gc.collect() + torch.cuda.synchronize() + torch._C._host_emptyCache() + before = torch.cuda.memory.host_memory_stats() + allocate = allocate_direct if mode == "direct" else PinMemPageCacheReqTokenIndex._alloc_pinned_tensor + samples, held = [], [] + for _ in range(live_count): + start = CLOCK() + tensor = allocate(token_num) + samples.append((CLOCK() - start) / 1000) + held.append(tensor) + after = torch.cuda.memory.host_memory_stats() + host_calls = after["num_host_alloc"] - before["num_host_alloc"] + expected_calls = live_count + if mode == "slab" and token_num <= SLAB_TOKENS: + expected_calls = math.ceil(live_count / (SLAB_TOKENS // token_num)) + assert host_calls == expected_calls, (mode, token_num, host_calls, expected_calls) + assert all(tensor.is_pinned() and tensor.numel() == token_num for tensor in held) + return { + "samples_us": samples, + "summary": summarize(samples), + "host_allocations": host_calls, + "allocated_bytes": after["allocated_bytes.current"] - before["allocated_bytes.current"], + "host_alloc_time_us": after["host_alloc_time.total"] - before["host_alloc_time.total"], + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--cpu", type=int) + parser.add_argument("--rounds", type=int, default=10) + parser.add_argument("--repeats", type=int, default=2000) + args = parser.parse_args() + if args.cpu is not None: + os.sched_setaffinity(0, {args.cpu}) + args.output_dir.mkdir(parents=True, exist_ok=True) + torch.set_num_threads(1) + torch.cuda.init() + seed = allocate_direct(1) + del seed + torch._C._host_emptyCache() + start = CLOCK() + seed = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(1) + initial_slab_us = (CLOCK() - start) / 1000 + backing = seed._base + assert backing.numel() == SLAB_TOKENS and backing.is_pinned() + del seed + gc.collect() + gc.disable() + + sizes = (0, 128, 1024, 2048, 8192, 32768, 262144, SLAB_TOKENS // 2, SLAB_TOKENS, 2 * SLAB_TOKENS) + modes = ("direct_cached", "slab_slice", "slab_cached") + result = { + "metadata": { + "torch": torch.__version__, + "torch_git": torch.version.git_version, + "python": platform.python_version(), + "gpu": torch.cuda.get_device_name(0), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "cpu_affinity": sorted(os.sched_getaffinity(0)), + "torch_threads": torch.get_num_threads(), + "allocator_env": {name: os.environ.get(name) for name in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF")}, + "source_sha256": hashlib.sha256( + Path(PinMemPageCacheReqTokenIndex._alloc_pinned_tensor.__code__.co_filename).read_bytes() + ).hexdigest(), + "slab_tokens": SLAB_TOKENS, + "slab_bytes": SLAB_TOKENS * 4, + "scope": "CPU allocation wall time; excludes release, fill, H2D, imports, CUDA init and GC.", + "cached_modes": { + "direct_cached": "torch.empty with warm HOST size buckets", + "slab_slice": "slab slicing; backing reset outside timer when exhausted, no live or in-flight data", + "slab_cached": "actual shared allocator, including warm HOST-cache slab replacements", + }, + "growth": "Retain every returned tensor; clear unused HOST cache before each cohort; include fresh slabs.", + }, + "initial_slab_allocation_us": initial_slab_us, + "cached": {str(size): {mode: {"rounds_us": []} for mode in modes} for size in sizes}, + "growth": {}, + } + + # 预热所有直接申请桶，再准备两个可交替复用的 slab；验证计时区间没有调用 host 分配 API。 + for token_num in sizes: + seed = allocate_direct(token_num) + del seed + warm = [allocate_direct(SLAB_TOKENS) for _ in range(2)] + del warm + before = torch.cuda.memory.host_memory_stats()["num_host_alloc"] + for repeat in range(args.rounds): + for token_num in sizes if repeat % 2 == 0 else reversed(sizes): + for mode in modes if repeat % 2 == 0 else reversed(modes): + samples = measure_cached(mode, token_num, args.repeats, backing) + result["cached"][str(token_num)][mode]["rounds_us"].append(samples) + print(f"Cached round {repeat + 1}/{args.rounds} complete", flush=True) + result["cached_host_allocations"] = torch.cuda.memory.host_memory_stats()["num_host_alloc"] - before + assert result["cached_host_allocations"] == 0 + for data in result["cached"].values(): + for measured in data.values(): + rounds = measured["rounds_us"] + measured["summary"] = summarize([sample for samples in rounds for sample in samples]) + measured["round_medians_us"] = [statistics.median(samples) for samples in rounds] + + for token_num in (1024, 8192, 32768, 262144, 1048576, 2 * SLAB_TOKENS): + live_count = min(256, max(16, 16 * 1024 * 1024 // token_num)) + data = {mode: {"rounds": []} for mode in ("direct", "slab")} + data["live_count"] = live_count + for repeat in range(args.rounds): + for mode in ("direct", "slab") if repeat % 2 == 0 else ("slab", "direct"): + data[mode]["rounds"].append(measure_growth(mode, token_num, live_count)) + for mode in ("direct", "slab"): + rounds = data[mode]["rounds"] + data[mode]["summary"] = summarize([sample for batch in rounds for sample in batch["samples_us"]]) + data[mode]["round_mean_us"] = [batch["summary"]["mean_us"] for batch in rounds] + result["growth"][str(token_num)] = data + print(f"Uncached growth for {token_num} tokens x {live_count} complete", flush=True) + + PinMemPageCacheReqTokenIndex._pinned_buffer = None + PinMemPageCacheReqTokenIndex._pinned_offset = 0 + gc.enable() + (args.output_dir / "result.json").write_text(json.dumps(result, indent=2) + "\n") + with (args.output_dir / "timings.csv").open("w", newline="") as output: + writer = csv.writer(output) + fields = ("count", "min_us", "max_us", "mean_us", "median_us", "p95_us") + writer.writerow(("scenario", "tokens", "bytes", "mode", *fields)) + for scenario in ("cached", "growth"): + for tokens, data in result[scenario].items(): + for mode, measured in data.items(): + if mode != "live_count": + writer.writerow( + (scenario, tokens, int(tokens) * 4, mode, *[measured["summary"][f] for f in fields]) + ) + print(f"Results: {args.output_dir}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/unit_tests/common/test_kv_cache_allocator.py b/unit_tests/common/test_kv_cache_allocator.py index a083428403..f2fcfc4b41 100644 --- a/unit_tests/common/test_kv_cache_allocator.py +++ b/unit_tests/common/test_kv_cache_allocator.py @@ -245,7 +245,7 @@ def test_memory_manager_passes_page_size_and_resizes_token_capacity(make_allocat assert_capacity(manager.allocator, manager.size) -@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) @pytest.mark.parametrize("page_size", [1, 64]) def test_request_free_indexes_recycle_pages_after_reset_batch_and_slot_reuse( make_allocator, monkeypatch, mode, page_size @@ -253,15 +253,16 @@ def test_request_free_indexes_recycle_pages_after_reset_batch_and_slot_reuse( # 延迟导入，沿用在线推理初始化顺序，不增加另一套请求/allocator 夹具。 from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext - from lightllm.server.router.model_infer.infer_req_ext1 import create_req_token_index + from lightllm.server.router.model_infer.infer_req_ext1 import PinMemPageCacheReqTokenIndex, create_req_token_index + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_page_offsets", None) allocator = make_allocator(6 * page_size, page_size) allocated = allocator.alloc(allocator.size) first = torch.cat([allocated[4 * page_size : 5 * page_size], allocated[page_size : 2 * page_size]]).pin_memory() second = allocated[3 * page_size : 4 * page_size] expected = torch.cat([first, second]) table = torch.full((2, 2 * page_size), -1, dtype=torch.int32, device="cuda") - manager = SimpleNamespace(req_to_token_indexs=table) + manager = SimpleNamespace(req_to_token_indexs=table, page_size=page_size) context = InferenceContext.__new__(InferenceContext) context.req_manager = manager context.radix_cache = None diff --git a/unit_tests/common/test_req_manager_cpu.py b/unit_tests/common/test_req_manager_cpu.py index c36d3f7eee..0e8e64e906 100644 --- a/unit_tests/common/test_req_manager_cpu.py +++ b/unit_tests/common/test_req_manager_cpu.py @@ -11,7 +11,7 @@ from lightllm.common.req_manager.base import ReqManager, _ReqLinkedList from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext -from lightllm.server.router.model_infer.infer_req_ext1 import create_req_token_index +from lightllm.server.router.model_infer.infer_req_ext1 import PinMemPageCacheReqTokenIndex, create_req_token_index from lightllm.server.router.model_infer import infer_batch from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from unit_tests.common.cpu_index_helpers import bind_cpu_indexes, make_cpu_index_manager @@ -23,6 +23,7 @@ def reset_batch_copy_availability(monkeypatch): monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", True) monkeypatch.setattr(infer_batch.g_infer_context, "req_manager", None) + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_page_offsets", None) def _make_manager(device="cpu"): @@ -143,11 +144,12 @@ def forbid_cpu_export(*args, **kwargs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) @pytest.mark.parametrize("hybrid", [False, True]) @pytest.mark.parametrize("resuming", [False, True]) def test_prefix_match_appends_to_initialized_or_released_index_state(monkeypatch, mode, hybrid, resuming): manager = _make_manager("cuda") + manager.page_size = 4 req = InferReq.__new__(InferReq) req.req_idx = 0 req.req_manager = manager @@ -199,7 +201,7 @@ def unexpected_reset(): req._match_radix_cache() assert req.token_index.token_index_len == 4 - assert req.token_index.gpu_index_submitted_len == (0 if mode == "pin_mem_cache" else 4) + assert req.token_index.gpu_index_submitted_len == (4 if mode == "pure_gpu" else 0) assert req.cur_kv_len == req.hold_kv_len == 4 assert req.shared_kv_node is node assert req.shm_req.prompt_cache_len == req.shm_req.shm_cur_kv_len == 4 diff --git a/unit_tests/common/test_req_token_index.py b/unit_tests/common/test_req_token_index.py index e8eee3542d..49ccbd34dc 100644 --- a/unit_tests/common/test_req_token_index.py +++ b/unit_tests/common/test_req_token_index.py @@ -7,7 +7,9 @@ from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 from lightllm.common.req_manager import base from lightllm.server.router.model_infer.infer_req_ext1 import ( + ReqTokenIndex, PinMemCacheReqTokenIndex, + PinMemPageCacheReqTokenIndex, PureGpuReqTokenIndex, create_req_token_index, ) @@ -16,15 +18,23 @@ from lightllm.server.core.objs.start_args_type import StartArgs -MODES = ["pin_mem_cache", "pure_gpu"] +MODES = ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"] requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and pinned CPU memory") -def make_req(token_table, mode="pin_mem_cache", req_idx=0, input_len=0): +@pytest.fixture(autouse=True) +def reset_page_offsets(monkeypatch): + # 各用例模拟不同启动配置；隔离进程级页偏移，同一用例内的请求仍共享。 + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_page_offsets", None) + + +def make_req(token_table, mode="pin_mem_cache", req_idx=0, input_len=0, page_size=1): req = InferReq.__new__(InferReq) req.req_idx = req_idx - req.req_manager = SimpleNamespace(req_to_token_indexs=token_table) - req.args = StartArgs() if mode is None else StartArgs(req_token_index_mode=mode) + req.req_manager = SimpleNamespace(req_to_token_indexs=token_table, page_size=page_size) + req.args = ( + StartArgs(page_size=page_size) if mode is None else StartArgs(req_token_index_mode=mode, page_size=page_size) + ) req.shm_req = SimpleNamespace(input_len=input_len) return req @@ -34,6 +44,7 @@ def make_req(token_table, mode="pin_mem_cache", req_idx=0, input_len=0): [ (None, PinMemCacheReqTokenIndex), ("pin_mem_cache", PinMemCacheReqTokenIndex), + ("pin_mem_page_cache", PinMemPageCacheReqTokenIndex), ("pure_gpu", PureGpuReqTokenIndex), ], ) @@ -46,6 +57,9 @@ def test_factory_uses_request_mode_and_row_view_without_retaining_request(mode, assert not hasattr(InferReq, "move_to_gpu_index_len") assert isinstance(helper, expected_type) + assert hasattr(helper, "_alloc_pinned_tensor") == (expected_type is PinMemPageCacheReqTokenIndex) + assert not hasattr(ReqTokenIndex, "_pinned_buffer") + assert not hasattr(ReqTokenIndex, "_pinned_offset") assert helper._index_stream is index_stream assert helper.req_idx == req.req_idx assert helper.gpu_indexes.data_ptr() == table[req.req_idx].data_ptr() @@ -76,29 +90,85 @@ def test_start_args_and_cli_select_mode(mode): @requires_cuda @pytest.mark.parametrize("mode", MODES) @pytest.mark.parametrize("pinned", [False, True]) -def test_append_accepts_cpu_input_and_only_pins_when_gpu_upload_requires_it(monkeypatch, mode, pinned): +def test_append_accepts_cpu_input_and_allocates_pinned_storage_when_needed(monkeypatch, mode, pinned): table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") helper = create_req_token_index(make_req(table, mode), index_stream=torch.cuda.current_stream()) indexes = torch.arange(4, dtype=torch.int32, pin_memory=pinned) - pin_memory = torch.Tensor.pin_memory - pinned_inputs = [] + alloc_pinned = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor + pinned_sizes = [] - def record_pin_memory(tensor, *args, **kwargs): - pinned_inputs.append(tensor) - return pin_memory(tensor, *args, **kwargs) + def record_alloc_pinned(need_size): + pinned_sizes.append(need_size) + return alloc_pinned(need_size) - monkeypatch.setattr(torch.Tensor, "pin_memory", record_pin_memory) + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_alloc_pinned_tensor", staticmethod(record_alloc_pinned)) helper.append_token_indexes(indexes) # 普通 CPU 输入已经复制到自有 storage，可在 append 返回后立即改写。 if not pinned: indexes.fill_(-99) - assert len(pinned_inputs) == (1 if mode == "pure_gpu" and not pinned else 0) + assert pinned_sizes == [] assert helper.token_index_len == 4 assert helper.get_cpu_token_indexes(0, 4).tolist() == [0, 1, 2, 3] assert helper.get_gpu_token_indexes(0, 4).cpu().tolist() == [0, 1, 2, 3] - if mode == "pin_mem_cache": + if mode != "pure_gpu": assert helper._cpu_buffer.is_pinned() + assert helper._cpu_buffer.untyped_storage().nbytes() == helper._cpu_buffer.nbytes + + +@requires_cuda +def test_shared_pinned_allocator_preserves_slices_across_slab_replacement(monkeypatch): + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_pinned_buffer", None) + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_pinned_offset", 0) + buffer_token_num = 1024 * 1024 + first = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(5) + second = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(3) + first.fill_(11) + second.fill_(22) + assert second.data_ptr() == first.data_ptr() + first.nbytes + assert first.is_pinned() and second.is_pinned() + assert first.dtype == second.dtype == torch.int32 + assert not first.requires_grad and not second.requires_grad + + buffer = PinMemPageCacheReqTokenIndex._pinned_buffer + empty = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(0) + oversized = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(buffer_token_num + 1) + assert empty.numel() == 0 and oversized.numel() == buffer_token_num + 1 and oversized.is_pinned() + assert PinMemPageCacheReqTokenIndex._pinned_buffer is buffer + assert PinMemPageCacheReqTokenIndex._pinned_offset == 8 + assert oversized.untyped_storage().data_ptr() != first.untyped_storage().data_ptr() + + replacement = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(buffer_token_num - 7) + replacement.fill_(33) + assert replacement.untyped_storage().data_ptr() != first.untyped_storage().data_ptr() + assert first.tolist() == [11] * 5 and second.tolist() == [22] * 3 + assert PinMemPageCacheReqTokenIndex._pinned_buffer.numel() == buffer_token_num + assert PinMemPageCacheReqTokenIndex._pinned_offset == buffer_token_num - 7 + + +@requires_cuda +def test_shared_pinned_allocator_rollover_keeps_released_h2d_source_valid(monkeypatch): + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_pinned_buffer", None) + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_pinned_offset", 0) + buffer_token_num = 1024 * 1024 + source = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(buffer_token_num) + source.fill_(11) + source_ref = weakref.ref(source) + target = torch.empty(buffer_token_num, dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + torch.cuda._sleep(200_000_000) + target.copy_(source, non_blocking=True) + del source + assert source_ref() is None and not stream.query() + + # 连续换掉多个 slab，旧 storage 释放后仍须等在途 H2D 结束才能被重新使用。 + for value in range(22, 26): + replacement = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(buffer_token_num) + replacement.fill_(value) + del replacement + stream.synchronize() + assert target.eq(11).all().item() @requires_cuda @@ -171,27 +241,33 @@ def unexpected_upload(*args): backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) with pytest.raises(AssertionError, match="All requests in a batch must use the same index stream"): backend._copy_req_token_indexes_to_gpu(reqs) - assert [req.token_index.gpu_index_submitted_len for req in reqs] == [0 if mode == "pin_mem_cache" else 4] * 2 + assert [req.token_index.gpu_index_submitted_len for req in reqs] == [4 if mode == "pure_gpu" else 0] * 2 @requires_cuda @pytest.mark.parametrize("mode", MODES) @pytest.mark.parametrize("page_size", [1, 3, 4, 64]) -def test_token_units_allow_partial_page_appends_and_snapshot_ownership(mode, page_size): - row = torch.full((3 * page_size + 5,), -1, dtype=torch.int32, device="cuda") - req = make_req(row.unsqueeze(0), mode) +def test_token_units_and_snapshot_ownership(mode, page_size): + row = torch.full((4 * page_size,), -1, dtype=torch.int32, device="cuda") + req = make_req(row.unsqueeze(0), mode, page_size=page_size) helper = create_req_token_index(req, index_stream=torch.cuda.default_stream()) req.token_index = helper count = row.numel() - original = (torch.arange(count, dtype=torch.int32) * 3).pin_memory() - split = page_size + 1 + positions = torch.arange(count, dtype=torch.int32) + expected = positions.div(page_size, rounding_mode="floor") * (3 * page_size) + positions.remainder(page_size) + original = expected.pin_memory() + split = page_size if mode == "pin_mem_page_cache" else page_size + 1 req.token_index.append_token_indexes(original[:split]) req.token_index.append_token_indexes(torch.empty(0, dtype=torch.int32)) assert helper.token_index_len == split req.token_index.append_token_indexes(original[split:]) assert helper.token_index_len == count - expected = torch.arange(count, dtype=torch.int32) * 3 - for start, end in [(0, count), (1, count - 1), (page_size, page_size + 2), (count, count)]: + ranges = ( + [(0, count), (page_size, 3 * page_size), (count, count)] + if mode == "pin_mem_page_cache" + else [(0, count), (1, count - 1), (page_size, page_size + 2), (count, count)] + ) + for start, end in ranges: assert torch.equal(req.token_index.get_cpu_token_indexes(start, end), expected[start:end]) snapshot = req.token_index.get_cpu_token_indexes(0, count) # pure_gpu 的阻塞 CPU getter 已等待 H2D 完成，此后才可主动改写上传源。 @@ -202,19 +278,70 @@ def test_token_units_allow_partial_page_appends_and_snapshot_ownership(mode, pag gpu = req.token_index.get_gpu_token_indexes(0, count) assert gpu.data_ptr() == row.data_ptr() assert torch.equal(gpu.cpu(), expected) - if mode == "pin_mem_cache": + if mode != "pure_gpu": cpu_buffer = helper._cpu_buffer helper.reset() assert helper.token_index_len == helper.gpu_index_submitted_len == 0 - if mode == "pin_mem_cache": + if mode != "pure_gpu": assert helper._cpu_buffer is None else: assert not hasattr(helper, "_cpu_buffer") - helper.append_token_indexes(torch.zeros(count, dtype=torch.int32, pin_memory=True)) - if mode == "pin_mem_cache": + replacement = torch.arange(count, dtype=torch.int32, pin_memory=True) + helper.append_token_indexes(replacement) + if mode != "pure_gpu": assert helper._cpu_buffer is not cpu_buffer - assert helper.get_gpu_token_indexes(0, count).cpu().eq(0).all() - assert torch.equal(snapshot, torch.arange(count, dtype=torch.int32) * 3) + assert torch.equal(helper.get_gpu_token_indexes(0, count).cpu(), replacement) + assert torch.equal(snapshot, expected) + + +@requires_cuda +@pytest.mark.parametrize("page_size", [1, 3, 4, 64]) +def test_page_cache_compresses_page_starts_and_uploads_aligned_token_ranges(page_size): + table = torch.full((1, 8 * page_size), -1, dtype=torch.int32, device="cuda") + req = make_req(table, "pin_mem_page_cache", input_len=4 * page_size, page_size=page_size) + helper = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + page_starts = torch.tensor([9, 2, 15, 4], dtype=torch.int32) * page_size + expected = (page_starts[:, None] + torch.arange(page_size, dtype=torch.int32)).reshape(-1) + source = expected.clone() + start = 0 + for end in (page_size, 2 * page_size, 3 * page_size, expected.numel()): + helper.append_token_indexes(source[start:end]) + start = end + source.fill_(-99) + + assert helper._cpu_buffer.is_pinned() and helper._cpu_buffer.numel() == 8 + assert helper._cpu_buffer.nbytes * page_size == table.nbytes + assert torch.equal(helper._cpu_buffer[:4], page_starts) + for start, end in ((0, expected.numel()), (page_size, 3 * page_size), (page_size, page_size)): + assert torch.equal(helper.get_cpu_token_indexes(start, end), expected[start:end]) + assert torch.equal(helper.get_free_token_indexes(start, end), expected[start:end]) + + snapshot = helper.get_cpu_token_indexes(0, expected.numel()) + upload = helper.get_upload_indexes() + assert upload.is_pinned() and torch.equal(upload, expected) + assert helper.gpu_index_submitted_len == 0 + if page_size > 1: + with pytest.raises(AssertionError, match="start must be page-aligned"): + helper.get_cpu_token_indexes(1, page_size) + with pytest.raises(AssertionError, match="end must be page-aligned"): + helper.get_free_token_indexes(0, page_size - 1) + with pytest.raises(AssertionError, match="start must be page-aligned"): + helper.get_cpu_token_indexes(1, 1) + + # 模拟完整前缀页已经提交，待上传接口从下一个页边界开始展开。 + submitted = page_size + table[0, :submitted].copy_(upload[:submitted], non_blocking=True) + helper.gpu_index_submitted_len = submitted + pending = helper.get_upload_indexes() + assert pending.is_pinned() and torch.equal(pending, expected[submitted:]) + assert helper.gpu_index_submitted_len == submitted + assert torch.equal(helper.get_gpu_token_indexes(0, expected.numel()).cpu(), expected) + assert helper.get_upload_indexes().numel() == 0 + assert table[0, expected.numel() :].eq(-1).all() + + helper.reset() + helper.append_token_indexes(torch.arange(expected.numel(), dtype=torch.int32)) + assert torch.equal(snapshot, expected) and torch.equal(upload, expected) @requires_cuda @@ -268,7 +395,7 @@ def test_reset_rebuild_keeps_cpu_appends_and_gpu_submission_separate(mode): assert pending is None else: assert pending.tolist() == [0, 1, 2, 3] - assert helper.gpu_index_submitted_len == (0 if mode == "pin_mem_cache" else 4) + assert helper.gpu_index_submitted_len == (4 if mode == "pure_gpu" else 0) before = helper.get_cpu_token_indexes(0, 4) helper.get_gpu_token_indexes(0, 4) pending = helper.get_upload_indexes() @@ -279,18 +406,20 @@ def test_reset_rebuild_keeps_cpu_appends_and_gpu_submission_separate(mode): helper.reset() helper.append_token_indexes(torch.tensor([0, 31], dtype=torch.int32, pin_memory=True)) assert helper.token_index_len == 2 - assert helper.gpu_index_submitted_len == (0 if mode == "pin_mem_cache" else 2) - if mode == "pin_mem_cache": + assert helper.gpu_index_submitted_len == (2 if mode == "pure_gpu" else 0) + if mode != "pure_gpu": assert helper.gpu_indexes[:4].cpu().tolist() == [0, 1, 2, 3] helper.append_token_indexes(torch.tensor([32, 3], dtype=torch.int32, pin_memory=True)) assert helper.token_index_len == 4 - assert helper.gpu_index_submitted_len == (0 if mode == "pin_mem_cache" else 4) + assert helper.gpu_index_submitted_len == (4 if mode == "pure_gpu" else 0) assert helper.get_gpu_token_indexes(0, 4).cpu().tolist() == [0, 31, 32, 3] assert helper.gpu_index_submitted_len == 4 assert before.tolist() == [0, 1, 2, 3] @requires_cuda +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache"]) +@pytest.mark.parametrize("page_size", [1, 64]) @pytest.mark.parametrize( ("input_len", "max_length", "initial_capacity", "next_end", "grown_capacity"), [ @@ -301,49 +430,66 @@ def test_reset_rebuild_keeps_cpu_appends_and_gpu_submission_separate(mode): (1025, 5500, 4096, 5000, 5500), ], ) -def test_pin_buffer_reserves_input_length_and_grows_with_headroom( - input_len, max_length, initial_capacity, next_end, grown_capacity +def test_pin_buffer_capacity_and_lifecycle( + mode, page_size, input_len, max_length, initial_capacity, next_end, grown_capacity ): - req = make_req(torch.empty((1, max_length), dtype=torch.int32, device="cuda"), input_len=input_len) - helper = PinMemCacheReqTokenIndex(req, index_stream=torch.cuda.default_stream()) + req = make_req( + torch.empty((1, max_length), dtype=torch.int32, device="cuda"), mode, input_len=input_len, page_size=page_size + ) + helper = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + storage_page_size = page_size if mode == "pin_mem_page_cache" else 1 + if mode == "pin_mem_page_cache": + next_end = (next_end + page_size - 1) // page_size * page_size + initial_elements = grown_elements = (max_length + page_size - 1) // page_size + else: + initial_elements = initial_capacity + grown_elements = grown_capacity assert helper.input_len == input_len assert helper._cpu_buffer is None helper.append_token_indexes(torch.empty(0, dtype=torch.int32)) assert helper._cpu_buffer is None - # 首次只追加一个索引，仍按完整输入长度额外预留 2K 输出 token。 - first_index = torch.zeros(1, dtype=torch.int32, pin_memory=True) + # 首次追加最小存储单位；完整模式预留 2K 输出，页模式一次申请到最大长度。 + first_index = torch.arange(storage_page_size, dtype=torch.int32, pin_memory=True) helper.append_token_indexes(first_index) - assert helper._cpu_buffer.numel() == initial_capacity + assert helper._cpu_buffer.numel() == initial_elements assert helper._cpu_buffer.is_pinned() and helper._cpu_buffer.data_ptr() != first_index.data_ptr() - helper.append_token_indexes(torch.arange(1, initial_capacity, dtype=torch.int32, pin_memory=True)) + helper.append_token_indexes(torch.arange(storage_page_size, initial_capacity, dtype=torch.int32, pin_memory=True)) old_buffer = helper._cpu_buffer helper.append_token_indexes(torch.arange(initial_capacity, next_end, dtype=torch.int32, pin_memory=True)) - assert helper._cpu_buffer.numel() == grown_capacity and helper._cpu_buffer.is_pinned() - assert helper._cpu_buffer.data_ptr() != old_buffer.data_ptr() + assert helper._cpu_buffer.numel() == grown_elements and helper._cpu_buffer.is_pinned() + if mode == "pin_mem_page_cache": + assert helper._cpu_buffer is old_buffer + else: + assert helper._cpu_buffer.data_ptr() != old_buffer.data_ptr() snapshot = helper.get_cpu_token_indexes(0, next_end) assert snapshot.tolist() == list(range(next_end)) grown_buffer = helper._cpu_buffer helper.reset() assert helper.input_len == input_len and helper._cpu_buffer is None assert helper.token_index_len == helper.gpu_index_submitted_len == 0 - helper.append_token_indexes(torch.tensor([99], dtype=torch.int32, pin_memory=True)) + replacement = torch.arange(99, 99 + storage_page_size, dtype=torch.int32, pin_memory=True) + helper.append_token_indexes(replacement) assert helper._cpu_buffer is not grown_buffer - assert helper._cpu_buffer.numel() == initial_capacity - assert helper.token_index_len == 1 - assert helper.get_cpu_token_indexes(0, helper.token_index_len).tolist() == [99] + assert helper._cpu_buffer.numel() == initial_elements + assert helper.token_index_len == storage_page_size + assert torch.equal(helper.get_cpu_token_indexes(0, helper.token_index_len), replacement) assert snapshot.tolist() == list(range(next_end)) @requires_cuda +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache"]) +@pytest.mark.parametrize("page_size", [1, 64]) @pytest.mark.parametrize("use_default_stream", [False, True]) -def test_pin_buffer_reset_during_upload_keeps_old_data_until_h2d_completes(use_default_stream): +def test_pin_buffer_reset_during_upload_keeps_old_data_until_h2d_completes(mode, page_size, use_default_stream): count = 8192 row = torch.full((count,), -1, dtype=torch.int32, device="cuda") index_stream = torch.cuda.default_stream() if use_default_stream else torch.cuda.Stream() index_stream.wait_stream(torch.cuda.current_stream()) - helper = PinMemCacheReqTokenIndex(make_req(row.unsqueeze(0), input_len=count), index_stream=index_stream) + helper = create_req_token_index( + make_req(row.unsqueeze(0), mode, input_len=count, page_size=page_size), index_stream=index_stream + ) original = torch.arange(count, dtype=torch.int32, pin_memory=True) - replacement = torch.full((count,), -99, dtype=torch.int32, pin_memory=True) + replacement = (torch.arange(count, dtype=torch.int32).remainder(page_size) + count * 2).pin_memory() helper.append_token_indexes(original) source_ref = weakref.ref(helper._cpu_buffer) @@ -448,15 +594,17 @@ def test_manager_owns_only_gpu_table_and_initializes_hold_page(monkeypatch, page @requires_cuda @pytest.mark.parametrize("batch_available", [False, True]) @pytest.mark.parametrize("mode", MODES) +@pytest.mark.parametrize("page_size", [1, 64]) @pytest.mark.parametrize("index_stream_kind", ["default", "shared"]) def test_backend_uploads_pinned_views_without_request_events_or_keepalive_list( - monkeypatch, mode, batch_available, index_stream_kind + monkeypatch, mode, page_size, batch_available, index_stream_kind ): from lightllm.common import cuda_wrapper from lightllm.server.router.model_infer.mode_backend import base_backend - table = torch.full((2, 16), -1, dtype=torch.int32, device="cuda") - expected = (torch.arange(16, dtype=torch.int32).view(2, 8) + 32).pin_memory() + count = max(8, page_size) + table = torch.full((2, 2 * count), -1, dtype=torch.int32, device="cuda") + expected = (torch.arange(2 * count, dtype=torch.int32).view(2, count) + 32).pin_memory() index_streams = [torch.cuda.default_stream()] * 2 if index_stream_kind == "shared": index_streams = [torch.cuda.Stream()] * 2 @@ -464,14 +612,14 @@ def test_backend_uploads_pinned_views_without_request_events_or_keepalive_list( index_stream.wait_stream(torch.cuda.default_stream()) reqs = [] for row, index_stream in enumerate(index_streams): - req = make_req(table, mode, req_idx=row) + req = make_req(table, mode, req_idx=row, page_size=page_size) req.token_index = create_req_token_index(req, index_stream=index_stream) with torch.cuda.stream(index_stream): # 批量上传须先等待旧写入；pure_gpu 在本流追加，调用流尚未承接依赖。 torch.cuda._sleep(60_000_000) - table[row, :8].fill_(-2) + table[row, :count].fill_(-2) req.token_index.append_token_indexes(expected[row]) - req.hold_kv_len = 8 + req.hold_kv_len = count reqs.append(req) monkeypatch.setattr(base_backend, "g_infer_context", SimpleNamespace(req_manager=reqs[0].req_manager)) monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", batch_available) @@ -486,10 +634,13 @@ def test_backend_uploads_pinned_views_without_request_events_or_keepalive_list( event_type = torch.cuda.Event def record_batch(dst_tensors, src_tensors): - # 直接借用各请求缓存的待上传区间，不能再创建聚合快照或改成逐请求上传。 - assert [tensor.data_ptr() for tensor in src_tensors] == [ - req.token_index.get_upload_indexes().data_ptr() for req in reqs - ] + # 完整缓存直接借用视图，页缓存先展开；两者都应一次批量提交正确的 token 索引。 + assert all(tensor.is_pinned() for tensor in src_tensors) + assert all(torch.equal(tensor, expected[row]) for row, tensor in enumerate(src_tensors)) + if mode == "pin_mem_cache": + assert [tensor.data_ptr() for tensor in src_tensors] == [ + req.token_index.get_upload_indexes().data_ptr() for req in reqs + ] batch_calls.append(len(src_tensors)) return copy_batch(dst_tensors, src_tensors) @@ -506,22 +657,22 @@ def record_event(*args, **kwargs): for req, index_stream in zip(reqs, index_streams): with torch.cuda.stream(index_stream): # 不等待后台线程：已标记提交的索引必须能在请求自己的索引流立即使用。 - observed.append(req.token_index.get_gpu_token_indexes(0, 8).clone()) - if mode == "pin_mem_cache": + observed.append(req.token_index.get_gpu_token_indexes(0, count).clone()) + if mode != "pure_gpu": for req in reqs: req.token_index.reset() - req.token_index.append_token_indexes(torch.zeros(8, dtype=torch.int32, pin_memory=True)) + req.token_index.append_token_indexes(torch.arange(count, dtype=torch.int32, pin_memory=True)) # 两个请求只建立一个 batch 保活事件；pure_gpu 与 PyTorch 回退无需该事件。 - assert len(events) == (1 if mode == "pin_mem_cache" and batch_available else 0) - assert batch_calls == ([2] if mode == "pin_mem_cache" else []) + assert len(events) == (1 if mode != "pure_gpu" and batch_available else 0) + assert batch_calls == ([] if mode == "pure_gpu" else [2]) pool = cuda_wrapper.PinnedH2DBatchCopier._release_pool if pool is not None: pool.submit(lambda: None).result(timeout=10) - assert [req.token_index.gpu_index_submitted_len for req in reqs] == [0 if mode == "pin_mem_cache" else 8] * 2 + assert [req.token_index.gpu_index_submitted_len for req in reqs] == [count if mode == "pure_gpu" else 0] * 2 assert not hasattr(backend, "_pending_cpu_index_uploads") assert not hasattr(backend, "_index_upload_stream") assert all(not hasattr(req.token_index, "_gpu_write_event") for req in reqs) - assert torch.equal(table[:, :8].cpu(), expected) + assert torch.equal(table[:, :count].cpu(), expected) for actual, expected_row, index_stream in zip(observed, expected, index_streams): with torch.cuda.stream(index_stream): assert torch.equal(actual.cpu(), expected_row) 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 b3d504956f..409a9656e1 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 @@ -13,7 +13,7 @@ from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as multi_level_kv_cache_impl from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes -from lightllm.server.router.model_infer.infer_req_ext1 import create_req_token_index +from lightllm.server.router.model_infer.infer_req_ext1 import PinMemPageCacheReqTokenIndex, create_req_token_index def test_cache_tiers_reassignment_is_rejected(): @@ -168,9 +168,10 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): assert req.hybrid_len_to_big_page_id == {} -@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires pinned CPU buffers") def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch, mode): + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_page_offsets", None) table = torch.full((1, 136), -1, dtype=torch.int32) gpu_table = torch.full_like(table, -2) table[0, :4] = torch.arange(4, dtype=torch.int32) @@ -262,10 +263,11 @@ def unexpected_copy(*args, **kwargs): assert gpu_table[0, 132:].eq(-2).all() -@pytest.mark.parametrize("mode", ["pin_mem_cache", "pure_gpu"]) +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) @pytest.mark.parametrize("use_default_index_stream", [False, True]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_cpu_cache_offload_getter_orders_transfer_stream(monkeypatch, mode, use_default_index_stream): + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_page_offsets", None) default_stream = torch.cuda.default_stream() index_stream = default_stream if use_default_index_stream else torch.cuda.Stream() index_stream.wait_stream(default_stream) @@ -273,7 +275,7 @@ def test_cpu_cache_offload_getter_orders_transfer_stream(monkeypatch, mode, use_ gpu_indexes = torch.full((64,), -1, dtype=torch.int32, device="cuda") req = SimpleNamespace( req_idx=0, - req_manager=SimpleNamespace(req_to_token_indexs=gpu_indexes.unsqueeze(0)), + req_manager=SimpleNamespace(req_to_token_indexs=gpu_indexes.unsqueeze(0), page_size=64), args=SimpleNamespace(req_token_index_mode=mode), cache_tiers=(CacheTier.CPU,), cur_kv_len=64, From e77f3d158353706ad430c570e2ec2f8c5a915149 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 07:03:26 +0000 Subject: [PATCH 14/22] Fix partial token reads in page-cached request indexes --- docs/kv_cache_page_size.md | 5 +- .../router/model_infer/infer_req_ext1.py | 12 ++--- unit_tests/common/test_req_token_index.py | 29 +++++------ .../test_paged_side_allocations.py | 50 +++++++++++++++++++ .../test_prompt_logprobs_cache.py | 43 ++++++++++++++++ 5 files changed, 114 insertions(+), 25 deletions(-) diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index 1187655fd8..ed2aad2f36 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -61,7 +61,8 @@ Helper 保存该流，后续不随调用点的当前流切换： 一次性申请 `ceil(max_sequence_length / page_size)` 个页首索引，后续追加无需扩容或复制旧缓存。 页模式继承完整模式的 GPU 上传、释放快照和 reset 流程，覆盖 CPU 索引的存储和展开。 压缩依赖已有的页内连续索引约束；追加完整页，输入从下标 0 按 page_size 取页首。 -页模式展开 CPU 快照、释放索引或待上传索引时，断言 start 和 end 均为页边界，直接返回完整展开结果。 +页模式读取支持页内任意 token 区间：内部展开覆盖区间的完整页，再按页内偏移截取精确范围。 +CPU 快照独立于压缩缓存，PD 尾段传输和最终 metadata 无需自行处理页对齐。 缓冲区在首次非空写入时按需申请，重置时保留输入长度， 丢弃已申请的缓冲区引用，由 PyTorch 在引用和异步传输依赖结束后回收。 页模式上传时先将待提交区间展开到临时 pinned token 索引，使用 `PinMemPageCacheReqTokenIndex._alloc_pinned_tensor` @@ -77,7 +78,7 @@ offset 只前进，不回绕；空间不足时换用新 slab 并从零开始， `get_free_token_indexes(start, end)` 和 `get_gpu_token_indexes(start, end)`，均使用 token 单位， 追加接口将完整输入接到已记录索引的尾部，非空输入在三种模式下均断言为 CPU Tensor；空输入直接返回。 CPU 缓存模式直接复制到自身的 pinned 缓冲区，无需先 pin 输入；GPU 模式仅在输入未 pinned 时内部转换，已有 pinned 输入直接复用。 -前缀匹配和索引重建路径直接传入 CPU 索引，不再自行 pin。GPU 视图读取支持页内任意区间；页模式 CPU 展开和 allocator 释放使用完整页区间。 +前缀匹配和索引重建路径直接传入 CPU 索引，不再自行 pin。读取接口均支持页内任意区间；提交给 allocator 的释放区间仍须覆盖完整页。 CPU getter 返回独立快照，缓存、PD 元数据和异步任务可跨请求释放持有；GPU getter 返回当前请求的 GPU 表视图，不应跨请求槽位复用持有。Helper 保存外部传入的索引流，GPU 索引读写固定在该流， 调用方保证 token 区间有效；GPU getter 返回前让调用流等待索引流，没有新增上传时也需承接已有写入的依赖。 diff --git a/lightllm/server/router/model_infer/infer_req_ext1.py b/lightllm/server/router/model_infer/infer_req_ext1.py index 7f384bd81a..517e7003b6 100644 --- a/lightllm/server/router/model_infer/infer_req_ext1.py +++ b/lightllm/server/router/model_infer/infer_req_ext1.py @@ -203,16 +203,14 @@ def _reserve(self, end: int): ) def _expand_token_indexes(self, start: int, end: int, pin_memory: bool) -> torch.Tensor: - """展开页对齐区间内的 token 索引，返回值不借用压缩缓存。""" - assert start % self.page_size == 0, "start must be page-aligned" - assert end % self.page_size == 0, "end must be page-aligned" + """展开覆盖读取区间的完整页，再截取精确 token 范围，返回值不借用压缩缓存。""" if start == end: return torch.empty( 0, dtype=self.gpu_indexes.dtype, device="cpu", pin_memory=pin_memory, requires_grad=False ) first_page = start // self.page_size - last_page = end // self.page_size - token_num = end - start + last_page = (end + self.page_size - 1) // self.page_size + token_num = (last_page - first_page) * self.page_size if pin_memory: # 上传使用一次性的 slab 区间，后续追加、扩容或 reset 均不会覆写在途 H2D 的源。 token_indexes = self._alloc_pinned_tensor(token_num) @@ -221,7 +219,9 @@ def _expand_token_indexes(self, start: int, end: int, pin_memory: bool) -> torch page_starts = self._cpu_buffer[first_page:last_page] # 将每个 page 的首索引加上页内偏移，直接展开到 token 级别的目标缓冲区。 torch.add(page_starts[:, None], self._page_offsets, out=token_indexes.view(-1, self.page_size)) - return token_indexes + # 读取可从页内任意位置开始或结束，页内偏移只影响返回范围，不改变压缩页缓存。 + offset = start % self.page_size + return token_indexes[offset : offset + end - start] @staticmethod def _alloc_pinned_tensor(need_size: int) -> torch.Tensor: diff --git a/unit_tests/common/test_req_token_index.py b/unit_tests/common/test_req_token_index.py index 49ccbd34dc..71575669f4 100644 --- a/unit_tests/common/test_req_token_index.py +++ b/unit_tests/common/test_req_token_index.py @@ -262,11 +262,7 @@ def test_token_units_and_snapshot_ownership(mode, page_size): assert helper.token_index_len == split req.token_index.append_token_indexes(original[split:]) assert helper.token_index_len == count - ranges = ( - [(0, count), (page_size, 3 * page_size), (count, count)] - if mode == "pin_mem_page_cache" - else [(0, count), (1, count - 1), (page_size, page_size + 2), (count, count)] - ) + ranges = [(0, count), (1, count - 1), (page_size, page_size + 2), (count, count)] for start, end in ranges: assert torch.equal(req.token_index.get_cpu_token_indexes(start, end), expected[start:end]) snapshot = req.token_index.get_cpu_token_indexes(0, count) @@ -296,7 +292,7 @@ def test_token_units_and_snapshot_ownership(mode, page_size): @requires_cuda @pytest.mark.parametrize("page_size", [1, 3, 4, 64]) -def test_page_cache_compresses_page_starts_and_uploads_aligned_token_ranges(page_size): +def test_page_cache_compresses_page_starts_and_reads_partial_token_ranges(page_size): table = torch.full((1, 8 * page_size), -1, dtype=torch.int32, device="cuda") req = make_req(table, "pin_mem_page_cache", input_len=4 * page_size, page_size=page_size) helper = create_req_token_index(req, index_stream=torch.cuda.default_stream()) @@ -312,7 +308,14 @@ def test_page_cache_compresses_page_starts_and_uploads_aligned_token_ranges(page assert helper._cpu_buffer.is_pinned() and helper._cpu_buffer.numel() == 8 assert helper._cpu_buffer.nbytes * page_size == table.nbytes assert torch.equal(helper._cpu_buffer[:4], page_starts) - for start, end in ((0, expected.numel()), (page_size, 3 * page_size), (page_size, page_size)): + for start, end in ( + (0, expected.numel()), + (1, expected.numel() - 1), + (page_size - 1, 2 * page_size + 1), + (1, 2), + (1, 1), + (page_size, page_size), + ): assert torch.equal(helper.get_cpu_token_indexes(start, end), expected[start:end]) assert torch.equal(helper.get_free_token_indexes(start, end), expected[start:end]) @@ -320,16 +323,8 @@ def test_page_cache_compresses_page_starts_and_uploads_aligned_token_ranges(page upload = helper.get_upload_indexes() assert upload.is_pinned() and torch.equal(upload, expected) assert helper.gpu_index_submitted_len == 0 - if page_size > 1: - with pytest.raises(AssertionError, match="start must be page-aligned"): - helper.get_cpu_token_indexes(1, page_size) - with pytest.raises(AssertionError, match="end must be page-aligned"): - helper.get_free_token_indexes(0, page_size - 1) - with pytest.raises(AssertionError, match="start must be page-aligned"): - helper.get_cpu_token_indexes(1, 1) - - # 模拟完整前缀页已经提交，待上传接口从下一个页边界开始展开。 - submitted = page_size + # 模拟已提交前缀结束于页内，待上传接口仍返回精确的剩余 token 区间。 + submitted = page_size + 1 table[0, :submitted].copy_(upload[:submitted], non_blocking=True) helper.gpu_index_submitted_len = submitted pending = helper.get_upload_indexes() 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 index 1719180a8e..ad37c2c3a9 100644 --- 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 @@ -8,6 +8,7 @@ # 在独立运行本测试文件时形成模块级循环导入。 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.infer_req_ext1 import PinMemPageCacheReqTokenIndex, create_req_token_index from lightllm.server.router.model_infer.mode_backend import base_backend 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 @@ -85,6 +86,55 @@ def create_task(**kwargs): assert table[0, :12].tolist() == list(range(12)) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and pinned CPU memory") +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("input_len", [128, 130]) +def test_pd_prefill_transfers_exact_logical_tail_with_real_index_helper(monkeypatch, mode, input_len): + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_page_offsets", None) + monkeypatch.setattr(prefill_impl.g_infer_context, "is_hybrid_att_model", False) + page_size = 64 + token_num = (input_len + page_size - 1) // page_size * page_size + req = SimpleNamespace( + req_idx=0, + req_id=1, + req_manager=SimpleNamespace( + page_size=page_size, req_to_token_indexs=torch.empty((1, token_num), dtype=torch.int32, device="cuda") + ), + args=SimpleNamespace(req_token_index_mode=mode), + shm_req=SimpleNamespace(input_len=input_len), + sampling_param=SimpleNamespace( + pd_master_node_id=1, + pd_decode_node=SimpleNamespace(agent_name="decode", agent_metadata=None, num_pages=1, page_reg_desc=None), + ), + cur_kv_len=input_len, + hold_kv_len=token_num, + pd_trans_kv_start_index=0, + pd_trans_device_id=0, + pd_task_num=0, + ) + req.token_index = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + page_starts = torch.tensor([9, 2, 15], dtype=torch.int32)[: token_num // page_size] * page_size + expected = (page_starts[:, None] + torch.arange(page_size, dtype=torch.int32)).reshape(-1) + req.token_index.append_token_indexes(expected) + backend = prefill_impl.PDChunkedPrefillForPrefillNode.__new__(prefill_impl.PDChunkedPrefillForPrefillNode) + backend.args = SimpleNamespace(page_size=page_size, pd_kv_page_size=128) + backend.is_master_in_dp = True + backend.dp_rank_in_node = 0 + tasks = [] + backend.info_queue = SimpleNamespace(put=tasks.append) + + backend._prefill_chuncked_handle_func(req, next_token_id=7, next_token_prob=-0.5, output_len=1) + + assert req.pd_task_num == len(tasks) + assert req.pd_trans_kv_start_index == input_len + assert [(task.start_kv_index, task.end_kv_index) for task in tasks] == ( + [(0, 128)] if input_len == 128 else [(0, 128), (128, 130)] + ) + assert [index for task in tasks for index in task.mem_indexes] == expected[:input_len].tolist() + assert tasks[-1].first_gen_token_id == 7 + assert tasks[-1].first_gen_token_logprob == -0.5 + + 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) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py index 4c4aa6c46a..7450e23dd4 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_prompt_logprobs_cache.py @@ -9,6 +9,7 @@ from lightllm.server.core.objs.sampling_params import SamplingParams from lightllm.server.router.model_infer import infer_batch from lightllm.server.router.model_infer.infer_batch import InferReq, InferSamplingParams +from lightllm.server.router.model_infer.infer_req_ext1 import PinMemPageCacheReqTokenIndex, create_req_token_index from lightllm.server.router.model_infer.mode_backend import base_backend from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes @@ -127,3 +128,45 @@ def test_final_metadata_extracts_owned_token_indexes_through_req_helper(monkeypa req.token_index.reset() req.token_index.append_token_indexes(torch.arange(6, dtype=torch.int32)) assert indexes.tolist() == [8, 9, 2, 3, 14] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and pinned CPU memory") +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("input_len", [128, 130]) +@pytest.mark.parametrize("kind", ["prompt_logprobs", "routed_experts"]) +def test_final_metadata_reads_partial_pages_with_real_index_helper(monkeypatch, mode, input_len, kind): + from lightllm.server.router.model_infer import infer_req_ext + + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_page_offsets", None) + page_size = 64 + token_num = (input_len + page_size - 1) // page_size * page_size + req = SimpleNamespace( + req_idx=0, + req_manager=SimpleNamespace( + page_size=page_size, req_to_token_indexs=torch.empty((1, token_num), dtype=torch.int32, device="cuda") + ), + args=SimpleNamespace(req_token_index_mode=mode), + shm_req=SimpleNamespace(input_len=input_len), + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=3)), + cur_kv_len=input_len, + cur_output_len=0, + ) + req.token_index = create_req_token_index(req, index_stream=torch.cuda.default_stream()) + page_starts = torch.tensor([9, 2, 15], dtype=torch.int32)[: token_num // page_size] * page_size + expected = (page_starts[:, None] + torch.arange(page_size, dtype=torch.int32)).reshape(-1) + req.token_index.append_token_indexes(expected) + manager = SimpleNamespace(extract=Mock(return_value="metadata")) + metadata = infer_req_ext.FinalTokenMetadataExt(req) + if kind == "prompt_logprobs": + monkeypatch.setattr(infer_req_ext.PromptLogprobsCaptureManager, "get_instance", lambda: manager) + result = metadata.collect_prompt_logprobs() + else: + monkeypatch.setattr(infer_req_ext.MoeRouteInfoManager, "get_instance", lambda: manager) + result = metadata.collect_routed_experts() + + assert result == "metadata" + indexes = manager.extract.call_args.args[0] + assert torch.equal(indexes, expected[: input_len - 1]) + req.token_index.reset() + req.token_index.append_token_indexes(torch.arange(token_num, dtype=torch.int32)) + assert torch.equal(indexes, expected[: input_len - 1]) From 9c0107b8abe855b723387203faadcb4ddfcc531f Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 07:26:24 +0000 Subject: [PATCH 15/22] fix doc --- docs/kv_cache_page_size.md | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index ed2aad2f36..80387a9fa3 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -195,3 +195,36 @@ diverse mode 仍只支持 `page_size=1`，启动阶段会对其他取值明确 `[HOLD, HOLD+1, ..., HOLD+page_size-1]` 循环填充，而不是重复同一个 token 索引。 - Radix 子节点用“首个完整 token 页”作为键，避免不同序列仅首 token 相同造成页级分支冲突。 - 非法的 `page_size < 1` 以及尚未支持的功能组合在模型加载前失败。 + +## 附录：选择批量异步 H2D 的原因 + +当前实现采用“CPU 写入 pinned 索引 → CUDA 异步 H2D → GPU kernel 读取 device 请求表”的流程。 +选择显式拷贝的一项工程依据是：根据以往的使用经验，CPU 写入 pinned 内存后，直接让 GPU kernel +读取这块 host 内存的方案，曾在极端时序下观察到读取值异常。因此，本实现通过 CUDA 内存传输接口 +建立明确的数据交接过程，模型 kernel 统一读取 GPU 上的 `req_to_token_indexs`。 + +排查上述现象时，需要考虑 CPU 存储缓冲、写合并、cache 一致性以及 CPU/GPU 访问顺序对数据可见性的影响。 +“CPU 写入仍滞留于写缓冲，尚未对设备可见”是历史排查中考虑的可能原因，本文没有将其确认为所有平台上的根因。 +具体的 cache 维护与可见性实现取决于硬件和驱动，正确性应依赖 CUDA API 的访问顺序与完成语义， +不假定每次异步拷贝都会执行固定的 CPU cache flush，或必须先把数据写回 DRAM 才能被设备正确读取。 +CUDA 也支持正确同步的 mapped pinned memory；其文档明确要求应用协调 host/device 访问，避免读写竞争。 +参见 [CUDA 13.0：Mapped Memory](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-c-programming-guide/index.html#mapped-memory)。 + +本实现的正确性由以下约束共同保证： + +1. CPU 完成源索引写入后才提交 H2D；上传完成前不再改写已提交的源区间。 +2. `cudaMemcpyBatchAsync` 使用 `cudaMemcpySrcAccessOrderStream`，源读取遵循上传流顺序； + API 返回只表示已经提交，源内存须保持有效到实际传输结束。 +3. GPU 消费流通过已有的流等待承接上传依赖，随后读取 device 请求表。 +4. 原生 batch 路径由完成事件和后台任务持有 pinned 源引用；PyTorch `copy_` 回退路径 + 由其 host allocator 跟踪异步使用，避免传输完成前回收复用源内存。 + +源访问顺序的定义参见 [CUDA 13.0 Runtime API：Memory Management](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-runtime-api/group__CUDART__MEMORY.html)。 +批量接口用于将多个请求的索引上传合并为一次 API 提交，并直接使用各请求的 pinned 视图，避免额外的 CPU 聚合拷贝； +这些正确性约束同样适用于普通的异步 H2D，batch 接口本身不能替代源保活、访问排序和禁止在途改写。 + +PyTorch 文档介绍了使用 pinned CPU Tensor 配合 `to(..., non_blocking=True)` / `cuda(non_blocking=True)` 进行异步传输的方式， +并说明 pinned 源在传输完成前被 CPU 改写可能导致 GPU 数据错误。本实现遵循这一显式传输范式， +`cudaMemcpyBatchAsync` 则是多个索引区间的批量提交选择。 +参见 [PyTorch 2.11：Use pinned memory buffers](https://docs.pytorch.org/docs/2.11/notes/cuda.html#use-pinned-memory-buffers) +和 [PyTorch：non_blocking 与 pin_memory 使用指南](https://docs.pytorch.org/tutorials/intermediate/pinmem_nonblock.html#other-copy-directions-gpu-cpu-cpu-mps)。 From 760e8aaeb92b24c4bd5649b8fa19059bff34c12f Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 07:48:22 +0000 Subject: [PATCH 16/22] Validate request token index bounds and page appends Align CPU read bounds across index modes and reject partial-page appends before modifying page cache state. Extend regression coverage and add a standalone pinned allocation rounding and reuse prototype. --- .../router/model_infer/infer_req_ext1.py | 10 +- .../benchmark_pinned_allocation_rounding.py | 127 ++++++++++++++++++ unit_tests/common/test_req_token_index.py | 34 ++++- .../mode_backend/test_diverse_backend.py | 1 + 4 files changed, 163 insertions(+), 9 deletions(-) create mode 100644 test/benchmark/benchmark_pinned_allocation_rounding.py diff --git a/lightllm/server/router/model_infer/infer_req_ext1.py b/lightllm/server/router/model_infer/infer_req_ext1.py index 517e7003b6..bb4ff0aa0e 100644 --- a/lightllm/server/router/model_infer/infer_req_ext1.py +++ b/lightllm/server/router/model_infer/infer_req_ext1.py @@ -35,13 +35,14 @@ def append_token_indexes(self, mem_indexes: torch.Tensor): 非空输入必须是 CPU Tensor；CPU 缓存模式直接复制到自己的 pinned 缓存。 pure_gpu 按需 pin 后异步上传；已 pinned 的输入在 H2D 完成前不得改写，但可释放引用。 + pin_mem_page_cache 仅接受从页边界追加完整页，调用方保证每页内物理 token 索引连续。 前缀替换先 reset，再按顺序追加重建。 """ pass @abstractmethod def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: - """返回独立 CPU 快照，允许缓存或传输任务跨请求生命周期持有。""" + """读取有效长度内任意 token 区间，返回可跨请求生命周期持有的独立 CPU 快照。""" pass @abstractmethod @@ -97,6 +98,7 @@ def append_token_indexes(self, mem_indexes: torch.Tensor): self.token_index_len = end def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + assert end <= self.token_index_len, "Read range exceeds recorded token indexes" if start == end: return torch.empty(0, dtype=self.gpu_indexes.dtype, device="cpu", requires_grad=False) return self._cpu_buffer[start:end].clone() @@ -173,13 +175,15 @@ def __init__(self, infer_req: "InferReq", index_stream: torch.cuda.Stream): ) def append_token_indexes(self, mem_indexes: torch.Tensor): - """从页边界追加 token 索引；调用方保证同一逻辑页内索引连续。""" + """从页边界追加完整页；调用方保证同一逻辑页内物理 token 索引连续。""" append_len = mem_indexes.numel() if append_len == 0: return assert mem_indexes.device.type == "cpu", "mem_indexes must be a CPU Tensor" start = self.token_index_len end = start + append_len + assert start % self.page_size == 0, "Append start must be page-aligned" + assert end % self.page_size == 0, "Append end must be page-aligned" self._reserve(end) # 输入从页首开始，每隔 page_size 个 token 保存一个页首索引。 first_page = start // self.page_size @@ -188,6 +192,7 @@ def append_token_indexes(self, mem_indexes: torch.Tensor): self.token_index_len = end def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + assert end <= self.token_index_len, "Read range exceeds recorded token indexes" return self._expand_token_indexes(start, end, pin_memory=False) def get_upload_indexes(self) -> torch.Tensor: @@ -268,6 +273,7 @@ def append_token_indexes(self, mem_indexes: torch.Tensor): self.gpu_index_submitted_len = end def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: + assert end <= self.token_index_len, "Read range exceeds recorded token indexes" # 阻塞到 D2H 完成，GPU 索引拷回后得到独立 CPU 快照。 with torch.cuda.stream(self._index_stream): return self.gpu_indexes[start:end].detach().cpu() diff --git a/test/benchmark/benchmark_pinned_allocation_rounding.py b/test/benchmark/benchmark_pinned_allocation_rounding.py new file mode 100644 index 0000000000..d77bb032e6 --- /dev/null +++ b/test/benchmark/benchmark_pinned_allocation_rounding.py @@ -0,0 +1,127 @@ +"""Verify pinned allocation rounding and cache reuse in a standalone CUDA process. + +Example: + CUDA_VISIBLE_DEVICES=4 python test/benchmark/benchmark_pinned_allocation_rounding.py \ + --output /tmp/pinned_allocation_rounding.json +""" + +import argparse +import json +import os +from pathlib import Path +import weakref + +import torch + + +MIB = 1024 * 1024 +STAT_KEYS = ("allocated_bytes.current", "active_bytes.current", "num_host_alloc", "num_host_free") + + +def snapshot(): + stats = torch.cuda.memory.host_memory_stats() + return {key: stats[key] for key in STAT_KEYS} + + +def measure(label, numel, dtype): + # 仅清理本独立进程的 host 缓存，确保首次申请无法命中已有尺寸桶。 + torch._C._host_emptyCache() + before = snapshot() + tensor = torch.empty(numel, dtype=dtype, device="cpu", pin_memory=True, requires_grad=False) + allocated = snapshot() + requested_bytes = tensor.nbytes + expected_block_bytes = 1 << (requested_bytes - 1).bit_length() + block_bytes = allocated["allocated_bytes.current"] - before["allocated_bytes.current"] + pointer = tensor.data_ptr() + storage_bytes = tensor.untyped_storage().nbytes() + assert tensor.is_pinned() + assert block_bytes == expected_block_bytes, (label, block_bytes, expected_block_bytes) + assert allocated["active_bytes.current"] - before["active_bytes.current"] == block_bytes + assert allocated["num_host_alloc"] - before["num_host_alloc"] == 1 + assert storage_bytes == requested_bytes + + # 没有 H2D 或共享视图；删除最后一个引用后，整块内存应进入空闲缓存。 + tensor_ref = weakref.ref(tensor) + del tensor + assert tensor_ref() is None + released = snapshot() + # 此安装版本的无事件 free 路径未扣减 active 计数；用地址复用和实际分配次数验证回收。 + assert released["allocated_bytes.current"] == allocated["allocated_bytes.current"] + + # 按桶的完整大小再申请，验证与原申请复用同一块内存，不再调用底层申请。 + reused = torch.empty(block_bytes, dtype=torch.uint8, device="cpu", pin_memory=True, requires_grad=False) + reuse_stats = snapshot() + same_pointer = reused.data_ptr() == pointer + assert same_pointer + assert reuse_stats["num_host_alloc"] == allocated["num_host_alloc"] + del reused + torch._C._host_emptyCache() + cleared = snapshot() + assert cleared["allocated_bytes.current"] == before["allocated_bytes.current"] + assert cleared["num_host_free"] - before["num_host_free"] == 1 + return { + "case": label, + "dtype": str(dtype), + "numel": numel, + "tensor_bytes": requested_bytes, + "storage_bytes": storage_bytes, + "allocator_block_bytes": block_bytes, + "tensor_mib": requested_bytes / MIB, + "allocator_block_mib": block_bytes / MIB, + "reuse_request_bytes": block_bytes, + "reuse_same_pointer": same_pointer, + "reuse_new_host_allocations": reuse_stats["num_host_alloc"] - allocated["num_host_alloc"], + "stats": { + "before": before, + "allocated": allocated, + "released": released, + "reused": reuse_stats, + "cache_cleared": cleared, + }, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + # reserve segment 会改变缓存清理行为，本原型验证默认 host allocator 的独立分配路径。 + for name in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF"): + assert not os.environ.get(name), f"Run this standalone prototype with {name} unset" + torch.cuda.init() + warmup = torch.empty(1, dtype=torch.uint8, pin_memory=True) + del warmup + cases = [ + ("64 MiB uint8", 64 * MIB, torch.uint8), + ("64 MiB + 1 byte uint8", 64 * MIB + 1, torch.uint8), + ("65 MiB uint8", 65 * MIB, torch.uint8), + ("128 MiB uint8", 128 * MIB, torch.uint8), + ("65 million bytes (decimal)", 65_000_000, torch.uint8), + ("65M int32 elements (binary M)", 65 * MIB, torch.int32), + ] + results = [measure(*case) for case in cases] + output = { + "torch": torch.__version__, + "torch_commit": torch.version.git_version, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "allocator_environment": { + name: os.environ.get(name) for name in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF") + }, + "active_counter_note": "This build does not decrement active counters on event-free release; " + "reuse is verified by pointer equality and unchanged num_host_alloc.", + "results": results, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(output, indent=2) + "\n") + for result in results: + print( + f"{result['case']}: tensor={result['tensor_mib']:.6f} MiB, " + f"allocator={result['allocator_block_mib']:.0f} MiB, " + f"reuse_same_pointer={result['reuse_same_pointer']}, " + f"reuse_new_host_allocations={result['reuse_new_host_allocations']}" + ) + print(f"All assertions passed. Results: {args.output}") + + +if __name__ == "__main__": + main() diff --git a/unit_tests/common/test_req_token_index.py b/unit_tests/common/test_req_token_index.py index 71575669f4..95d1247e73 100644 --- a/unit_tests/common/test_req_token_index.py +++ b/unit_tests/common/test_req_token_index.py @@ -186,18 +186,38 @@ def test_append_rejects_gpu_input_without_changing_index_state(mode): @requires_cuda @pytest.mark.parametrize("mode", MODES) -def test_gpu_getter_rejects_unrecorded_indexes_without_advancing_progress(mode): - table = torch.full((1, 8), -7, dtype=torch.int32, device="cuda") - helper = create_req_token_index(make_req(table, mode), index_stream=torch.cuda.current_stream()) - helper.append_token_indexes(torch.tensor([10, 11], dtype=torch.int32, pin_memory=True)) +@pytest.mark.parametrize("getter_name", ["get_cpu_token_indexes", "get_gpu_token_indexes", "get_free_token_indexes"]) +@pytest.mark.parametrize("start,end", [(0, 65), (65, 65)]) +def test_read_getters_reject_unrecorded_indexes_without_advancing_progress(mode, getter_name, start, end): + table = torch.full((1, 128), -7, dtype=torch.int32, device="cuda") + helper = create_req_token_index(make_req(table, mode, page_size=64), index_stream=torch.cuda.current_stream()) + helper.append_token_indexes(torch.arange(64, dtype=torch.int32, pin_memory=True)) submitted_len = helper.gpu_index_submitted_len + getter = getattr(helper, getter_name) with pytest.raises(AssertionError, match="Read range exceeds"): - helper.get_gpu_token_indexes(0, 4) + getter(start, end) - assert helper.token_index_len == 2 + assert helper.token_index_len == 64 assert helper.gpu_index_submitted_len == submitted_len - assert helper.get_gpu_token_indexes(0, 2).cpu().tolist() == [10, 11] + assert getter(1, 3).cpu().tolist() == [1, 2] + assert getter(64, 64).numel() == 0 + + +@pytest.mark.parametrize("start,append_len,boundary", [(0, 63, "end"), (1, 63, "start")]) +def test_page_cache_rejects_partial_page_appends_before_mutating_state(start, append_len, boundary): + table = torch.full((1, 128), -7, dtype=torch.int32) + helper = create_req_token_index(make_req(table, "pin_mem_page_cache", page_size=64), index_stream=object()) + # 模拟不对齐的已有长度，验证从页内追加也会在分配或写入之前失败。 + helper.token_index_len = start + + with pytest.raises(AssertionError, match=f"Append {boundary} must be page-aligned"): + helper.append_token_indexes(torch.arange(append_len, dtype=torch.int32)) + + assert helper.token_index_len == start + assert helper.gpu_index_submitted_len == 0 + assert helper._cpu_buffer is None + assert table.eq(-7).all() @requires_cuda 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 index fe8e3547a0..105e59231a 100644 --- 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 @@ -46,6 +46,7 @@ def test_copy_master_kv_updates_slave_hold_kv_len(): master_req = SimpleNamespace( req_idx=0, cur_kv_len=4, + hold_kv_len=4, cur_output_len=1, shared_kv_node=new_shared_node, ) From d91d3cda27038be53273f27949c4ea81b5110e94 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 08:09:22 +0000 Subject: [PATCH 17/22] Fix decode capacity checks before optional token reservation Admit requests using required KV capacity before assigning optional decode headroom from the remaining batch budget. Keep one allocation per decode request and cover capacity boundaries, request ordering, and PD cleanup behavior. --- .../model_infer/mode_backend/base_backend.py | 20 +++- unit_tests/common/test_req_manager_page.py | 97 +++++++++++++++++++ 2 files changed, 112 insertions(+), 5 deletions(-) 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 aa1a11094e..15b306c444 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -747,6 +747,7 @@ def _get_classed_reqs( finished_reqs = [] prefill_reqs = [] decode_reqs = [] + decode_allocations = [] # 单轮最多处理少量因 token 容量不足而无法继续的请求，避免一次性影响大量请求。 # 普通 Decode 请求进入暂停队列等待恢复；PD Decode 请求则强制提前结束并进入清理流程。 @@ -789,12 +790,10 @@ def _get_classed_reqs( if is_decode: # 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) + # 只扣除必需容量；等本轮分类结束后，再从剩余预算中分配可选预留。 + if alloc_token_num > 0: + decode_allocations.append((req_obj, alloc_token_num)) decode_reqs.append(req_obj) can_alloc_token_num -= alloc_token_num else: @@ -840,6 +839,17 @@ def _get_classed_reqs( req_obj.wait_pause = True wait_pause_count += 1 + # 所有接纳请求的必需容量均已计入预算，额外预留不能挤占后续 decode 或 prefill。 + # page_size 较小时，decode 频繁分配 mem index 会增加 CPU overhead；因此在剩余预算允许时， + # 为 decode 请求额外预留一些 token，供后续步骤复用，减少分配次数并降低 CPU overhead。 + for req_obj, alloc_token_num in decode_allocations: + if self.args.page_size < 8: + extra_token_num = min(8, can_alloc_token_num) // self.args.page_size * self.args.page_size + alloc_token_num += extra_token_num + can_alloc_token_num -= extra_token_num + # 将必需容量和可选预留合并申请，保留每个请求一次 alloc 的行为。 + self._alloc_req_kv_mem(req_obj, alloc_token_num) + # 先由控制器确定请求需要写入的缓存层级，再按是否包含 CPU cache 决定是否发起 offload。 cache_controller = g_infer_context.cache_placement_controller new_finished_reqs = [req for req in finished_reqs if req.cpu_cache_task_status.is_not_started()] diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index 1dcf637ff8..b3a92969c2 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -259,6 +259,103 @@ def test_decode_scheduler_reserves_cpu_request_table(monkeypatch): assert context.req_manager.mem_manager.alloc_sizes == [12] +def _make_capacity_scheduler(monkeypatch, page_size, available, run_mode, support_overlap=False): + context, backend = _make_context(monkeypatch) + context.args.page_size = context.req_manager.page_size = page_size + context.req_manager.token_indexes_cpu = torch.full((2, 256), -1, dtype=torch.int32) + context.req_manager.req_to_token_indexs = torch.full((2, 256), -1, dtype=torch.int32) + backend.args.enable_cpu_cache = False + backend.args.enable_prefill_decode_mixed = False + backend.args.run_mode = run_mode + backend.support_overlap = support_overlap + backend.disable_chunked_prefill = True + backend.batch_max_tokens = 256 + backend.is_master_in_dp = True + backend.logger = SimpleNamespace(info=lambda message: None) + backend._timer_merge_radix_tree = lambda: None + backend._reorder_pd_high_priority_reqs = lambda reqs: reqs + backend._reorder_long_prefill_reqs = lambda reqs: reqs + context.get_can_alloc_token_num = lambda: available + context.cache_placement_controller = SimpleNamespace(set_req_cache_way=lambda reqs: None) + context.finished_reqs = [] + context.filter_reqs = lambda finished_reqs: context.finished_reqs.extend(finished_reqs) + context.pause_reqs = lambda reqs, is_master_in_dp: None + return context, backend + + +def _make_capacity_req(req_idx, context, is_decode=True): + req = _make_req(req_idx) + req.req_id = req_idx + req.args = context.args + req.cur_kv_len = req.hold_kv_len = context.args.page_size + req.mtp_step = 0 + req.filter_mark = req.wait_pause = req.paused = req.infer_aborted = False + req.finish_status = infer_batch.FinishStatus() + req.cpu_cache_task_status = SimpleNamespace(is_not_started=lambda: True) + req.is_slave_req = lambda: False + req.get_cur_total_len = lambda: req.cur_kv_len + (1 if is_decode else 2 * context.args.page_size) + req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) + req.prefill_need_token_num = MethodType(InferReq.prefill_need_token_num, req) + context.req_manager.append_token_indexes(req_idx, torch.arange(req.hold_kv_len, dtype=torch.int32)) + return req + + +@pytest.mark.parametrize("page_size", [1, 3, 4, 7, 8, 64]) +@pytest.mark.parametrize("run_mode", ["normal", "decode"]) +@pytest.mark.parametrize("extra_budget", [0, 1, 8]) +@pytest.mark.parametrize("request_types", [(True,), (True, True), (True, False), (False, True)]) +def test_decode_headroom_only_uses_budget_after_all_required_allocations( + monkeypatch, page_size, run_mode, extra_budget, request_types +): + required = sum(page_size if is_decode else 2 * page_size for is_decode in request_types) + context, backend = _make_capacity_scheduler(monkeypatch, page_size, required + extra_budget, run_mode) + reqs = [_make_capacity_req(index, context, is_decode) for index, is_decode in enumerate(request_types)] + context.req_manager.mem_manager.next_index = len(reqs) * page_size + backend._filter_not_ready_reqs = lambda req_ids: reqs + + prefill, decode = backend._get_classed_reqs(req_ids=list(range(len(reqs)))) + + assert prefill == [req for req, is_decode in zip(reqs, request_types) if not is_decode] + assert decode == [req for req, is_decode in zip(reqs, request_types) if is_decode] + assert context.finished_reqs == [] + remaining = extra_budget + for req, is_decode in zip(reqs, request_types): + extra = min(8, remaining) // page_size * page_size if is_decode and page_size < 8 else 0 + remaining -= extra + assert req.hold_kv_len == page_size + (page_size if is_decode else 2 * page_size) + extra + assert req.token_index.token_index_len == req.hold_kv_len + assert req.token_index.gpu_index_submitted_len == req.hold_kv_len + assert not req.wait_pause and not req.filter_mark + assert not hasattr(req, "finished_by_pd_decode_capacity") + # 必需量和可选预留合并申请，不额外增加分配次数，且不会超出总预算。 + allocations = context.req_manager.mem_manager.alloc_sizes + assert len(allocations) == len(reqs) + assert sum(allocations) == required + extra_budget - remaining + + +@pytest.mark.parametrize("page_size", [1, 4, 64]) +@pytest.mark.parametrize("run_mode", ["normal", "decode"]) +@pytest.mark.parametrize("support_overlap", [False, True]) +def test_decode_still_rejects_requests_when_required_capacity_is_insufficient( + monkeypatch, page_size, run_mode, support_overlap +): + context, backend = _make_capacity_scheduler(monkeypatch, page_size, page_size - 1, run_mode, support_overlap) + req = _make_capacity_req(0, context) + backend._filter_not_ready_reqs = lambda req_ids: [req] + + assert backend._get_classed_reqs(req_ids=[0]) == ([], []) + + assert context.req_manager.mem_manager.alloc_sizes == [] + assert req.hold_kv_len == page_size + if run_mode == "normal": + assert req.wait_pause + assert context.finished_reqs == [] + else: + assert req.finished_by_pd_decode_capacity + assert req.filter_mark is support_overlap + assert context.finished_reqs == ([] if support_overlap else [req]) + + @pytest.mark.parametrize( "page_size, expected_alloc_token_num", [(1, 9), (3, 9), (4, 12), (7, 14), (8, 8)], From f281a1af4a07a2f2cd102bb580a3e0b8e9f61fb8 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 11:03:27 +0000 Subject: [PATCH 18/22] Refactor inference preparation and hybrid cache release --- docs/kv_cache_page_size.md | 20 +- lightllm/common/basemodel/basemodel.py | 172 +++---- .../server/router/model_infer/infer_batch.py | 142 ++--- .../router/model_infer/infer_req_ext1.py | 75 ++- .../model_infer/mode_backend/base_backend.py | 39 +- .../common/basemodel/test_model_output.py | 58 ++- .../common/test_infer_batch_cache_release.py | 486 ++++++++++++++++++ unit_tests/common/test_req_manager_page.py | 109 +++- unit_tests/common/test_req_token_index.py | 35 +- 9 files changed, 881 insertions(+), 255 deletions(-) create mode 100644 unit_tests/common/test_infer_batch_cache_release.py diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index 80387a9fa3..e6c6ea49c4 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -38,8 +38,8 @@ KV 存储和请求表，但资源的申请、缓存和释放以完整物理页 `ReqManager` 只持有 GPU `req_to_token_indexs`，不再分配全局 CPU 页表。 每个 `InferReq` 通过成员 `token_index` 持有一个 `ReqTokenIndex`，实现位于 `server/router/model_infer/infer_req_ext1.py`。 索引操作直接使用 `req.token_index.append_token_indexes(...)` 等接口，`InferReq` 不再提供同名转发方法。 -基类仅在构造函数中初始化公共对象和状态，其余方法只定义抽象接口； -各子类实现或复用读写、上传和重置逻辑。 +基类初始化公共对象和状态，并统一维护上传提交协议及索引流等待； +各子类实现或复用存储、读写和重置逻辑。 扩展构造函数接收 `infer_req` 和必填的 `index_stream`，仅提取并保存固定的 `req_idx`、GPU 请求表行视图及索引流， 避免反向持有请求对象形成循环引用；索引流由外部选择，Helper 内部不查询默认流或捕获构造时的当前流。 CPU 缓存实现同时记录输入长度 `input_len`，helper 在共享请求初始化后、前缀匹配前构造。 @@ -114,17 +114,19 @@ PyTorch `copy_(non_blocking=True)` 会登记 pinned 源的流使用，旧存储 不改变提交进度；无待上传内容时返回空 tensor， 纯 GPU 模式返回 `None`。CPU 缓存模式的 GPU getter 校验 `end <= token_index_len` 后按需提交完整待上传区间， 纯 GPU 模式校验 `end <= gpu_index_submitted_len`；读取均不允许推进超出已有索引的提交进度。 -批量入口遇到存在未提交区间、但没有 CPU 上传源的情况会直接断言，不再替 GPU 模式补写提交进度。 +批量入口不会替 GPU 模式补写提交进度；若缺少待上传数据，由后续拷贝操作报错，提交进度保持不变。 例如 CPU 缓存模式追加 64 个索引后，两个字段分别为 `64/0`；提交到 GPU 后为 `64/64`； 再追加 32 个索引后为 `96/64`，下次仅需提交 `[64,96)`。需要替换映射时，`reset()` 将两者清为 `0/0`， 重新追加 96 个完整索引后为 `96/0`，再从头提交 GPU 更新。纯 GPU 模式在追加接口中 直接提交写表，所以追加后两者相同。 -`_get_classed_reqs` 返回前，backend 收集各请求待上传的 pinned 视图，直接使用 -`cudaMemcpyBatchAsync` 提交，不再申请聚合缓冲区或额外复制 CPU 索引。 -同一批请求必须共用一个索引流，backend 对此进行断言；随后只准备源、目标视图列表， -在索引流上下文调用批量拷贝并更新提交进度。返回前仅在调用流不同于索引流时排入等待， -没有新上传时也承接已有写入的依赖。backend 不再持有上传流，也不逐条交接多条索引流。 +`_get_classed_reqs` 返回前，backend 将本批索引对象交给 `ReqTokenIndex.submit_pending_uploads`， +并传入批量拷贝函数；backend 不再读写索引流、GPU 行和提交进度。 +统一入口检查所有请求共用一个索引流，准备源、目标视图列表，在索引流提交拷贝， +仅在整批提交成功后更新提交进度。单请求 GPU getter 复用该协议，默认通过 PyTorch `copy_` 提交； +backend 路径使用 `cudaMemcpyBatchAsync`，不申请聚合缓冲区或额外复制 CPU 索引。 +返回前仅在调用流不同于索引流时排入等待，没有新上传时也承接已有写入的依赖。 +backend 不再持有上传流，也不逐条交接多条索引流。 原生 API 不接受 legacy 默认流，适配由 `PinnedH2DBatchCopier` 内部处理：按设备延迟创建并复用非默认上传流， 先让上传流等待默认流，提交后再让默认流等待上传流；非默认索引流直接提交，PyTorch 回退直接使用当前流。 这些操作只排入 GPU 依赖，不阻塞 CPU。 @@ -179,6 +181,8 @@ HOLD 请求直接在 GPU 表初始化为预留物理页的循环索引。 - Decode：请求通过调度后，若尾页仍有预留槽位则不再访问 allocator；容量不足时先补齐完整页再执行。 - Prefix Cache 命中：命中长度天然是整页，`cur_kv_len` 与 `hold_kv_len` 同时初始化为命中长度。 - 完成/暂停：完整逻辑页可进入 Radix Cache；重复页和未完成尾页按物理页展开后释放。 + 混合注意力在 Prefill 尚未完成时，也可插入已保存大页 checkpoint 且 KV 已计算完成的最长前缀； + 缓存长度不超过目标缓存边界，其余 KV 统一释放；大页 checkpoint 随前缀移交缓存，重复 checkpoint 在插入时回收。 - Attention：FA3/FlashInfer 页表每项由物理 token 页首索引除以 `page_size` 得到，KV buffer 视为 `[page_count, page_size, ...]`。 diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index c89a88d90c..2d48237b2b 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -400,6 +400,57 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) return infer_state + def _prepare_input_ids(self, model_input: ModelInput): + """普通与 overlap 入口统一准备真实 token，随后再由状态准备流程 padding。""" + if model_input.is_prefill: + 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, + ) + elif 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；padding 随后填入 dummy token。 + model_input.input_ids = torch.empty((0,), dtype=torch.int64, device=model_input.b_req_idx.device) + + def _create_prefill_inferstate( + self, + model_input: ModelInput, + handle_token_num: int, + microbatch_index: int = 0, + mem_indexes_ready_event=None, + ): + """统一 prefill padding 和状态初始化，允许在 KV 索引就绪处记录事件。""" + model_input = self._create_padded_prefill_model_input(model_input, handle_token_num) + infer_state = self._create_inferstate(model_input, microbatch_index) + if mem_indexes_ready_event is not None: + mem_indexes_ready_event.record() + 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 + ): + """统一 decode padding、graph 标记和初始化顺序。""" + model_input = self._create_padded_decode_model_input(model_input, batch_size) + infer_state = self._create_inferstate(model_input, microbatch_index) + # Attention backend 在初始化时需要知道是否正在 capture。 + infer_state.is_cuda_graph = 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 @@ -541,15 +592,7 @@ 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, - ) + self._prepare_input_ids(model_input) origin_handle_token_num = model_input.input_ids.shape[0] origin_batch_size = model_input.batch_size @@ -565,16 +608,10 @@ 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, mem_indexes_ready_event=prefill_mem_indexes_ready_event + ) model_output = self._context_forward(infer_state=infer_state) model_output = self._create_unpad_prefill_model_output( @@ -589,21 +626,7 @@ 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, - ) + self._prepare_input_ids(model_input) origin_batch_size = model_input.batch_size # 空 DP rank 也需要完整的 dummy verify group；TP/SP 切分不能破坏 MTP 分组。 @@ -621,13 +644,7 @@ 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 会根据该标记准备 CUDA Graph capture 专用状态， - # 因此必须在 init_att_state 之前完成赋值。 - infer_state.is_cuda_graph = 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) if use_cuda_graph: if need_capture: @@ -739,15 +756,7 @@ 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, - ) + self._prepare_input_ids(model_input) return self._microbatch_overlap_prefill_cuda(model_input0, model_input1) def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input1: ModelInput): @@ -765,20 +774,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() + 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) prefill_mem_indexes_ready_event = torch.cuda.Event() prefill_mem_indexes_ready_event.record() @@ -808,19 +805,7 @@ 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, - ) + self._prepare_input_ids(model_input) return self._microbatch_overlap_decode_cuda(model_input0, model_input1) def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1: ModelInput): @@ -831,21 +816,16 @@ 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 = 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 = 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) + infer_state1 = self._create_decode_inferstate(model_input1, infer_batch_size, need_capture, 1) + + if use_cuda_graph: if need_capture: model_output0, model_output1 = self.graph.capture_decode( self._overlap_tpsp_token_forward, @@ -858,23 +838,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/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 5d1f377245..522477c93d 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -168,32 +168,23 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer def _full_att_free_req(self, free_token_index: 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) 待释放区间，再处理可缓存前缀的插入和去重。 + # 只缓存完整页，不足一页的有效 KV 和预留尾部均归还 allocator。 + # 待释放索引可以借用 GPU 行；缓存值须独立于请求槽位及其后续复用。 free_token_index.append(req.token_index.get_free_token_indexes(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") - # radix cache 跨请求槽位复用持有索引，因此使用独立 CPU 快照。 + key = torch.tensor(req.get_input_token_ids()[:cache_kv_len], dtype=torch.int64, device="cpu") value = req.token_index.get_cpu_token_indexes(0, cache_kv_len) - prefix_len, _ = self.radix_cache.insert(key, value) - old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len - free_token_index.append(req.token_index.get_free_token_indexes(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 + self._release_duplicate_prefix(free_token_index, req, prefix_len) def _hybrid_att_free_req(self, free_token_index: 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 big_page_num = args.linear_att_page_block_num + big_page_token_num = hash_page_size * big_page_num shared_kv_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len - tail_big_page_token_num = ( - req.hybrid_cache_len // (hash_page_size * big_page_num) * (hash_page_size * big_page_num) - ) + completed_kv_len = min(req.cur_kv_len, req.hybrid_cache_len) + completed_big_page_token_num = completed_kv_len // big_page_token_num * big_page_token_num page_num = req.hybrid_cache_len // hash_page_size assert req.hybrid_cache_len >= shared_kv_len if req.tail_small_page_buffer_id is not None: @@ -205,79 +196,58 @@ 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(req.token_index.get_free_token_indexes(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 = req.token_index.get_cpu_token_indexes(0, req.cur_kv_len) - 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(req.token_index.get_free_token_indexes(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(req.token_index.get_free_token_indexes(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 completed_big_page_token_num > shared_kv_len: + # Prefill 在大页边界保存 checkpoint，因此按已完成长度对齐即可。 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 = req.token_index.get_cpu_token_indexes(0, req.cur_kv_len) - 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(req.token_index.get_free_token_indexes(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(req.token_index.get_free_token_indexes(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 + assert len(req.hybrid_len_to_big_page_id) > 0 + cache_kv_len = completed_big_page_token_num + elif shared_kv_len <= req.cur_kv_len: + # 没有新 checkpoint，区分共享小页、共享大页和空前缀。 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 - return + assert len(req.hybrid_len_to_big_page_id) == 0 + if req.shared_kv_node is not None and not req.shared_kv_node.is_big_page_node(): + # 大页关闭时，前缀匹配会直接保留共享小页。请求只持有节点引用， + # 没有可移交给 insert 的新小页 checkpoint；节点引用只保护 KV， + # 其 CPU page buffer 仍可能被独立淘汰，因此不能依赖它仍然存在。 + # 即使 buffer 仍在，也不能借用其 id 重新插入，否则重复插入会将 + # 缓存自身持有的槽位当作重复 checkpoint 释放。 + # 这里只回收请求私有的尾部 KV 并解除引用，共享 KV 和 checkpoint + # 继续由 radix cache 管理，避免改变原有的小页插入协议。 + self._free_req_mem_without_radix_insert(free_token_index, req) + return + # 共享大页可在没有新 checkpoint 时重复插入；空前缀也可直接插入。 + cache_kv_len = shared_kv_len + else: + assert False, f"error state: cur_kv_len: {req.cur_kv_len}" - assert False, f"error state: cur_kv_len: {req.cur_kv_len}" + free_token_index.append(req.token_index.get_free_token_indexes(cache_kv_len, req.hold_kv_len)) + req.cur_kv_len = cache_kv_len + key = torch.tensor(req.get_input_token_ids()[:cache_kv_len], dtype=torch.int64, device="cpu") + value = req.token_index.get_cpu_token_indexes(0, cache_kv_len) + page_num = cache_kv_len // hash_page_size + state_idxs = [None] * page_num + if req.tail_small_page_buffer_id is not None: + 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=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, + ) + assert len(req.hybrid_len_to_big_page_id) == 0 + self._release_duplicate_prefix(free_token_index, req, prefix_len) + + def _release_duplicate_prefix(self, free_token_index: List, req: "InferReq", prefix_len: int): + """回收插入时发现的重复 KV 页，并解除请求对共享节点的引用。""" + old_prefix_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len + free_token_index.append(req.token_index.get_free_token_indexes(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 def _save_promptcache_kvbuffer(self): """ diff --git a/lightllm/server/router/model_infer/infer_req_ext1.py b/lightllm/server/router/model_infer/infer_req_ext1.py index bb4ff0aa0e..2dcc694b25 100644 --- a/lightllm/server/router/model_infer/infer_req_ext1.py +++ b/lightllm/server/router/model_infer/infer_req_ext1.py @@ -1,7 +1,7 @@ """InferReq 的 token 索引扩展，仅保存请求槽位和索引状态。""" from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, List, Optional, Sequence import torch @@ -76,6 +76,52 @@ def reset(self): """重置有效索引长度和 GPU 提交进度，释放 CPU 缓存引用。""" pass + @staticmethod + def submit_pending_uploads( + token_indexes: Sequence["ReqTokenIndex"], + copy_batch: Optional[Callable[[List[torch.Tensor], List[torch.Tensor]], None]] = None, + ) -> None: + """统一准备待上传区间、提交拷贝并确认进度，随后让调用流承接依赖。 + + copy_batch 接收 GPU 目标和 pinned CPU 源列表，必须在当前索引流提交拷贝， + 并保护异步源的生命周期；未提供时使用 PyTorch copy_。 + 只有整批提交成功才推进进度；提交不表示传输完成，无新上传也需等待已有写入。 + """ + if not token_indexes: + return + index_stream = token_indexes[0]._index_stream + assert all( + token_index._index_stream == index_stream for token_index in token_indexes + ), "All requests in a batch must use the same index stream" + cpu_slices, gpu_slices, pending_uploads = [], [], [] + for token_index in token_indexes: + start, end = token_index.gpu_index_submitted_len, token_index.token_index_len + if start < end: + source = token_index.get_upload_indexes() + cpu_slices.append(source) + gpu_slices.append(token_index.gpu_indexes[start:end]) + pending_uploads.append((token_index, end)) + if pending_uploads: + with torch.cuda.stream(index_stream): + if copy_batch is None: + for destination, source in zip(gpu_slices, cpu_slices): + destination.copy_(source, non_blocking=True) + else: + copy_batch(gpu_slices, cpu_slices) + for token_index, end in pending_uploads: + token_index._mark_gpu_indexes_submitted(end) + token_indexes[0]._wait_for_index_stream() + + def _mark_gpu_indexes_submitted(self, end: int) -> None: + """只确认已成功提交的连续前缀，不把 GPU 完成状态混入长度计数。""" + assert self.gpu_index_submitted_len <= end <= self.token_index_len + self.gpu_index_submitted_len = end + + def _wait_for_index_stream(self) -> None: + current_stream = torch.cuda.current_stream(self._index_stream.device) + if current_stream != self._index_stream: + current_stream.wait_stream(self._index_stream) + class PinMemCacheReqTokenIndex(ReqTokenIndex): """按请求按需申请完整 token 级 pinned CPU 缓存，GPU 更新可批量提交。""" @@ -106,19 +152,10 @@ def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: """在索引流提交全部待上传索引，让调用流承接依赖后返回 GPU 视图。""" assert end <= self.token_index_len, "Read range exceeds recorded token indexes" - current_stream = torch.cuda.current_stream(self._index_stream.device) - with torch.cuda.stream(self._index_stream): - if self.gpu_index_submitted_len < end: - source_indexes = self.get_upload_indexes() - upload_end = self.token_index_len - # 缓存只追加，reset 会丢弃旧 buffer，已有上传区间不会被 CPU 重写。 - # copy_ 向 PyTorch pinned allocator 登记流使用，即使 reset 释放引用， - # 旧内存也会等 H2D 完成后才允许回收复用，因此可以直接上传缓存视图。 - self.gpu_indexes[self.gpu_index_submitted_len : upload_end].copy_(source_indexes, non_blocking=True) - self.gpu_index_submitted_len = upload_end - # 没有新上传时，已有写入也可能尚未完成；调用流仍需等待索引流。 - if current_stream != self._index_stream: - current_stream.wait_stream(self._index_stream) + if self.gpu_index_submitted_len < end: + self.submit_pending_uploads([self]) + else: + self._wait_for_index_stream() return self.gpu_indexes[start:end] def get_upload_indexes(self) -> torch.Tensor: @@ -264,13 +301,11 @@ def append_token_indexes(self, mem_indexes: torch.Tensor): end = start + append_len # copy_ 会向 PyTorch pinned allocator 登记流使用，H2D 完成前不会回收源内存。 # 普通 CPU 输入已形成独立 pinned 副本；直接复用的 pinned 输入在传输完成前不能改写。 - current_stream = torch.cuda.current_stream(self._index_stream.device) with torch.cuda.stream(self._index_stream): self.gpu_indexes[start:end].copy_(mem_indexes, non_blocking=True) - if current_stream != self._index_stream: - current_stream.wait_stream(self._index_stream) self.token_index_len = end - self.gpu_index_submitted_len = end + self._mark_gpu_indexes_submitted(end) + self._wait_for_index_stream() def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: assert end <= self.token_index_len, "Read range exceeds recorded token indexes" @@ -281,9 +316,7 @@ def get_cpu_token_indexes(self, start: int, end: int) -> torch.Tensor: def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: # append 已在索引流提交写入，读取只校验范围，不推进提交进度。 assert end <= self.gpu_index_submitted_len, "Read range exceeds submitted GPU indexes" - current_stream = torch.cuda.current_stream(self._index_stream.device) - if current_stream != self._index_stream: - current_stream.wait_stream(self._index_stream) + self._wait_for_index_stream() return self.gpu_indexes[start:end] def get_upload_indexes(self) -> None: diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 15b306c444..b2f9e5ae8d 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -11,6 +11,7 @@ from lightllm.utils.log_utils import init_logger from lightllm.models import get_draft_model_class, get_model from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack +from lightllm.server.router.model_infer.infer_req_ext1 import ReqTokenIndex from lightllm.server.router.token_load import TokenLoad from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async @@ -663,40 +664,10 @@ def _alloc_req_kv_mem( return mem_indexes def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): - """在本批请求共用的索引流上批量上传，调用流按需等待索引流。""" - # 第一阶段：空批直接返回，确认请求共用一个索引流，并记录当前调用流。 - if not run_reqs: - return - index_streams = {req.token_index._index_stream for req in run_reqs} - assert len(index_streams) == 1, "All requests in a batch must use the same index stream" - index_stream = index_streams.pop() - current_stream = torch.cuda.current_stream(index_stream.device) - # 第二阶段：准备尚未提交的 CPU token 索引和 GPU 目标视图，记录本次提交终点。 - cpu_index_slices, gpu_index_slices = [], [] - pending_index_uploads = [] - for req in run_reqs: - token_index = req.token_index - start, end = token_index.gpu_index_submitted_len, token_index.token_index_len - if start < end: - source_indexes = token_index.get_upload_indexes() - assert source_indexes is not None, ( - "GPU index submission progress does not match recorded length; " - "pure_gpu must submit indexes when appending" - ) - cpu_index_slices.append(source_indexes) - gpu_index_slices.append(token_index.gpu_indexes[start:end]) - pending_index_uploads.append((token_index, end)) - # 第三阶段：在索引流上批量提交 H2D，默认流适配和 pinned 源保活由拷贝封装处理。 - if cpu_index_slices: - with torch.cuda.stream(index_stream): - cuda_memcpy_pinned_h2d_batch_async(gpu_index_slices, cpu_index_slices) - # 第四阶段：成功提交后再更新进度；提交完成不代表 GPU 传输已经完成。 - for token_index, end in pending_index_uploads: - token_index.gpu_index_submitted_len = end - # 第五阶段：调用流与索引流不同时建立等待依赖，保证后续 GPU 访问能读到更新。 - # 各模式在没有新上传时，已提交的索引写入也可能尚未完成，仍需承接依赖。 - if current_stream != index_stream: - current_stream.wait_stream(index_stream) + """选择本批请求；索引对象统一管理上传区间、提交进度和流依赖。""" + ReqTokenIndex.submit_pending_uploads( + [req.token_index for req in run_reqs], copy_batch=cuda_memcpy_pinned_h2d_batch_async + ) def _get_classed_reqs( self, diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index ba95189e90..388378bb6f 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -134,7 +134,7 @@ def pad_once(model_input, new_batch_size): graph_flags_at_att_init = [] - 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, @@ -234,3 +234,59 @@ def forward(state): assert seen_sizes == [expected_size] * (2 if overlap else 1) torch.testing.assert_close(output.logits, expected) torch.testing.assert_close(torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1), expected) + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_prefill_preparation_preserves_index_event_and_initialization_order(monkeypatch, overlap): + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = SimpleNamespace(enable_prefill_decode_mixed=False, enable_tpsp_mix_mode=True) + model.tp_world_size_ = 2 + model.prefill_graph = None + calls = [] + + class ReadyEvent: + def record(self): + calls.append("ready") + + monkeypatch.setattr(basemodel.torch.cuda, "Event", ReadyEvent) + monkeypatch.setattr(basemodel.dist_group_manager, "clear_deepep_buffer", lambda: None) + + def pad_input(model_input, handle_token_num): + assert handle_token_num == 4 + return model_input + + def create_state(model_input, microbatch_index=0): + calls.append(("select_indexes", microbatch_index)) + return SimpleNamespace( + init_some_extra_state=lambda _: calls.append(("extra", microbatch_index)), + init_att_state=lambda: calls.append(("attention", microbatch_index)), + ) + + def forward(*args, **kwargs): + calls.append("forward") + return ModelOutput(logits=torch.ones(1, 4)) + + model._create_padded_prefill_model_input = pad_input + model._create_inferstate = create_state + model._context_forward = forward + model._overlap_tpsp_context_forward = lambda *args, **kwargs: (forward(), forward()) + model_input = SimpleNamespace(is_prefill=True, input_ids=torch.arange(3), batch_size=1) + if overlap: + outputs = model._microbatch_overlap_prefill_cuda(model_input, model_input) + expected = [ + ("select_indexes", 0), + ("extra", 0), + ("attention", 0), + ("select_indexes", 1), + ("extra", 1), + ("attention", 1), + "ready", + "forward", + "forward", + ] + assert outputs[0].prefill_mem_indexes_ready_event is outputs[1].prefill_mem_indexes_ready_event + else: + outputs = (model._prefill(model_input),) + expected = [("select_indexes", 0), "ready", ("extra", 0), ("attention", 0), "forward"] + assert calls == expected + assert all(isinstance(output.prefill_mem_indexes_ready_event, ReadyEvent) for output in outputs) diff --git a/unit_tests/common/test_infer_batch_cache_release.py b/unit_tests/common/test_infer_batch_cache_release.py new file mode 100644 index 0000000000..bb7e64b6bf --- /dev/null +++ b/unit_tests/common/test_infer_batch_cache_release.py @@ -0,0 +1,486 @@ +"""Exercise request release with real radix trees, KV allocation and checkpoint slot pools.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from sortedcontainers import SortedDict + +from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 +from lightllm.common.kv_cache_mem_manager import allocator as allocator_impl +from lightllm.common.state_cache_manager.base import StateCacheManager +from lightllm.server.multi_level_kv_cache import CacheTier +from lightllm.server.router.dynamic_prompt import hybrid_att_radix_cache, radix_cache +from lightllm.server.router.model_infer import infer_batch +from lightllm.server.router.model_infer.infer_req_ext1 import PinMemPageCacheReqTokenIndex, create_req_token_index + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires pinned memory and CUDA index tables") + + +class _LocalSharedArray: + def __init__(self, name, shape, dtype): + self.arr = np.zeros(shape, dtype=dtype) + + +class _StateSlots(StateCacheManager): + def __init__(self): + super().__init__(size=16) + self.states = torch.empty((self.size, 1)) + + def get_state_cache(self, buffer_idx): + return self.states[buffer_idx] + + +@pytest.fixture +def make_release_case(monkeypatch): + # Only replace interprocess counters; allocation, tree operations and slot accounting are production code. + monkeypatch.setattr(radix_cache, "SharedArray", _LocalSharedArray) + monkeypatch.setattr(hybrid_att_radix_cache, "SharedArray", _LocalSharedArray) + monkeypatch.setattr(allocator_impl, "SharedInt", lambda name: SimpleNamespace(set_value=lambda value: None)) + monkeypatch.setattr(allocator_impl, "get_current_rank_in_node", lambda: 0) + monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_page_offsets", None) + + def make_case(attention, mode, page_size, cache_case): + hybrid = attention != "full" + partial_prefill = attention == "hybrid_partial" + allocator = allocator_impl.KvCacheAllocator(64, page_size) + small_slots, big_slots = _StateSlots(), _StateSlots() + memory = SimpleNamespace( + page_size=page_size, + allocator=allocator, + alloc=allocator.alloc, + free=allocator.free, + big_page_buffers=big_slots, + ) + if hybrid: + cache = hybrid_att_radix_cache.HybridAttPagedRadixCache( + 64, + 0, + hash_page_size=4, + big_page_num=2, + kv_cache_mem_manager=memory, + small_page_buffers=small_slots, + ) + else: + cache = radix_cache.RadixCache(64, 0, mem_manager=memory, page_size=page_size) + key = torch.arange(32, dtype=torch.int64) + seeded_indexes = None + if cache_case != "uncached": + seeded_indexes = allocator.alloc(8).clone() + if hybrid: + cache.insert( + key[:8], + seeded_indexes, + block_hashs=[1, 2], + block_state_idxs=[None, None], + len_to_big_page_id=SortedDict({8: big_slots.alloc_one_state_cache()}), + ) + else: + cache.insert(key[:8], seeded_indexes) + + shared_node, shared_indexes = None, None + if cache_case == "shared": + if hybrid: + shared_node, shared_len, shared_indexes = cache.match_prefix( + key[:8], block_hashs=[1, 2], update_refs=True + ) + else: + shared_node, shared_len, shared_indexes = cache.match_prefix(key[:8], update_refs=True) + assert shared_len == cache.get_refed_tokens_num() == 8 + shared_len = 0 if shared_node is None else shared_node.node_prefix_total_len + context = infer_batch.InferenceContext() + context.args = SimpleNamespace(page_size=page_size, linear_att_hash_page_size=4, linear_att_page_block_num=2) + context.is_hybrid_att_model = hybrid + context.radix_cache = cache + context.req_manager = SimpleNamespace( + page_size=page_size, req_to_token_indexs=torch.full((1, 32), -1, dtype=torch.int32, device="cuda") + ) + req = SimpleNamespace( + req_idx=0, + req_manager=context.req_manager, + args=SimpleNamespace(req_token_index_mode=mode), + cur_kv_len=12 if partial_prefill else 13, + hold_kv_len=16, + cache_tiers=(CacheTier.GPU,), + shared_kv_node=shared_node, + hybrid_cache_len=24 if partial_prefill else 12, + tail_small_page_buffer_id=small_slots.alloc_one_state_cache() if attention == "hybrid_small" else None, + hybrid_len_to_big_page_id=SortedDict(), + get_input_token_ids=lambda: key.tolist(), + shm_req=SimpleNamespace( + input_len=25 if partial_prefill else 13, + shm_cur_kv_len=12 if partial_prefill else 13, + hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [1, 2, 3, 4, 5, 6]), + ), + ) + if hybrid and shared_len < 8: + req.hybrid_len_to_big_page_id[8] = big_slots.alloc_one_state_cache() + req.token_index = create_req_token_index(req, torch.cuda.current_stream()) + if shared_indexes is not None: + req.token_index.append_token_indexes(shared_indexes) + req.token_index.append_token_indexes(allocator.alloc(req.hold_kv_len - shared_len)) + cache_len = (13 // page_size * page_size) if not hybrid else (12 if attention == "hybrid_small" else 8) + expected_indexes = req.token_index.get_cpu_token_indexes(0, cache_len) + if seeded_indexes is not None: + expected_indexes[:8] = seeded_indexes + monkeypatch.setattr(infer_batch, "g_infer_context", context) + monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: context.args) + return context, req, allocator, small_slots, big_slots, expected_indexes + + return make_case + + +@pytest.mark.parametrize("attention", ["full", "hybrid_small", "hybrid_big", "hybrid_partial"]) +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("page_size", [1, 4]) +@pytest.mark.parametrize("cache_case", ["uncached", "duplicate", "shared"]) +def test_release_and_eviction_conserve_kv_pages_and_checkpoint_slots( + make_release_case, attention, mode, page_size, cache_case +): + context, req, allocator, small_slots, big_slots, expected_indexes = make_release_case( + attention, mode, page_size, cache_case + ) + cache = context.radix_cache + released = [] + context.free_a_req_mem(released, req) + allocator.free(torch.cat(released)) + # Reusing the request row must not change cached values, including pure_gpu release views. + context.req_manager.req_to_token_indexs.fill_(-99) + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == 0 + assert req.token_index.token_index_len == req.token_index.gpu_index_submitted_len == 0 + assert req.shared_kv_node is None and req.tail_small_page_buffer_id is None + assert not req.hybrid_len_to_big_page_id + cache_len = expected_indexes.numel() + assert cache.get_refed_tokens_num() == 0 + assert cache.root_node.ref_counter == 1 + assert cache.get_tree_total_tokens_num() == cache_len + assert allocator.can_use_mem_size == allocator.size - cache_len + assert big_slots.get_used_cache_num() == (1 if attention != "full" and cache_len > 0 else 0) + assert small_slots.get_used_cache_num() == (1 if attention == "hybrid_small" else 0) + + key = torch.arange(cache_len, dtype=torch.int64) + if attention == "full": + node, matched, indexes = cache.match_prefix(key, update_refs=True) + else: + node, matched, indexes = cache.match_prefix(key, block_hashs=[1, 2, 3][: cache_len // 4], update_refs=True) + assert matched == cache_len + if cache_len: + torch.testing.assert_close(indexes, expected_indexes, check_dtype=False) + cache.dec_node_ref_counter(node) + else: + assert node is None and indexes is None + + free_indexes = allocator.alloc(allocator.can_use_mem_size) + assert len(set(free_indexes.tolist())) == free_indexes.numel() + assert set(free_indexes.tolist()).isdisjoint(expected_indexes.tolist()) + allocator.free(free_indexes) + # Real eviction must return every physical KV page and both kinds of checkpoint slot exactly once. + cache.free_radix_cache_to_get_enough_token(allocator.size) + assert cache.get_tree_total_tokens_num() == cache.get_refed_tokens_num() == 0 + assert allocator.can_use_mem_size == allocator.size + assert sorted(allocator.alloc(allocator.size).tolist()) == list(range(allocator.size)) + for slots in (small_slots, big_slots): + assert slots.get_used_cache_num() == 0 + assert sorted(slots.free_list) == list(range(slots.size)) + + +@pytest.mark.parametrize( + "checkpoint_lens,computed_len,target_len,inserted_len", + [ + ([8, 16], 20, 24, 16), + ([8, 16, 24], 20, 24, 16), + ([8, 16], 16, 24, 16), + ([8], 7, 24, 0), + ([], 7, 24, 0), + ([8, 16], 20, 12, 8), + ([8, 16, 24], 24, 24, 24), + ([8], 20, 12, 8), + ], + ids=[ + "multiple", + "future", + "at_boundary", + "future_only", + "none", + "target_limit", + "final_big", + "decode_after_target", + ], +) +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("cache_case", ["uncached", "duplicate", "shared"]) +def test_partial_prefill_validates_checkpoints_and_caches_completed_prefix( + make_release_case, monkeypatch, checkpoint_lens, computed_len, target_len, inserted_len, mode, cache_case +): + context, req, allocator, small_slots, big_slots, first_page_indexes = make_release_case( + "hybrid_partial", mode, 4, cache_case + ) + cache = context.radix_cache + shared_len = 8 if cache_case == "shared" else 0 + req.cur_kv_len = max(computed_len, shared_len) + req.hybrid_cache_len = target_len + req.token_index.append_token_indexes(allocator.alloc(8)) + req.hold_kv_len = 24 + for length in list(req.hybrid_len_to_big_page_id): + if length not in checkpoint_lens: + big_slots.free_state_cache([req.hybrid_len_to_big_page_id.pop(length)]) + for length in checkpoint_lens: + if length > shared_len and length not in req.hybrid_len_to_big_page_id: + req.hybrid_len_to_big_page_id[length] = big_slots.alloc_one_state_cache() + + cache_len = max(inserted_len, 8 if cache_case != "uncached" else 0) + expected_indexes = req.token_index.get_cpu_token_indexes(0, cache_len) + if cache_case != "uncached": + expected_indexes[:8] = first_page_indexes + checkpoint = req.hybrid_len_to_big_page_id.get(inserted_len) + insert = Mock(wraps=cache.insert) + monkeypatch.setattr(cache, "insert", insert) + released = [] + no_new_big_page = inserted_len <= shared_len + out_of_range_checkpoint = any(length > max(shared_len, inserted_len) for length in req.hybrid_len_to_big_page_id) + if req.hybrid_len_to_big_page_id and (no_new_big_page or out_of_range_checkpoint): + # 无新大页时持有 checkpoint，或 checkpoint 超出插入前缀，均违反请求状态约束。 + checkpoints = dict(req.hybrid_len_to_big_page_id) + available_tokens = allocator.can_use_mem_size + cached_tokens = cache.get_tree_total_tokens_num() + with pytest.raises(AssertionError): + context.free_a_req_mem(released, req) + if no_new_big_page: + insert.assert_not_called() + assert released == [] + assert req.cur_kv_len == max(computed_len, shared_len) + assert req.hybrid_len_to_big_page_id == checkpoints + assert cache.get_tree_total_tokens_num() == cached_tokens + else: + assert insert.call_count == 1 + assert len(released) == 1 + assert req.cur_kv_len == inserted_len + # insert 可能已接管有效 checkpoint，未消费的条目仍由请求持有。 + assert set(req.hybrid_len_to_big_page_id.items()) <= set(checkpoints.items()) + assert cache.get_tree_total_tokens_num() <= inserted_len + assert req.token_index.token_index_len == req.hold_kv_len == 24 + assert allocator.can_use_mem_size == available_tokens + assert cache.get_refed_tokens_num() == shared_len + assert big_slots.get_used_cache_num() == cache.get_tree_total_tokens_num() // 8 + len( + req.hybrid_len_to_big_page_id + ) + return + context.free_a_req_mem(released, req) + allocator.free(torch.cat(released)) + context.req_manager.req_to_token_indexs.fill_(-99) + + assert insert.call_count == 1 + assert insert.call_args.args[0].numel() == max(shared_len, inserted_len) + assert req.cur_kv_len == req.hold_kv_len == req.token_index.token_index_len == 0 + assert req.shared_kv_node is None and not req.hybrid_len_to_big_page_id + assert cache.get_tree_total_tokens_num() == cache_len + assert cache.get_refed_tokens_num() == 0 + assert allocator.can_use_mem_size == allocator.size - cache_len + assert small_slots.get_used_cache_num() == 0 + assert big_slots.get_used_cache_num() == cache_len // 8 + node, matched, indexes = cache.match_prefix( + torch.arange(24, dtype=torch.int64), block_hashs=[1, 2, 3, 4, 5, 6], update_refs=True + ) + assert matched == cache_len + if cache_len: + torch.testing.assert_close(indexes, expected_indexes, check_dtype=False) + if cache_case == "uncached" or cache_len > 8: + assert node.big_page_buffer_idx == checkpoint + cache.dec_node_ref_counter(node) + else: + assert node is None and indexes is None + + cache.free_radix_cache_to_get_enough_token(allocator.size) + assert allocator.can_use_mem_size == allocator.size + assert sorted(allocator.alloc(allocator.size).tolist()) == list(range(allocator.size)) + assert big_slots.get_used_cache_num() == 0 + assert sorted(big_slots.free_list) == list(range(big_slots.size)) + + +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("reserved_len", [0, 8]) +def test_hybrid_release_before_first_forward_handles_reserved_capacity( + make_release_case, monkeypatch, mode, reserved_len +): + context, req, allocator, small_slots, big_slots, _ = make_release_case("hybrid_partial", mode, 4, "uncached") + allocator.free(req.token_index.get_free_token_indexes(0, req.hold_kv_len)) + req.token_index.reset() + big_slots.free_state_cache(list(req.hybrid_len_to_big_page_id.values())) + req.hybrid_len_to_big_page_id.clear() + req.cur_kv_len = 0 + req.hold_kv_len = reserved_len + if reserved_len: + req.token_index.append_token_indexes(allocator.alloc(reserved_len)) + insert = Mock(wraps=context.radix_cache.insert) + monkeypatch.setattr(context.radix_cache, "insert", insert) + + released = [] + context.free_a_req_mem(released, req) + if reserved_len: + allocator.free(torch.cat(released)) + assert insert.call_count == 1 + assert insert.call_args.args[0].numel() == insert.call_args.args[1].numel() == 0 + else: + insert.assert_not_called() + assert released == [] + assert context.radix_cache.get_tree_total_tokens_num() == context.radix_cache.get_refed_tokens_num() == 0 + assert context.radix_cache.root_node.ref_counter == 1 + assert req.cur_kv_len == req.hold_kv_len == req.token_index.token_index_len == 0 + assert allocator.can_use_mem_size == allocator.size + assert sorted(allocator.alloc(allocator.size).tolist()) == list(range(allocator.size)) + assert small_slots.get_used_cache_num() == big_slots.get_used_cache_num() == 0 + + +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("new_tokens", [0, 1]) +@pytest.mark.parametrize("evict_checkpoint", [False, True]) +def test_hybrid_release_skips_insert_for_shared_small_prefix( + make_release_case, monkeypatch, mode, new_tokens, evict_checkpoint +): + context, req, allocator, small_slots, big_slots, expected_indexes = make_release_case( + "hybrid_small", mode, 4, "uncached" + ) + # Use the real match path with big pages disabled, which retains a shared small node. + big_slots.free_state_cache(list(req.hybrid_len_to_big_page_id.values())) + req.hybrid_len_to_big_page_id.clear() + context.args.linear_att_page_block_num = 10000000 + context.radix_cache = hybrid_att_radix_cache.HybridAttPagedRadixCache( + 64, + 0, + hash_page_size=4, + big_page_num=context.args.linear_att_page_block_num, + kv_cache_mem_manager=context.radix_cache.mem_manager, + small_page_buffers=small_slots, + ) + req.args = SimpleNamespace( + req_token_index_mode=mode, + page_size=4, + linear_att_hash_page_size=4, + linear_att_page_block_num=context.args.linear_att_page_block_num, + max_req_total_len=32, + ) + req.sampling_param = SimpleNamespace(disable_prompt_cache=False) + req.get_cur_total_len = lambda: 13 + req.shm_req.shm_prompt_ids = SimpleNamespace(arr=np.arange(13)) + req.shm_req.hybrid_token_hash_list = SimpleNamespace(size=3, get_all=lambda: [1, 2, 3]) + restored_checkpoints = [] + context.req_manager.restore_small_page_state = lambda req: restored_checkpoints.append( + req.shared_kv_node.small_page_buffer_idx + ) + released = [] + context.free_a_req_mem(released, req) + allocator.free(torch.cat(released)) + cache = context.radix_cache + infer_batch.InferReq._hybrid_match_radix_cache(req) + node = req.shared_kv_node + shared_len = req.cur_kv_len + checkpoint = node.small_page_buffer_idx + assert shared_len == 12 and checkpoint is not None + assert not node.is_big_page_node() + assert restored_checkpoints == [checkpoint] + reserved_checkpoints = [] + if evict_checkpoint: + # Exhaust the slot pool, then recycle the checkpoint while its node is still referenced. + reserved_checkpoints = small_slots.alloc_state_cache(small_slots.get_free_cache_num()) + cache.free_one_small_page_buffer() + assert node.ref_counter > 0 and node.small_page_buffer_idx is None + reused_checkpoint = small_slots.alloc_one_state_cache() + assert reused_checkpoint == checkpoint + reserved_checkpoints.append(reused_checkpoint) + small_slots.states[reused_checkpoint].fill_(42) + req.cur_kv_len = shared_len + new_tokens + req.hold_kv_len = 16 + req.token_index.append_token_indexes(allocator.alloc(req.hold_kv_len - shared_len)) + insert = Mock(wraps=cache.insert) + monkeypatch.setattr(cache, "insert", insert) + + released = [] + context.free_a_req_mem(released, req) + allocator.free(torch.cat(released)) + context.req_manager.req_to_token_indexs.fill_(-99) + + insert.assert_not_called() + assert req.shared_kv_node is None + assert req.cur_kv_len == req.hold_kv_len == req.token_index.token_index_len == 0 + assert cache.get_tree_total_tokens_num() == shared_len + assert cache.get_refed_tokens_num() == 0 + assert allocator.can_use_mem_size == allocator.size - shared_len + assert big_slots.get_used_cache_num() == 0 + assert small_slots.get_used_cache_num() == (small_slots.size if evict_checkpoint else 1) + assert node.small_page_buffer_idx == (None if evict_checkpoint else checkpoint) + if evict_checkpoint: + assert small_slots.states[reused_checkpoint].item() == 42 + node, matched, indexes = cache.match_prefix( + torch.arange(12, dtype=torch.int64), block_hashs=[1, 2, 3], update_refs=True + ) + if evict_checkpoint: + assert node is None and indexes is None and matched == 0 + small_slots.free_state_cache(reserved_checkpoints) + else: + assert matched == shared_len + torch.testing.assert_close(indexes, expected_indexes, check_dtype=False) + cache.dec_node_ref_counter(node) + cache.free_radix_cache_to_get_enough_token(allocator.size) + assert allocator.can_use_mem_size == allocator.size + assert sorted(allocator.alloc(allocator.size).tolist()) == list(range(allocator.size)) + for slots in (small_slots, big_slots): + assert slots.get_used_cache_num() == 0 + assert sorted(slots.free_list) == list(range(slots.size)) + + +@pytest.mark.parametrize("prefix_len", [4, 12]) +def test_hybrid_new_small_prefix_still_requires_checkpoint(make_release_case, prefix_len): + context, req, _, _, big_slots, _ = make_release_case("hybrid_partial", "pin_mem_cache", 4, "uncached") + indexes = req.token_index.get_cpu_token_indexes(0, prefix_len) + big_checkpoints = req.hybrid_len_to_big_page_id if prefix_len == 12 else SortedDict() + original_checkpoints = dict(big_checkpoints) + with pytest.raises(AssertionError): + context.radix_cache.insert( + torch.arange(prefix_len, dtype=torch.int64), + indexes, + block_hashs=[1, 2, 3][: prefix_len // 4], + block_state_idxs=[None] * (prefix_len // 4), + len_to_big_page_id=big_checkpoints, + ) + assert context.radix_cache.get_tree_total_tokens_num() == 0 + assert context.radix_cache.root_node.children == {} + assert big_checkpoints == original_checkpoints + assert big_slots.get_used_cache_num() == 1 + + +@pytest.mark.parametrize("attention", ["hybrid_small", "hybrid_partial"]) +@pytest.mark.parametrize("failed_range", ["tail", "cached_prefix"]) +def test_hybrid_snapshot_failure_keeps_checkpoint_owned_by_request( + make_release_case, monkeypatch, attention, failed_range +): + context, req, allocator, small_slots, big_slots, _ = make_release_case(attention, "pin_mem_cache", 4, "uncached") + original_kv_len = req.cur_kv_len + cache_len = 12 if attention == "hybrid_small" else 8 + checkpoint = req.tail_small_page_buffer_id + big_checkpoints = dict(req.hybrid_len_to_big_page_id) + get_indexes = req.token_index.get_cpu_token_indexes + + def failed_snapshot(start, end): + if (start == 0) == (failed_range == "cached_prefix"): + raise RuntimeError("snapshot failed") + return get_indexes(start, end) + + with monkeypatch.context() as failure: + failure.setattr(req.token_index, "get_cpu_token_indexes", failed_snapshot) + with pytest.raises(RuntimeError, match="snapshot failed"): + context.free_a_req_mem([], req) + assert req.tail_small_page_buffer_id == checkpoint + assert req.hybrid_len_to_big_page_id == big_checkpoints + assert req.cur_kv_len == (original_kv_len if failed_range == "tail" else cache_len) + assert small_slots.get_used_cache_num() == (1 if attention == "hybrid_small" else 0) + assert big_slots.get_used_cache_num() == len(big_checkpoints) + assert context.radix_cache.get_tree_total_tokens_num() == 0 + released = [] + context.free_a_req_mem(released, req) + allocator.free(torch.cat(released)) + context.radix_cache.free_radix_cache_to_get_enough_token(allocator.size) + assert allocator.can_use_mem_size == allocator.size + assert small_slots.get_used_cache_num() == big_slots.get_used_cache_num() == 0 diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index b3a92969c2..62d43a47ad 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -8,6 +8,7 @@ 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 +from lightllm.server.multi_level_kv_cache import CacheTier from unit_tests.common.cpu_index_helpers import make_cpu_index_manager, bind_cpu_indexes @@ -485,10 +486,110 @@ def test_page_size_one_frees_all_preallocated_indexes(): assert req.token_index.token_index_len == req.token_index.gpu_index_submitted_len == 0 +@pytest.mark.parametrize("attention", ["full", "hybrid_small", "hybrid_big"]) +def test_cache_release_preserves_snapshots_and_reclaims_tail_duplicates_and_reference(monkeypatch, attention): + hybrid = attention != "full" + context = InferenceContext() + context.args = SimpleNamespace(page_size=4, linear_att_hash_page_size=4, linear_att_page_block_num=2) + context.is_hybrid_att_model = hybrid + context.req_manager = make_cpu_index_manager(torch.arange(16, dtype=torch.int32)[None, :], page_size=4) + cached = [] + released_refs = [] + + def insert(key, indexes, *, block_hashs=None, block_state_idxs=None, len_to_big_page_id=None): + cached.append((key, indexes, block_hashs, block_state_idxs)) + # 模拟 radix cache 接管请求已生成的大页 checkpoint。 + if hybrid: + len_to_big_page_id.clear() + return 8, None + + context.radix_cache = SimpleNamespace(insert=insert, dec_node_ref_counter=released_refs.append) + shared_node = SimpleNamespace(node_prefix_total_len=4) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=13 if hybrid else 11, + hold_kv_len=16, + cache_tiers=(CacheTier.GPU,), + shared_kv_node=shared_node, + hybrid_cache_len=12, + tail_small_page_buffer_id=7 if attention == "hybrid_small" else None, + hybrid_len_to_big_page_id={8: 9} if hybrid else {}, + get_input_token_ids=lambda: list(range(13)), + shm_req=SimpleNamespace( + shm_cur_kv_len=13, + hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [100, 200, 300]), + ), + ) + bind_cpu_indexes(req, context.req_manager) + monkeypatch.setattr(infer_batch, "g_infer_context", context) + monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: context.args) + released = [] + + context.free_a_req_mem(released, req) + + cache_len = 12 if attention == "hybrid_small" else 8 + assert [indexes.tolist() for indexes in released] == [list(range(cache_len, 16)), list(range(4, 8))] + assert released_refs == [shared_node] + assert req.shared_kv_node is None + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == 0 + assert req.token_index.token_index_len == req.token_index.gpu_index_submitted_len == 0 + assert req.tail_small_page_buffer_id is None and req.hybrid_len_to_big_page_id == {} + context.req_manager.token_indexes_cpu.fill_(-99) + key, indexes, block_hashs, block_state_idxs = cached[0] + assert key.tolist() == indexes.tolist() == list(range(cache_len)) + if hybrid: + assert block_hashs == [100, 200, 300][: cache_len // 4] + assert block_state_idxs == ([None, None, 7] if attention == "hybrid_small" else [None, None]) + else: + assert block_hashs is None and block_state_idxs is None + + +def _insert_empty_hybrid_prefix(key, indexes, *, block_hashs, block_state_idxs, len_to_big_page_id): + assert key.numel() == indexes.numel() == 0 + assert block_hashs == block_state_idxs == [] + assert not len_to_big_page_id + return 0, None + + +def test_hybrid_release_without_new_big_page_rejects_pending_checkpoints(monkeypatch): + context = InferenceContext() + context.args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) + context.is_hybrid_att_model = True + context.req_manager = make_cpu_index_manager(torch.arange(8, dtype=torch.int32)[None, :]) + freed_states = [] + context.radix_cache = SimpleNamespace( + insert=_insert_empty_hybrid_prefix, + big_page_buffers=SimpleNamespace(free_state_cache=freed_states.extend), + ) + 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={8: 9}, + shared_kv_node=None, + get_input_token_ids=lambda: list(range(8)), + shm_req=SimpleNamespace(hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [])), + ) + bind_cpu_indexes(req, context.req_manager) + monkeypatch.setattr(infer_batch, "g_infer_context", context) + monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: context.args) + released = [] + + with pytest.raises(AssertionError): + context._hybrid_att_free_req(released, req) + + assert released == [] + assert freed_states == [] + assert req.hybrid_len_to_big_page_id == {8: 9} + assert req.cur_kv_len == 5 and req.hold_kv_len == 8 + + 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.radix_cache = SimpleNamespace(insert=_insert_empty_hybrid_prefix) infer_context.req_manager = make_cpu_index_manager( token_table=torch.arange(16, dtype=torch.int32)[None, :], page_size=1 ) @@ -500,6 +601,8 @@ def test_linear_attention_frees_reserved_page_tail(monkeypatch): tail_small_page_buffer_id=None, hybrid_len_to_big_page_id={}, shared_kv_node=None, + get_input_token_ids=lambda: list(range(8)), + shm_req=SimpleNamespace(hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [])), ) bind_cpu_indexes(req, infer_context.req_manager) free_token_indexes = [] @@ -514,7 +617,7 @@ def test_linear_attention_frees_reserved_page_tail(monkeypatch): 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.radix_cache = SimpleNamespace(insert=_insert_empty_hybrid_prefix) infer_context.req_manager = make_cpu_index_manager( token_table=torch.arange(8, dtype=torch.int32)[None, :], page_size=1 ) @@ -526,6 +629,8 @@ def test_linear_attention_frees_preallocated_pages_before_first_forward(monkeypa tail_small_page_buffer_id=None, hybrid_len_to_big_page_id={}, shared_kv_node=None, + get_input_token_ids=lambda: list(range(8)), + shm_req=SimpleNamespace(hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [])), ) bind_cpu_indexes(req, infer_context.req_manager) free_token_indexes = [] diff --git a/unit_tests/common/test_req_token_index.py b/unit_tests/common/test_req_token_index.py index 95d1247e73..9dbedfbd78 100644 --- a/unit_tests/common/test_req_token_index.py +++ b/unit_tests/common/test_req_token_index.py @@ -232,7 +232,7 @@ def test_batch_upload_rejects_missing_gpu_submission_without_advancing_progress( monkeypatch.setattr(base_backend, "g_infer_context", SimpleNamespace(req_manager=req.req_manager)) backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) - with pytest.raises(AssertionError, match="GPU index submission progress does not match recorded length"): + with pytest.raises(TypeError): backend._copy_req_token_indexes_to_gpu([req]) assert req.token_index.gpu_index_submitted_len == 0 @@ -264,6 +264,39 @@ def unexpected_upload(*args): assert [req.token_index.gpu_index_submitted_len for req in reqs] == [4 if mode == "pure_gpu" else 0] * 2 +@requires_cuda +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache"]) +def test_partial_batch_submission_can_retry_and_share_progress_with_gpu_getter(mode): + table = torch.full((2, 16), -1, dtype=torch.int32, device="cuda") + helpers = [ + create_req_token_index(make_req(table, mode, req_idx=row, page_size=4), torch.cuda.current_stream()) + for row in range(2) + ] + for row, helper in enumerate(helpers): + helper.append_token_indexes(torch.arange(row * 8, row * 8 + 4, dtype=torch.int32)) + + def partial_copy(destinations, sources): + destinations[0].copy_(sources[0], non_blocking=True) + raise RuntimeError("partial submission failed") + + with pytest.raises(RuntimeError, match="partial submission failed"): + ReqTokenIndex.submit_pending_uploads(helpers, copy_batch=partial_copy) + assert [helper.gpu_index_submitted_len for helper in helpers] == [0, 0] + + # 重试须覆盖完整待提交区间；单请求 getter 随后只提交新追加部分。 + ReqTokenIndex.submit_pending_uploads(helpers) + helpers[0].append_token_indexes(torch.arange(4, 8, dtype=torch.int32)) + assert helpers[0].get_gpu_token_indexes(0, 8).cpu().tolist() == list(range(8)) + assert helpers[1].get_gpu_token_indexes(0, 4).cpu().tolist() == list(range(8, 12)) + assert [helper.gpu_index_submitted_len for helper in helpers] == [8, 4] + + def unexpected_copy(*args): + pytest.fail("Getter and batch submission must share the same upload progress") + + ReqTokenIndex.submit_pending_uploads(helpers, copy_batch=unexpected_copy) + assert table[0, 8:].eq(-1).all() and table[1, 4:].eq(-1).all() + + @requires_cuda @pytest.mark.parametrize("mode", MODES) @pytest.mark.parametrize("page_size", [1, 3, 4, 64]) From bae2aaf7e3ed43a3d078e8a40d33c7284c6a6613 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 15:00:50 +0000 Subject: [PATCH 19/22] Refactor request cache prefix state updates --- .../server/router/model_infer/infer_batch.py | 78 ++++++++++++------- .../common/test_infer_batch_cache_release.py | 3 +- 2 files changed, 53 insertions(+), 28 deletions(-) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 522477c93d..6cd22cc389 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -623,14 +623,7 @@ def _match_radix_cache(self): key = key[0 : len(key) - 1] # 最后一个不需要，因为需要一个额外的token，让其在prefill的时候输出下一个token的值 share_node, kv_len, value_tensor = g_infer_context.radix_cache.match_prefix(key, update_refs=True) if share_node is not None: - self.shared_kv_node = share_node - ready_cache_len = share_node.node_prefix_total_len - # 新请求索引为空，暂停恢复前也已统一释放并 reset，可直接追加命中前缀。 - self.token_index.append_token_indexes(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) self.shm_req.shm_cur_kv_len = self.cur_kv_len return @@ -662,11 +655,7 @@ def _hybrid_match_radix_cache(self): 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 - self.token_index.append_token_indexes(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 的命中长度 + self._set_cache_prefix(share_node, value_tensor) assert self.tail_small_page_buffer_id is None # 恢复 hybrid checkpoint g_infer_context.req_manager.restore_big_page_state( @@ -676,11 +665,7 @@ def _hybrid_match_radix_cache(self): # 小页匹配 if big_page_is_disable: # 如果 大页本质是被禁用的，可以直接使用小页的匹配结果 - self.shared_kv_node = share_node - ready_cache_len = share_node.node_prefix_total_len - self.token_index.append_token_indexes(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 的命中长度 + self._set_cache_prefix(share_node, value_tensor) assert self.tail_small_page_buffer_id is None # 恢复 hybrid checkpoint g_infer_context.req_manager.restore_small_page_state( @@ -698,8 +683,6 @@ def _hybrid_match_radix_cache(self): 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) - self.token_index.append_token_indexes(value_tensor[:cur_big_page_tokens]) - self.token_index.append_token_indexes(tail_mems) # 将 对应的 value_tensors 中的 kv 数据 拷贝到 tail_mems 中对应的数据去 radix_cache.mem_manager.operator.copy_mem_to_mem( @@ -718,20 +701,19 @@ def _hybrid_match_radix_cache(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 的命中长度 + self._set_cache_prefix( + shared_node=big_page_shared_node, + shared_prefix_indexes=value_tensor[:cur_big_page_tokens], + private_tail_indexes=tail_mems, + ) else: # 没有充足的token 容量时， 直接找到最接近的大页，进行大页恢复 share_node = radix_cache.deref_to_first_big_page_node(node=share_node) 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 - self.token_index.append_token_indexes(value_tensor[: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 的命中长度 + self._set_cache_prefix(share_node, value_tensor[:ready_cache_len]) assert self.tail_small_page_buffer_id is None # 恢复 hybrid checkpoint g_infer_context.req_manager.restore_big_page_state( @@ -751,6 +733,48 @@ def _hybrid_match_radix_cache(self): g_infer_context.req_manager.init_hybrid_attention_state(req=self) return + def _set_cache_prefix( + self, + shared_node: Optional[Union[TreeNode, HybridAttPagedTreeNode]], + shared_prefix_indexes: torch.Tensor, + private_tail_indexes: Optional[torch.Tensor] = None, + ): + """登记已命中的 KV 前缀，统一更新请求索引、持有长度和缓存命中信息。 + + Args: + shared_node: 请求最终保留引用的 radix cache 节点。普通匹配时为命中节点； + hybrid 小页尾部复制到私有 KV 后，为回退到的最近大页节点。 + 若没有可保留的共享大页，则为 None。节点引用计数由调用方的匹配、回退流程维护。 + shared_prefix_indexes: 共享前缀对应的物理 KV 槽位索引，为按逻辑 token 顺序排列的 + 一维 CPU Tensor，每个元素代表一个 token 槽位，覆盖请求从位置 0 开始的前缀。 + 长度对应 shared_node 的完整前缀长度；没有共享节点时传空 Tensor。 + private_tail_indexes: 紧接共享前缀的请求私有 KV 槽位索引，同样是一维 CPU Tensor。 + 用于 hybrid 小页命中后，将尾部 KV 复制到新分配槽位的场景；这里传入复制目标 + 的索引。其 KV 由当前请求独占，不计入 shared_node 的共享长度。无需私有尾部时 + 使用 None。两段索引均使用 token 单位，多 token 页模式下须覆盖完整物理页。 + + 使用场景： + 普通前缀命中、hybrid 大页命中，以及大页禁用时直接复用小页，都只传共享前缀。 + hybrid 小页命中且允许重建尾部时，调用方先提交尾部 KV 拷贝、恢复小页 checkpoint， + 再回退共享节点，最后分别传入保留的共享前缀和私有尾部索引，无需先拼接 Tensor。 + 容量不足、只能回退到大页时，只传该大页覆盖的共享前缀。 + + 调用时请求索引必须为空，适用于新请求或已完成暂停释放的请求。 + cur_kv_len、hold_kv_len 和 prompt_cache_len 均设为两段索引的总长度，因为私有尾部 + 也来自已计算的缓存 KV。例如共享前缀为 8、私有尾部为 4 时，命中长度为 12， + 共享节点仍只保护前 8 个 token。shm_cur_kv_len 由外层匹配方法统一发布。 + """ + # 新请求或已完成暂停释放的请求，其索引状态为空，无需再次 reset。 + self.shared_kv_node = shared_node + self.token_index.append_token_indexes(shared_prefix_indexes) + cache_len = shared_prefix_indexes.numel() + if private_tail_indexes is not None: + self.token_index.append_token_indexes(private_tail_indexes) + cache_len += private_tail_indexes.numel() + self.cur_kv_len = self.hold_kv_len = int(cache_len) + assert self.hold_kv_len % self.args.page_size == 0 + self.shm_req.prompt_cache_len = self.cur_kv_len + def is_master_req(self): """ diverse 模式下，判断当前请求是否为独立主请求，其进行prefill后，将 diff --git a/unit_tests/common/test_infer_batch_cache_release.py b/unit_tests/common/test_infer_batch_cache_release.py index bb7e64b6bf..50be25b5b9 100644 --- a/unit_tests/common/test_infer_batch_cache_release.py +++ b/unit_tests/common/test_infer_batch_cache_release.py @@ -1,6 +1,6 @@ """Exercise request release with real radix trees, KV allocation and checkpoint slot pools.""" -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from unittest.mock import Mock import numpy as np @@ -374,6 +374,7 @@ def test_hybrid_release_skips_insert_for_shared_small_prefix( context.free_a_req_mem(released, req) allocator.free(torch.cat(released)) cache = context.radix_cache + req._set_cache_prefix = MethodType(infer_batch.InferReq._set_cache_prefix, req) infer_batch.InferReq._hybrid_match_radix_cache(req) node = req.shared_kv_node shared_len = req.cur_kv_len From b771127726fdd16b0f1ea0a332b31b3d3189f582 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 15:23:09 +0000 Subject: [PATCH 20/22] Refactor request KV allocation ownership and index pool naming --- lightllm/common/req_manager/base.py | 6 ++-- .../server/router/model_infer/infer_batch.py | 14 +++++++++ .../model_infer/mode_backend/base_backend.py | 21 ++----------- .../dp_backend/dp_shared_kv_trans.py | 4 +-- .../mode_backend/multi_level_kv_cache.py | 2 +- .../pd/decode_node_impl/decode_impl.py | 2 +- unit_tests/common/test_req_manager_cpu.py | 8 ++--- unit_tests/common/test_req_manager_page.py | 31 ++++++++++--------- .../mode_backend/test_multi_level_kv_cache.py | 4 +-- .../test_paged_side_allocations.py | 5 +-- 10 files changed, 48 insertions(+), 49 deletions(-) diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index 49ace35c07..2b13c70107 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -11,7 +11,7 @@ logger = init_logger("lightllm.common.req_manager") -class _ReqLinkedList: +class _ReqIndexPool: def __init__(self, max_request_num): self._free_indexes = deque(range(max_request_num)) self.marks = [0 for _ in range(max_request_num)] @@ -48,7 +48,7 @@ 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 self.page_size = page_size aligned_sequence_length = (max_sequence_length + page_size - 1) // page_size * page_size @@ -107,5 +107,5 @@ def free_token(self, free_token_index): return 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/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 6cd22cc389..ea8033c455 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -960,6 +960,20 @@ def _kv_cache_alloc_need(self, target_kv_len: int) -> int: assert alloc_token_num % page_size == 0 return alloc_token_num + def alloc_kv_mem(self, alloc_token_num: int) -> Optional[torch.Tensor]: + """按完整页追加 KV 容量，更新请求索引和持有长度；零申请返回 None。""" + 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) + + mem_indexes = self.req_manager.mem_manager.alloc(alloc_token_num) + self.token_index.append_token_indexes(mem_indexes) + self.hold_kv_len += alloc_token_num + return mem_indexes + 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 b2f9e5ae8d..243427f1f5 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -646,23 +646,6 @@ def remaining_prefill_tokens(req: InferReq) -> int: return ready_reqs # 一些可以复用的通用功能函数 - def _alloc_req_kv_mem( - self, - req_obj: InferReq, - alloc_token_num: int, - ) -> 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) - - mem_indexes = g_infer_context.req_manager.mem_manager.alloc(alloc_token_num) - req_obj.token_index.append_token_indexes(mem_indexes) - req_obj.hold_kv_len += alloc_token_num - return mem_indexes - def _copy_req_token_indexes_to_gpu(self, run_reqs: List[InferReq]): """选择本批请求；索引对象统一管理上传区间、提交进度和流依赖。""" ReqTokenIndex.submit_pending_uploads( @@ -801,7 +784,7 @@ def _get_classed_reqs( 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) + req_obj.alloc_kv_mem(alloc_token_num) prefill_tokens += token_num prefill_reqs.append(req_obj) can_alloc_token_num -= alloc_token_num @@ -819,7 +802,7 @@ def _get_classed_reqs( alloc_token_num += extra_token_num can_alloc_token_num -= extra_token_num # 将必需容量和可选预留合并申请，保留每个请求一次 alloc 的行为。 - self._alloc_req_kv_mem(req_obj, alloc_token_num) + req_obj.alloc_kv_mem(alloc_token_num) # 先由控制器确定请求需要写入的缓存层级，再按是否包含 CPU cache 决定是否发起 offload。 cache_controller = g_infer_context.cache_placement_controller 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 d7873584ff..ba569eb352 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 @@ -80,10 +80,10 @@ def build_shared_kv_trans_tasks( 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) + mem_indexes = req.alloc_kv_mem(alloc_token_num) assert mem_indexes is not None # mem_indexes 只描述需要复制的逻辑 KV；页尾预留槽位已经由 - # _alloc_req_kv_mem 写入请求表，但不参与本次跨 DP 传输。 + # alloc_kv_mem 记录到请求索引，但不参与本次跨 DP 传输。 mem_indexes = mem_indexes[:trans_size] 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]) 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 8aa9c37491..a1db9c740b 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 @@ -103,7 +103,7 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): assert ready_page_num <= len(page_list) need_pages = page_list[ready_page_num:] # 只取需要的页面 - allocated_mem_indexes = self.backend._alloc_req_kv_mem(req, need_token_num) + allocated_mem_indexes = req.alloc_kv_mem(need_token_num) assert allocated_mem_indexes is not None if self.need_sync_compute_stream(): 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 f0ef204c53..3306f6ba9d 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 @@ -123,7 +123,7 @@ def _decode_node_gen_trans_tasks(self, req_obj: InferReq): 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) + mem_indexes = req_obj.alloc_kv_mem(need_mem_size) assert mem_indexes is not None # 传输只覆盖真实 KV；最后一个模型页面中尚未使用的部分继续留在 # req_to_token_indexs 中，供后续 decode 直接切片使用。 diff --git a/unit_tests/common/test_req_manager_cpu.py b/unit_tests/common/test_req_manager_cpu.py index 0e8e64e906..30137cb04f 100644 --- a/unit_tests/common/test_req_manager_cpu.py +++ b/unit_tests/common/test_req_manager_cpu.py @@ -8,7 +8,7 @@ # Use the production import order to avoid the existing package initialization cycle. from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.req_manager.base import ReqManager, _ReqLinkedList +from lightllm.common.req_manager.base import ReqManager, _ReqIndexPool from lightllm.common.cuda_wrapper import cuda_memcpy_pinned_h2d_batch_async from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext from lightllm.server.router.model_infer.infer_req_ext1 import PinMemPageCacheReqTokenIndex, create_req_token_index @@ -33,12 +33,12 @@ def _make_manager(device="cpu"): ) infer_batch.g_infer_context.req_manager = manager manager.max_request_num = 2 - manager.req_list = _ReqLinkedList(2) + manager.req_list = _ReqIndexPool(2) return manager def test_req_index_allocation_preserves_order_capacity_and_reuse(): - indexes = _ReqLinkedList(3) + indexes = _ReqIndexPool(3) assert indexes.is_all_free() and indexes.can_alloc_size == 3 assert [indexes.alloc() for _ in range(3)] == [0, 1, 2] assert indexes.can_alloc_size == 0 and not indexes.is_all_free() @@ -55,7 +55,7 @@ def test_req_index_allocation_preserves_order_capacity_and_reuse(): def test_req_index_free_rejects_unallocated_and_duplicate_indexes(): - indexes = _ReqLinkedList(2) + indexes = _ReqIndexPool(2) with pytest.raises(AssertionError): indexes.free(0) assert indexes.alloc() == 0 diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index 62d43a47ad..3bc4314847 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -83,6 +83,7 @@ def copy_indexes(run_reqs): def _make_req(req_idx): req = InferReq.__new__(InferReq) req.req_idx = req_idx + req.args = infer_batch.g_infer_context.args bind_cpu_indexes(req, infer_batch.g_infer_context.req_manager) req.cur_kv_len = req.hold_kv_len = 0 req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) @@ -90,45 +91,45 @@ def _make_req(req_idx): def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypatch): - context, backend = _make_context(monkeypatch) + context, _ = _make_context(monkeypatch) req = _make_req(0) - mem_indexes = backend._alloc_req_kv_mem(req, alloc_token_num=4) + mem_indexes = req.alloc_kv_mem(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.get_cpu_token_indexes(0, 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 req.alloc_kv_mem(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) + req.alloc_kv_mem(alloc_token_num=4) assert req.hold_kv_len == 8 assert context.req_manager.mem_manager.alloc_sizes == [4, 4] assert context.req_manager.get_cpu_token_indexes(0, 0, 8).tolist() == list(range(8)) def test_reservation_fills_each_request_table_row(monkeypatch): - _, backend = _make_context(monkeypatch) + _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) + req0.alloc_kv_mem(alloc_token_num=4) + req1.alloc_kv_mem(alloc_token_num=4) assert generic_pre_process.g_infer_context.req_manager.get_cpu_token_indexes(0, 0, 4).tolist() == [0, 1, 2, 3] assert generic_pre_process.g_infer_context.req_manager.get_cpu_token_indexes(1, 0, 4).tolist() == [4, 5, 6, 7] assert req0.hold_kv_len == req1.hold_kv_len == 4 -def test_alloc_req_kv_mem_defers_gpu_table_update(monkeypatch): - context, backend = _make_context(monkeypatch) +def test_alloc_kv_mem_defers_gpu_table_update(monkeypatch): + context, _ = _make_context(monkeypatch) req = _make_req(0) - backend._alloc_req_kv_mem(req, alloc_token_num=4) - backend._alloc_req_kv_mem(req, alloc_token_num=4) + req.alloc_kv_mem(alloc_token_num=4) + req.alloc_kv_mem(alloc_token_num=4) assert context.req_manager.get_cpu_token_indexes(0, 0, 8).tolist() == list(range(8)) assert context.req_manager.req_to_token_indexs.eq(-1).all() @@ -409,7 +410,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step def test_decode_reserves_mtp_headroom(monkeypatch): - context, backend = _make_context(monkeypatch) + context, _ = _make_context(monkeypatch) req = _make_req(0) req.cur_kv_len = 3 req.hold_kv_len = 4 @@ -425,7 +426,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): 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) + req.alloc_kv_mem(alloc_token_num) model_input, run_reqs = generic_pre_process.prepare_decode_inputs([req]) @@ -437,7 +438,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): - context, backend = _make_context(monkeypatch) + context, _ = _make_context(monkeypatch) context.args.page_size = 1 context.req_manager.page_size = 1 context.req_manager.token_indexes_cpu = torch.full((2, 16), -1, dtype=torch.int32) @@ -454,7 +455,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): 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) + req.alloc_kv_mem(alloc_token_num) model_input, _ = generic_pre_process.prepare_decode_inputs([req]) 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 409a9656e1..24e8650f3e 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 @@ -186,7 +186,7 @@ def alloc(need_size): next_index += need_size return mem_indexes - def alloc_req_kv_mem(req, alloc_token_num): + def alloc_kv_mem(alloc_token_num): mem_indexes = alloc(alloc_token_num) new_hold_kv_len = req.hold_kv_len + alloc_token_num req.token_index.append_token_indexes(mem_indexes) @@ -206,7 +206,6 @@ def load(mem_indexes, **kwargs): mem_manager=SimpleNamespace(operator=operator), req_manager=make_cpu_index_manager(token_table=table, req_to_token_indexs=gpu_table, page_size=4), ), - _alloc_req_kv_mem=alloc_req_kv_mem, ) module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) module.args = SimpleNamespace(page_size=4, cpu_cache_token_page_size=132) @@ -229,6 +228,7 @@ def load(mem_indexes, **kwargs): token_hash_page_len_list=SimpleNamespace(get_all=lambda: [132]), ), sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + alloc_kv_mem=alloc_kv_mem, ) bind_cpu_indexes(req, backend.model.req_manager) req.args.req_token_index_mode = mode 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 index ad37c2c3a9..d31f2f1e17 100644 --- 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 @@ -45,7 +45,6 @@ def test_pd_decode_reserves_model_pages_but_transfers_only_logical_kv(monkeypatc backend.args = SimpleNamespace(pd_kv_page_size=4, page_size=4) backend.model = SimpleNamespace(req_manager=make_cpu_index_manager(token_table=table, page_size=1)) 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( @@ -71,6 +70,8 @@ def create_task(**kwargs): page_size=4, ) + req.alloc_kv_mem = lambda size: _reserve_into_table(table, req, alloc_sizes, next_index, size) + backend._decode_node_gen_trans_tasks(req) assert alloc_sizes == [8] @@ -250,7 +251,7 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat 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) + req.alloc_kv_mem = lambda size: _reserve_into_table(table, 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 From 7c14bca563911b7adbbd6f3484b5781669b6384e Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Fri, 23 Oct 2026 15:46:26 +0000 Subject: [PATCH 21/22] Refactor hybrid radix cache prefix restoration --- .../server/router/model_infer/infer_batch.py | 131 ++++++++---------- .../common/test_infer_batch_cache_release.py | 112 ++++++++++++++- 2 files changed, 166 insertions(+), 77 deletions(-) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index ea8033c455..c447f79f67 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -629,6 +629,7 @@ def _match_radix_cache(self): return def _hybrid_match_radix_cache(self): + """匹配完整 hash 页，恢复命中的 KV 和 checkpoint，再发布请求长度。""" assert ( g_infer_context.is_hybrid_att_model is True ), "current _hybrid_match_radix_cache only supports hybrid attention models, to do..." @@ -641,84 +642,15 @@ def _hybrid_match_radix_cache(self): match_block_num = match_tokens // hash_page_size block_hashs = block_hashs[:match_block_num] assert len(block_hashs) == self.shm_req.hybrid_token_hash_list.size - big_page_token_num = hash_page_size * self.args.linear_att_page_block_num - big_page_is_disable = big_page_token_num > self.args.max_req_total_len if enable_prompt_cache and match_tokens > 1 and len(block_hashs) > 0 and self.cur_kv_len == 0: input_token_ids = self.shm_req.shm_prompt_ids.arr[0 : self.get_cur_total_len()] key = torch.tensor(input_token_ids[0:match_tokens], dtype=torch.int64, device="cpu") assert len(key) == len(block_hashs) * hash_page_size - share_node, kv_len, value_tensor = g_infer_context.radix_cache.match_prefix( + shared_node, _, token_indexes = g_infer_context.radix_cache.match_prefix( key, block_hashs=block_hashs, update_refs=True ) - if share_node is not None: - # 与普通前缀匹配一致：新请求或已完成暂停释放的请求，其索引状态为空。 - assert self.tail_small_page_buffer_id is None - if share_node.is_big_page_node(): - # 大页匹配 - self._set_cache_prefix(share_node, value_tensor) - 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._set_cache_prefix(share_node, value_tensor) - assert self.tail_small_page_buffer_id is None - # 恢复 hybrid checkpoint - g_infer_context.req_manager.restore_small_page_state( - 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) - - # 将 对应的 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._set_cache_prefix( - shared_node=big_page_shared_node, - shared_prefix_indexes=value_tensor[:cur_big_page_tokens], - private_tail_indexes=tail_mems, - ) - else: - # 没有充足的token 容量时， 直接找到最接近的大页，进行大页恢复 - share_node = radix_cache.deref_to_first_big_page_node(node=share_node) - if share_node is not None: - assert share_node.is_big_page_node() - # 大页匹配 - ready_cache_len = share_node.node_prefix_total_len - self._set_cache_prefix(share_node, value_tensor[:ready_cache_len]) - 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 - ) + if shared_node is not None: + self._restore_hybrid_cache_match(shared_node, token_indexes) self.shm_req.shm_cur_kv_len = self.cur_kv_len @@ -729,10 +661,63 @@ def _hybrid_match_radix_cache(self): assert self.hold_kv_len % self.args.page_size == 0 if self.cur_kv_len == 0: - # 说明没有任何命中 g_infer_context.req_manager.init_hybrid_attention_state(req=self) return + def _restore_hybrid_cache_match(self, shared_node: HybridAttPagedTreeNode, token_indexes: torch.Tensor): + """直接复用共享前缀，或将小页尾部转为私有 KV；容量不足时回退到大页。""" + # 新请求或已完成暂停释放的请求，其索引和私有 checkpoint 状态为空。 + assert self.tail_small_page_buffer_id is None + big_page_token_num = self.args.linear_att_hash_page_size * self.args.linear_att_page_block_num + big_pages_disabled = big_page_token_num > self.args.max_req_total_len + if shared_node.is_big_page_node() or big_pages_disabled: + self._restore_hybrid_shared_prefix(shared_node, token_indexes) + return + + radix_cache = g_infer_context.radix_cache + shared_kv_len = shared_node.node_prefix_total_len + big_page_prefix_len = shared_kv_len // big_page_token_num * big_page_token_num + private_tail_len = shared_kv_len - big_page_prefix_len + + # 保持原有容量条件：只有可用量严格大于尾部需求时才重建小页尾部。 + if g_infer_context.get_can_alloc_token_num() <= private_tail_len: + big_page_node = radix_cache.deref_to_first_big_page_node(node=shared_node) + if big_page_node is not None: + assert big_page_node.is_big_page_node() + self._restore_hybrid_shared_prefix(big_page_node, token_indexes[: big_page_node.node_prefix_total_len]) + return + + assert private_tail_len % self.args.page_size == 0 + radix_cache.free_radix_cache_to_get_enough_token(need_token_num=private_tail_len) + private_tail_indexes = radix_cache.mem_manager.alloc(need_size=private_tail_len) + cached_tail_indexes = token_indexes[big_page_prefix_len:shared_kv_len] + radix_cache.mem_manager.operator.copy_mem_to_mem(cached_tail_indexes, private_tail_indexes) + self.prompt_selected_logprobs.copy_capture_slots_if_needed( + source_indexes=cached_tail_indexes, destination_indexes=private_tail_indexes + ) + + # 小页恢复接口从请求读取节点；必须先恢复 checkpoint，再解除小页引用。 + self.shared_kv_node = shared_node + g_infer_context.req_manager.restore_small_page_state(req=self) + self.shared_kv_node = None + big_page_node = radix_cache.deref_to_first_big_page_node(node=shared_node) + self._set_cache_prefix( + shared_node=big_page_node, + shared_prefix_indexes=token_indexes[:big_page_prefix_len], + private_tail_indexes=private_tail_indexes, + ) + + def _restore_hybrid_shared_prefix(self, shared_node: HybridAttPagedTreeNode, token_indexes: torch.Tensor): + """登记直接共享的 KV 前缀，并恢复同一节点的 checkpoint。""" + self._set_cache_prefix(shared_node, token_indexes) + assert self.tail_small_page_buffer_id is None + if shared_node.is_big_page_node(): + g_infer_context.req_manager.restore_big_page_state( + big_page_buffer_idx=shared_node.big_page_buffer_idx, req=self + ) + else: + g_infer_context.req_manager.restore_small_page_state(req=self) + def _set_cache_prefix( self, shared_node: Optional[Union[TreeNode, HybridAttPagedTreeNode]], diff --git a/unit_tests/common/test_infer_batch_cache_release.py b/unit_tests/common/test_infer_batch_cache_release.py index 50be25b5b9..7b0f6d97f7 100644 --- a/unit_tests/common/test_infer_batch_cache_release.py +++ b/unit_tests/common/test_infer_batch_cache_release.py @@ -1,6 +1,6 @@ """Exercise request release with real radix trees, KV allocation and checkpoint slot pools.""" -from types import MethodType, SimpleNamespace +from types import SimpleNamespace from unittest.mock import Mock import numpy as np @@ -98,7 +98,8 @@ def make_case(attention, mode, page_size, cache_case): context.req_manager = SimpleNamespace( page_size=page_size, req_to_token_indexs=torch.full((1, 32), -1, dtype=torch.int32, device="cuda") ) - req = SimpleNamespace( + req = infer_batch.InferReq.__new__(infer_batch.InferReq) + req.__dict__.update( req_idx=0, req_manager=context.req_manager, args=SimpleNamespace(req_token_index_mode=mode), @@ -334,6 +335,110 @@ def test_hybrid_release_before_first_forward_handles_reserved_capacity( assert small_slots.get_used_cache_num() == big_slots.get_used_cache_num() == 0 +@pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +@pytest.mark.parametrize("shared_prefix_len", [0, 8]) +@pytest.mark.parametrize("available_tokens", [3, 4, 5]) +def test_hybrid_small_page_match_preserves_tail_capacity_and_checkpoint_order( + make_release_case, monkeypatch, mode, shared_prefix_len, available_tokens +): + context, req, allocator, small_slots, big_slots, _ = make_release_case("hybrid_small", mode, 4, "uncached") + matched_len = shared_prefix_len + 4 + if shared_prefix_len == 0: + big_slots.free_state_cache(list(req.hybrid_len_to_big_page_id.values())) + req.hybrid_len_to_big_page_id.clear() + req.cur_kv_len = req.hybrid_cache_len = matched_len + small_checkpoint = req.tail_small_page_buffer_id + cached_indexes = req.token_index.get_cpu_token_indexes(0, matched_len) + released = [] + context.free_a_req_mem(released, req) + allocator.free(torch.cat(released)) + free_tokens_before_match = allocator.can_use_mem_size + + req.args = SimpleNamespace( + page_size=4, + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + max_req_total_len=32, + ) + req.sampling_param = SimpleNamespace(disable_prompt_cache=False) + req.get_cur_total_len = lambda: matched_len + 1 + req.shm_req.shm_prompt_ids = SimpleNamespace(arr=np.arange(matched_len + 1)) + block_hashs = list(range(1, matched_len // 4 + 1)) + req.shm_req.hybrid_token_hash_list = SimpleNamespace(size=len(block_hashs), get_all=lambda: block_hashs) + req.shm_req.prompt_cache_len = 0 + monkeypatch.setattr(context, "get_can_alloc_token_num", lambda: available_tokens) + + operations = [] + copy_kv = Mock(side_effect=lambda *args: operations.append("copy_kv")) + context.radix_cache.mem_manager.operator = SimpleNamespace(copy_mem_to_mem=copy_kv) + copy_logprobs = Mock(side_effect=lambda **kwargs: operations.append("copy_logprobs")) + req.prompt_selected_logprobs = SimpleNamespace(copy_capture_slots_if_needed=copy_logprobs) + + def restore_small_page(req): + assert req.shared_kv_node.small_page_buffer_idx == small_checkpoint + assert req.shared_kv_node.ref_counter > 0 + operations.append("restore_small") + + context.req_manager.restore_small_page_state = restore_small_page + context.req_manager.restore_big_page_state = Mock() + context.req_manager.init_hybrid_attention_state = Mock() + deref = context.radix_cache.deref_to_first_big_page_node + + def release_small_page_reference(node): + operations.append("deref") + return deref(node) + + monkeypatch.setattr(context.radix_cache, "deref_to_first_big_page_node", release_small_page_reference) + req._hybrid_match_radix_cache() + + rebuild_tail = available_tokens > 4 + expected_len = matched_len if rebuild_tail else shared_prefix_len + expected_indexes = cached_indexes[:shared_prefix_len] + if rebuild_tail: + source_indexes, private_indexes = copy_kv.call_args.args + torch.testing.assert_close(source_indexes, cached_indexes[shared_prefix_len:]) + assert private_indexes.numel() == 4 + assert not torch.isin(private_indexes, cached_indexes).any() + copy_logprobs.assert_called_once() + torch.testing.assert_close(copy_logprobs.call_args.kwargs["source_indexes"], source_indexes) + torch.testing.assert_close(copy_logprobs.call_args.kwargs["destination_indexes"], private_indexes) + expected_indexes = torch.cat([expected_indexes, private_indexes]) + assert operations == ["copy_kv", "copy_logprobs", "restore_small", "deref"] + else: + copy_kv.assert_not_called() + copy_logprobs.assert_not_called() + assert operations == ["deref"] + + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == expected_len + assert req.shm_req.prompt_cache_len == req.token_index.token_index_len == expected_len + torch.testing.assert_close(req.token_index.get_cpu_token_indexes(0, expected_len), expected_indexes) + assert allocator.can_use_mem_size == free_tokens_before_match - (4 if rebuild_tail else 0) + assert context.radix_cache.get_refed_tokens_num() == shared_prefix_len + if shared_prefix_len: + assert req.shared_kv_node.is_big_page_node() + assert req.shared_kv_node.node_prefix_total_len == shared_prefix_len + else: + assert req.shared_kv_node is None + if not rebuild_tail and shared_prefix_len: + context.req_manager.restore_big_page_state.assert_called_once_with( + big_page_buffer_idx=req.shared_kv_node.big_page_buffer_idx, req=req + ) + else: + context.req_manager.restore_big_page_state.assert_not_called() + if expected_len == 0: + context.req_manager.init_hybrid_attention_state.assert_called_once_with(req=req) + else: + context.req_manager.init_hybrid_attention_state.assert_not_called() + + released = [] + context.free_a_req_mem(released, req) + if released: + allocator.free(torch.cat(released)) + context.radix_cache.clear_tree_nodes() + assert allocator.can_use_mem_size == allocator.size + assert small_slots.get_used_cache_num() == big_slots.get_used_cache_num() == 0 + + @pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) @pytest.mark.parametrize("new_tokens", [0, 1]) @pytest.mark.parametrize("evict_checkpoint", [False, True]) @@ -374,8 +479,7 @@ def test_hybrid_release_skips_insert_for_shared_small_prefix( context.free_a_req_mem(released, req) allocator.free(torch.cat(released)) cache = context.radix_cache - req._set_cache_prefix = MethodType(infer_batch.InferReq._set_cache_prefix, req) - infer_batch.InferReq._hybrid_match_radix_cache(req) + req._hybrid_match_radix_cache() node = req.shared_kv_node shared_len = req.cur_kv_len checkpoint = node.small_page_buffer_idx From a5342fdbcc1c0c17407140d894aa3e029964a43a Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Sat, 24 Oct 2026 02:57:33 +0000 Subject: [PATCH 22/22] Reduce persistent pinned memory for request token indexes Store full request indexes in pageable CPU memory and copy pending uploads into shared pinned slab slices. Encapsulate the shared return buffer state in the base allocator. Update lifecycle tests, allocation benchmarks, CLI help, and design notes covering pinned memory scaling and remaining limits. --- docs/kv_cache_page_size.md | 141 ++++++++- lightllm/server/api_cli.py | 6 +- .../router/model_infer/infer_req_ext1.py | 69 ++--- .../benchmark_cpu_pinned_alloc_keepalive.py | 289 ++++++++++++++++++ .../benchmark_req_token_index_alloc.py | 27 +- unit_tests/common/test_req_token_index.py | 99 ++++-- 6 files changed, 528 insertions(+), 103 deletions(-) create mode 100644 test/benchmark/benchmark_cpu_pinned_alloc_keepalive.py diff --git a/docs/kv_cache_page_size.md b/docs/kv_cache_page_size.md index e6c6ea49c4..6508363edb 100644 --- a/docs/kv_cache_page_size.md +++ b/docs/kv_cache_page_size.md @@ -49,25 +49,25 @@ Helper 保存该流，后续不随调用点的当前流切换： | `--req_token_index_mode` | 存储与优点 | 代价 | | --- | --- | --- | -| `pin_mem_cache`（默认） | `PinMemCacheReqTokenIndex` 按请求直接申请 pinned CPU token 缓存；CPU 读取无需 D2H，GPU 更新可批量提交 | 占用 pinned CPU 内存，需要维护上传进度 | +| `pin_mem_cache`（默认） | `PinMemCacheReqTokenIndex` 按请求申请普通 CPU token 缓存，上传时仅为待提交区间创建 pinned 快照；CPU 读取无需 D2H，GPU 更新可批量提交 | 每次上传增加一次 CPU 拷贝和临时 pinned 区间申请，需要维护上传进度 | | `pin_mem_page_cache` | `PinMemPageCacheReqTokenIndex` 每个 KV 页仅保存一个 pinned CPU 首 token 索引，常驻索引项数约缩小为原来的 `1 / page_size` | CPU 读取和 GPU 上传时需要展开 token 索引，适用于 `page_size > 1` | | `pure_gpu` | `PureGpuReqTokenIndex` 直接维护 GPU 请求表，无常驻 CPU 索引缓存 | CPU 读取需要 D2H 并等待结果，频繁的 CPU 索引操作可能更慢 | 三种模式均保留相同的 GPU 请求表，显存占用基本相同。追加后的所需长度为 `end = token_index_len + mem_indexes.numel()`。 完整 token 模式的 CPU 缓存首次按 `max(input_len + 2K, end)` 向上对齐到 1K token 申请，为输出预留空间；后续容量不足时，在本次所需长度 `end` 上多预留 2K token，再向上对齐到 1K， -容量上限为 GPU 请求表行长度。两种 CPU 缓存模式均直接使用 `torch.empty(..., pin_memory=True, requires_grad=False)` -为每个请求独立申请常驻存储；完整模式保存每个 token，页模式在首次追加时按 GPU 请求表的最大长度 +容量上限为 GPU 请求表行长度。完整模式使用普通 CPU Tensor 保存每个 token，首次分配及扩容均不申请 pinned 内存。 +页模式的常驻存储仍使用 `torch.empty(..., pin_memory=True, requires_grad=False)`，在首次追加时按 GPU 请求表的最大长度 一次性申请 `ceil(max_sequence_length / page_size)` 个页首索引，后续追加无需扩容或复制旧缓存。 页模式继承完整模式的 GPU 上传、释放快照和 reset 流程，覆盖 CPU 索引的存储和展开。 压缩依赖已有的页内连续索引约束；追加完整页，输入从下标 0 按 page_size 取页首。 页模式读取支持页内任意 token 区间：内部展开覆盖区间的完整页，再按页内偏移截取精确范围。 CPU 快照独立于压缩缓存，PD 尾段传输和最终 metadata 无需自行处理页对齐。 缓冲区在首次非空写入时按需申请，重置时保留输入长度， -丢弃已申请的缓冲区引用，由 PyTorch 在引用和异步传输依赖结束后回收。 -页模式上传时先将待提交区间展开到临时 pinned token 索引，使用 `PinMemPageCacheReqTokenIndex._alloc_pinned_tensor` -按 1M token（4 MiB）申请 int32 slab。该接口及共享状态仅属于页缓存类，保存完整 buffer 和下一个空闲位置 `_pinned_offset`， -每次只创建返回值这一个切片。 +丢弃已申请的常驻缓冲区引用。上传使用独立暂存区，其引用和异步传输依赖由上传路径管理。 +完整模式的上传快照和页模式的展开暂存区统一使用基类 `ReqTokenIndex._alloc_return_pinned_tensor()`， +按 1M token（4 MiB）申请 int32 slab。函数内部通过函数属性 `return_buffer_state` 惰性保存完整 slab 和下一个空闲 offset， +两种模式共享同一分配进度，每次只创建返回值这一个切片。返回区间互不重叠，但可持有同一块底层 storage。 offset 只前进，不回绕；空间不足时换用新 slab 并从零开始，请求 reset 不重置共享 offset。 超出 slab 容量的申请独立分配，空申请不改变 slab 或 offset。旧 slab 的 storage 在所有切片引用和异步传输依赖结束后 由 PyTorch 回收。 @@ -77,7 +77,8 @@ offset 只前进，不回绕；空间不足时换用新 slab 并从零开始， 统一接口为 `append_token_indexes(mem_indexes)`、`get_cpu_token_indexes(start, end)`、 `get_free_token_indexes(start, end)` 和 `get_gpu_token_indexes(start, end)`，均使用 token 单位， 追加接口将完整输入接到已记录索引的尾部，非空输入在三种模式下均断言为 CPU Tensor；空输入直接返回。 -CPU 缓存模式直接复制到自身的 pinned 缓冲区，无需先 pin 输入；GPU 模式仅在输入未 pinned 时内部转换，已有 pinned 输入直接复用。 +CPU 缓存模式直接复制到自身的常驻缓冲区，无需先 pin 输入；完整模式为普通 CPU 内存，页模式为 pinned 内存。 +GPU 模式仅在输入未 pinned 时内部转换，已有 pinned 输入直接复用。 前缀匹配和索引重建路径直接传入 CPU 索引，不再自行 pin。读取接口均支持页内任意区间；提交给 allocator 的释放区间仍须覆盖完整页。 CPU getter 返回独立快照，缓存、PD 元数据和异步任务可跨请求释放持有；GPU getter 返回当前请求的 GPU 表视图，不应跨请求槽位复用持有。Helper 保存外部传入的索引流，GPU 索引读写固定在该流， @@ -105,13 +106,15 @@ CPU 缓存模式的追加接口只更新 CPU 索引和 `token_index_len`，不 首次前缀匹配时索引为空；暂停恢复前已由统一释放流程重置，两处匹配入口直接追加命中前缀，不重复 `reset()`。 释放或暂停请求时也由索引扩展的 `reset()` 统一清零两个长度。 CPU 模式重置时将 `_cpu_buffer` 设为 `None`，下一次追加重新申请，随后从零开始提交 GPU 更新。 -单独的 GPU getter 上传完整缓存视图或页缓存展开得到的暂存区：普通追加不会改写已有区间，扩容和重置使用新存储， -PyTorch `copy_(non_blocking=True)` 会登记 pinned 源的流使用，旧存储在 H2D 完成前不会被回收复用。 -因此不再为单个请求上传额外申请 pinned 快照或执行一次 CPU 拷贝，也不需要在 reset 中同步 CUDA。 +单独的 GPU getter 通过 `get_upload_indexes()` 准备独立 pinned 暂存区：完整模式复制待提交区间，页模式展开压缩索引。 +上传源不借用常驻缓存，常驻缓存的扩容和重置不影响在途 H2D，也不需要在 reset 中同步 CUDA。 +PyTorch `copy_(non_blocking=True)` 会登记 pinned 源的流使用，暂存区在 H2D 完成前不会被回收复用。 `get_upload_indexes()` 不接受区间参数，CPU 模式返回内部计算的 -`[gpu_index_submitted_len, token_index_len)` 的 pinned token 索引，完整模式返回缓存视图，页模式按需展开， -不改变提交进度；无待上传内容时返回空 tensor， +`[gpu_index_submitted_len, token_index_len)` 的 pinned token 索引。完整模式每次通过基类接口申请恰好覆盖该区间的切片， +同步复制普通 CPU 缓存中的索引后返回；连续调用返回互不重叠的快照，不覆写上次结果，不包含容量预留区间。 +切片的逻辑长度只覆盖待上传索引，其底层 storage 通常为共享的 4 MiB slab。 +页模式按需展开。两种模式均不改变提交进度；无待上传内容时返回空 tensor， 纯 GPU 模式返回 `None`。CPU 缓存模式的 GPU getter 校验 `end <= token_index_len` 后按需提交完整待上传区间， 纯 GPU 模式校验 `end <= gpu_index_submitted_len`；读取均不允许推进超出已有索引的提交进度。 批量入口不会替 GPU 模式补写提交进度；若缺少待上传数据，由后续拷贝操作报错，提交进度保持不变。 @@ -124,7 +127,7 @@ PyTorch `copy_(non_blocking=True)` 会登记 pinned 源的流使用，旧存储 并传入批量拷贝函数；backend 不再读写索引流、GPU 行和提交进度。 统一入口检查所有请求共用一个索引流，准备源、目标视图列表，在索引流提交拷贝， 仅在整批提交成功后更新提交进度。单请求 GPU getter 复用该协议，默认通过 PyTorch `copy_` 提交； -backend 路径使用 `cudaMemcpyBatchAsync`，不申请聚合缓冲区或额外复制 CPU 索引。 +backend 路径使用 `cudaMemcpyBatchAsync`，直接提交索引对象准备的 pinned 快照，不再额外聚合或复制。 返回前仅在调用流不同于索引流时排入等待，没有新上传时也承接已有写入的依赖。 backend 不再持有上传流，也不逐条交接多条索引流。 原生 API 不接受 legacy 默认流，适配由 `PinnedH2DBatchCopier` 内部处理：按设备延迟创建并复用非默认上传流， @@ -144,8 +147,8 @@ backend 不再持有上传流，也不逐条交接多条索引流。 后台 `event.synchronize()` 休眠等待并释放 GIL，完成后清空源列表，不向 Future 返回 Tensor； 主推理线程不等待后台任务，默认索引流由封装承接内部上传流的依赖，调用流按需等待索引流。 索引扩展不维护逐请求事件，backend 不维护保活列表，也不需要额外的小片段 H2D 或 C++ 扩展。 -该方式只保护引用生命周期，不阻止 CPU 主动改写数据。请求缓存只追加，reset 和扩容换用新存储， -因此可直接保活缓存视图；上传期间不能改写已提交的区间，无需独立聚合快照。 +该方式只保护引用生命周期，不阻止 CPU 主动改写上传源。完整模式的 pinned 快照和页模式的展开暂存区 +均独立于常驻缓存，reset 和扩容不会改写它们；调用方在传输完成前不能主动修改返回的 pinned Tensor。 单线程按提交顺序处理；多上传流共用线程池时可能延后回收，不改变 GPU 工作顺序。 正常进程退出时由线程池等待未完成任务，不取消保活任务。 @@ -232,3 +235,109 @@ PyTorch 文档介绍了使用 pinned CPU Tensor 配合 `to(..., non_blocking=Tru `cudaMemcpyBatchAsync` 则是多个索引区间的批量提交选择。 参见 [PyTorch 2.11：Use pinned memory buffers](https://docs.pytorch.org/docs/2.11/notes/cuda.html#use-pinned-memory-buffers) 和 [PyTorch：non_blocking 与 pin_memory 使用指南](https://docs.pytorch.org/tutorials/intermediate/pinmem_nonblock.html#other-copy-directions-gpu-cpu-cpu-mps)。 + +## 附录：分离常驻索引与 pinned 上传缓冲区 + +### 设计动机与容量估算 + +请求索引有两种不同的生命周期：完整映射需要保存到请求结束或重建，用于前缀匹配、缓存管理和资源释放； +H2D 源只需保存到本次上传完成。将完整映射长期保存在 pinned 内存中，会让短期传输需要的锁页资源 +随全部存活请求的索引容量增长。在当前离散 GPU 部署中，大量长期锁页会增加主机内存压力； +NVIDIA 也明确指出，过量分配 page-locked memory 会影响系统整体性能。 +参见 [NVIDIA CUDA FAQ：Host/Device transfer](https://developer.nvidia.com/cuda/faq)。 + +面向 V4.1 这类模型的后续部署，需要提前考虑上下文长度、并发请求数和推理进程数同时增长的情况。 +这里按部署容量分析，不假定 V4.1 的具体上下文上限、并发量或实测内存占用。 +对每个推理进程，设已分配缓存的请求数为 `R`，请求 `i` 的完整 token 缓存容量为 `C_i`， +每个索引为 `int32`，则完整索引的逻辑字节数为 `4 × Σ C_i`。 +如果每个请求各自申请 pinned 缓存，还需要计算底层分配块的容量取整。 +本次核对的 PyTorch 2.11 host caching allocator 将申请字节数向上取整到 2 的幂， +因此这部分常驻 pinned 块约占 `Σ round_pow2(4 × C_i)`，扩容期间的新旧块和空闲缓存另计。 +这是该版本 pinned allocator 的行为，不是所有 CPU Tensor 的通用分配规则。 +参见 [PyTorch 2.11：host_memory_stats](https://docs.pytorch.org/docs/2.11/generated/torch.cuda.memory.host_memory_stats.html)。 + +例如，假设输入长 `128 × 1024` token，GPU 请求表行长度允许再预留 2K token， +当前完整缓存策略会申请 `130 × 1024` 个索引，逻辑大小为 520 KiB；若直接申请 pinned tensor， +对应块会取整为 1 MiB。每进程 128 个这样的存活请求、同机 8 个进程分别保存这些映射时， +仅常驻索引就会占用约 1 GiB pinned 块，尚未计入上传、旧块和其他模块。 +多进程占用应按各进程的实际请求集合求和；TP 进程可能保存相同映射，DP 进程则应按各自请求数计算。 + +默认 `pin_mem_cache` 现将这部分长期存储改为普通 CPU Tensor，上述算例的常驻缓存逻辑大小合计为 520 MiB， +另加临时 pinned 上传区间。Linux 默认 CPU allocator 按请求大小进行对齐分配， +没有上述 host caching allocator 的统一 2 的幂取整步骤；系统 malloc 仍可能存在对齐、缓存和碎片开销。 +当前代码自己的 1K 对齐、2K 预留以及扩容复制也仍然存在，普通 CPU 缓存不能当作零成本。 +参见 [PyTorch 2.11：alloc_cpu 实现](https://github.com/pytorch/pytorch/blob/v2.11.0/c10/core/impl/alloc_cpu.cpp)。 + +### 为什么采用普通 CPU 常驻缓存和共享 pinned slab + +完整模式的 `get_upload_indexes()` 只复制 `[gpu_index_submitted_len, token_index_len)`， +已经提交的历史索引和 CPU 缓存预留空间都不会重复进入上传快照。 +这样保留了 CPU 读取无需 D2H 的能力，同时将 pinned 源的生命周期缩短到本次传输。 +代价是每次上传增加一次同步 CPU 拷贝；前缀重建后需要重新上传完整映射，不能总按少量 decode 增量估算。 + +如果每次快照都单独调用 `torch.empty(..., pin_memory=True)`，小额上传仍会频繁进入 host allocator， +且不同大小的快照分别受容量取整影响。`_alloc_return_pinned_tensor()` 为常规申请维护一个 +`1024 × 1024` 个 int32 元素的 slab，让多个请求取得互不重叠的区间，摊薄底层分配和取整开销。 +4 MiB 是当前的小额分配与保留粒度取舍，并非所有模型、并发量和上传大小下的最优值。 +空申请不创建 slab，单次超过 slab 容量时独立申请，避免把长期保留的工作块扩成一个大块。 + +完整模式将索引复制进这个区间，页模式将压缩页索引直接展开进这个区间； +两者共享相同的分配与生命周期约定。返回值是一次上传使用的快照，调用方不应把它保存为长期索引缓存。 +多个快照可以共享 storage，但不得共享可写区间。offset 只前进，空间不足时换新 slab， +请求 reset 不回绕 offset，避免后续 CPU 写入覆盖尚未完成的 H2D 源。 +PyTorch `copy_` 的 allocator 跟踪和原生 batch 路径的事件保活，继续负责传输结束前不回收源存储。 + +### 职责划分与状态归属 + +| 组件 | 负责的内容 | +| --- | --- | +| `PinMemCacheReqTokenIndex` | 普通 CPU 完整映射的追加、容量管理、快照生成和重置 | +| `PinMemPageCacheReqTokenIndex` | 页首映射的存储及 token 索引展开，复用公共上传协议 | +| `ReqTokenIndex` | 上传区间与提交进度、索引流依赖、临时 pinned 返回区间分配 | +| backend | 选择本批请求并调用统一提交入口 | +| `PinnedH2DBatchCopier` | CUDA 拷贝提交、原生接口适配及异步源保活 | + +这种划分让常驻存储策略可以变化，而 backend 无需了解 CPU 缓存布局或维护上传进度。 +`get_upload_indexes()` 只准备数据，`submit_pending_uploads()` 在整批提交成功后推进进度， +流完成与引用保活由传输层处理；缓存扩容和 reset 无需等待上传完成。 +页模式继承完整模式时覆盖所有依赖压缩布局的追加、读取和分配操作，复用 GPU getter、释放入口及 reset。 +保留 `PinMemCacheReqTokenIndex` 和 `pin_mem_cache` 名称以兼容现有调用和配置；名称中的 PinMem +现在对应上传缓冲区，常驻完整缓存的实际类型以实现、类说明和 CLI 帮助为准。 + +`return_buffer_state = (buffer, offset)` 放在分配函数对象上，由函数内部惰性建立和更新， +使两个字段的使用范围集中在唯一分配入口。它仍是进程内共享的持久状态，函数返回后不会消失， +移动字段本身不减少内存；主要收益来自常驻缓存改用普通 CPU 内存及上传区间共享 slab。 +该入口依赖同一推理线程串行申请，未提供多线程分配保护；未来若增加并行申请线程，需重新设计同步或状态归属。 + +### 内存边界与后续评估 + +4 MiB 仅表示此分配函数当前持有的一块 slab，不能解释为整个进程的 pinned 内存上限。 +上传缓冲区的实际占用还包括旧 slab 的存活切片、在途传输、超过 slab 的独立申请，以及 PyTorch 缓存的空闲块。 +一个小切片也会保留整块 slab；GPU 上传积压、后台回收延迟或调用方长期持有返回值，都会延长整块存储的存活时间。 +当前没有对同时存活的 slab 数量设置硬上限，也没有上传积压的内存背压。 +KV allocator 的返回缓冲区和其他模块的 pinned 申请应独立计入整机预算，不能只统计这个函数。 + +`pin_mem_page_cache` 的常驻缓存目前仍使用 pinned 内存，并在首次追加时按 GPU 行长度 `S` +申请 `ceil(S / page_size)` 个页首索引，而不是按本次实际长度增长。 +其常驻 pinned 块约为 `R × round_pow2(4 × ceil(S / page_size))`，上传暂存区另计。 +`page_size=1` 时没有压缩收益；即使页较大，配置很大的最大序列长度且并发很多时，也仍需评估这部分预分配。 +默认完整模式的改进不能代表页模式已经消除了长期锁页。若该项成为瓶颈，后续可评估普通 CPU 页缓存或按需增长， +并测量对页索引展开耗时的影响。 + +本次设计主要降低完整索引的长期 pinned 占用。首次长 prefill 或大规模重建时，普通 CPU 缓存与 pinned 快照会同时存在， +因此不保证主机总内存峰值下降，也不保证推理吞吐提升。 +GPU `req_to_token_indexs` 仍为 token 级表，约占 `4 × (max_request_num + 1) × S` 字节； +模型权重、KV Cache 和该 GPU 表的显存都不因这项改动缩小。 + +评估 V4.1 等模型的实际部署时，应同时观察长 prefill、稳态 decode、前缀重建和上传积压场景， +按主机汇总各推理进程的 CPU RSS、pinned 块占用及峰值，并测量索引准备、上传和端到端延迟。 +PyTorch 2.11 的 `host_memory_stats()` 中，`allocated_bytes.current` 包含活跃块和缓存块， +`active_bytes.current` 表示尚未变成可复用空闲块的容量，`num_host_alloc` 用于观察底层分配次数； +这些是进程级 host allocator 指标，不能直接归因到单个索引类，也不能用 GPU 的 `memory_allocated()` 替代。 +各 bucket 的峰值是汇总近似值，取整与指标定义参见上述 PyTorch 文档。 + +仓库中的 [CPU/pinned 保活分配基准](../test/benchmark/benchmark_cpu_pinned_alloc_keepalive.py) +保留所有申请结果，可观察不同大小的直接分配耗时和 pinned 块取整； +[索引分配基准](../test/benchmark/benchmark_req_token_index_alloc.py) 比较直接申请与 slab 切片路径。 +两者都不测量完整上传流程，前者也不写入普通 CPU Tensor 的全部页面，不能据此直接判断真实工作负载的 RSS 或推理性能。 +现有 [索引测试](../unit_tests/common/test_req_token_index.py) 覆盖快照独立性、跨模式共享、换块及在途 H2D 源的生命周期。 diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 4ad655f451..b4984088f3 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -502,9 +502,9 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: type=str, choices=["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"], default="pin_mem_cache", - help="""Storage mode for request token indexes. pin_mem_cache (default): allocates a pinned CPU cache - per request, allowing CPU reads without D2H transfers and batched GPU uploads, at the cost of pinned - memory usage and upload progress tracking. pin_mem_page_cache: stores one pinned CPU index per KV page, + help="""Storage mode for request token indexes. pin_mem_cache (default): allocates a pageable CPU cache + per request and copies pending indexes to a fresh pinned slab slice for each upload, allowing CPU reads without + D2H transfers and batched GPU uploads. pin_mem_page_cache: stores one pinned CPU index per KV page, reducing persistent index memory but expanding indexes for token reads and GPU uploads; page_size=1 provides no compression. pure_gpu: maintains the GPU request table directly without a persistent CPU index cache; CPU reads require D2H transfers and may incur synchronization waits. diff --git a/lightllm/server/router/model_infer/infer_req_ext1.py b/lightllm/server/router/model_infer/infer_req_ext1.py index 2dcc694b25..7235bab546 100644 --- a/lightllm/server/router/model_infer/infer_req_ext1.py +++ b/lightllm/server/router/model_infer/infer_req_ext1.py @@ -33,7 +33,7 @@ def __init__(self, infer_req: "InferReq", index_stream: torch.cuda.Stream): def append_token_indexes(self, mem_indexes: torch.Tensor): """将 token 索引追加到当前记录的末尾，追加长度由 mem_indexes.numel() 决定。 - 非空输入必须是 CPU Tensor；CPU 缓存模式直接复制到自己的 pinned 缓存。 + 非空输入必须是 CPU Tensor；CPU 缓存模式直接复制到自己的常驻缓存。 pure_gpu 按需 pin 后异步上传；已 pinned 的输入在 H2D 完成前不得改写，但可释放引用。 pin_mem_page_cache 仅接受从页边界追加完整页，调用方保证每页内物理 token 索引连续。 前缀替换先 reset，再按顺序追加重建。 @@ -54,8 +54,8 @@ def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: def get_upload_indexes(self) -> Optional[torch.Tensor]: """返回 [gpu_index_submitted_len, token_index_len) 的 pinned CPU token 索引，不更新提交进度。 - 完整 CPU 缓存借用内部视图，页缓存按需展开到独立暂存区，调用方均可直接批量上传。 - 缓存只追加，reset/扩容换用新存储，上传期间已有区间不变；GPU 模式返回 None。 + 完整 CPU 缓存每次复制到新申请的 pinned 区间，页缓存按需展开到独立暂存区。 + 返回值不借用常驻缓存，调用方可直接批量上传；GPU 模式返回 None。 """ pass @@ -122,9 +122,29 @@ def _wait_for_index_stream(self) -> None: if current_stream != self._index_stream: current_stream.wait_stream(self._index_stream) + @staticmethod + def _alloc_return_pinned_tensor(need_size: int) -> torch.Tensor: + """申请 int32 token 索引的 pinned 区间，slab 容量为 1M token（4 MiB）。""" + buffer_token_num = 1024 * 1024 + if need_size == 0 or need_size > buffer_token_num: + return torch.empty(need_size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True) + # 在函数对象上惰性保存 slab 和 offset，所有索引实例共享，由推理线程串行申请。 + allocator = ReqTokenIndex._alloc_return_pinned_tensor + buffer, start = getattr(allocator, "return_buffer_state", (None, 0)) + end = start + need_size + if buffer is None or end > buffer_token_num: + buffer = torch.empty( + buffer_token_num, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True + ) + start = 0 + end = need_size + # 每次只创建返回值这一个视图；offset 只前进，换新 slab 时才归零，不覆写在途 H2D 的源。 + allocator.return_buffer_state = (buffer, end) + return buffer[start:end] + class PinMemCacheReqTokenIndex(ReqTokenIndex): - """按请求按需申请完整 token 级 pinned CPU 缓存，GPU 更新可批量提交。""" + """按请求保存普通 CPU token 缓存，上传时复制待提交区间到独立 pinned 区间。""" def __init__(self, infer_req: "InferReq", index_stream: torch.cuda.Stream): super().__init__(infer_req, index_stream) @@ -159,17 +179,19 @@ def get_gpu_token_indexes(self, start: int, end: int) -> torch.Tensor: return self.gpu_indexes[start:end] def get_upload_indexes(self) -> torch.Tensor: - if self._cpu_buffer is None: - return torch.empty(0, dtype=self.gpu_indexes.dtype, device="cpu", pin_memory=True, requires_grad=False) - return self._cpu_buffer[self.gpu_index_submitted_len : self.token_index_len] + """每次为待上传区间创建 pinned 快照，不保留常驻缓存或更新提交进度。""" + start, end = self.gpu_index_submitted_len, self.token_index_len + upload_indexes = self._alloc_return_pinned_tensor(end - start) + if start < end: + upload_indexes.copy_(self._cpu_buffer[start:end]) + return upload_indexes def get_free_token_indexes(self, start: int, end: int) -> torch.Tensor: """返回用于释放的 CPU token 索引快照。""" return self.get_cpu_token_indexes(start, end) def reset(self): - # 丢弃旧缓存引用，下一次追加重新申请，避免 CPU 覆写仍被在途 H2D 读取的内存。 - # 旧存储的回收时机由 PyTorch pinned allocator 跟踪，不需要在这里同步 CUDA。 + # 上传使用独立 pinned 暂存区，常驻缓存可直接释放，无需等待在途 H2D。 self.token_index_len = self.gpu_index_submitted_len = 0 self._cpu_buffer = None @@ -186,9 +208,7 @@ def _reserve(self, end: int): reserve_len = end + 2 * 1024 # 按 1K token 向上对齐，最多申请到 GPU 请求表允许的长度。 new_capacity = min(self.gpu_indexes.numel(), (reserve_len + 1023) // 1024 * 1024) - new_buffer = torch.empty( - new_capacity, dtype=self.gpu_indexes.dtype, device="cpu", pin_memory=True, requires_grad=False - ) + new_buffer = torch.empty(new_capacity, dtype=self.gpu_indexes.dtype, device="cpu", requires_grad=False) if self.token_index_len: new_buffer[: self.token_index_len].copy_(self._cpu_buffer[: self.token_index_len]) self._cpu_buffer = new_buffer @@ -199,9 +219,6 @@ class PinMemPageCacheReqTokenIndex(PinMemCacheReqTokenIndex): # 同一进程使用统一的 page_size，所有请求共享首次初始化的页内偏移。 _page_offsets: Optional[torch.Tensor] = None - # 同一推理进程的页缓存实例共享 slab 和空闲 offset，由推理线程串行申请。 - _pinned_buffer: Optional[torch.Tensor] = None - _pinned_offset = 0 def __init__(self, infer_req: "InferReq", index_stream: torch.cuda.Stream): super().__init__(infer_req, index_stream) @@ -255,7 +272,7 @@ def _expand_token_indexes(self, start: int, end: int, pin_memory: bool) -> torch token_num = (last_page - first_page) * self.page_size if pin_memory: # 上传使用一次性的 slab 区间，后续追加、扩容或 reset 均不会覆写在途 H2D 的源。 - token_indexes = self._alloc_pinned_tensor(token_num) + token_indexes = self._alloc_return_pinned_tensor(token_num) else: token_indexes = torch.empty(token_num, dtype=self.gpu_indexes.dtype, device="cpu", requires_grad=False) page_starts = self._cpu_buffer[first_page:last_page] @@ -265,26 +282,6 @@ def _expand_token_indexes(self, start: int, end: int, pin_memory: bool) -> torch offset = start % self.page_size return token_indexes[offset : offset + end - start] - @staticmethod - def _alloc_pinned_tensor(need_size: int) -> torch.Tensor: - """申请 int32 token 索引的 pinned 区间，slab 容量为 1M token（4 MiB）。""" - buffer_token_num = 1024 * 1024 - if need_size == 0 or need_size > buffer_token_num: - return torch.empty(need_size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True) - buffer = PinMemPageCacheReqTokenIndex._pinned_buffer - start = PinMemPageCacheReqTokenIndex._pinned_offset - end = start + need_size - if buffer is None or end > buffer_token_num: - buffer = torch.empty( - buffer_token_num, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - PinMemPageCacheReqTokenIndex._pinned_buffer = buffer - start = 0 - end = need_size - # 每次只创建返回值这一个视图；offset 只前进，换新 slab 时才归零，不覆写在途 H2D 的源。 - PinMemPageCacheReqTokenIndex._pinned_offset = end - return buffer[start:end] - class PureGpuReqTokenIndex(ReqTokenIndex): """直接维护 GPU 请求表，不持有 CPU 索引缓存；CPU 读取会等待 D2H 完成。""" diff --git a/test/benchmark/benchmark_cpu_pinned_alloc_keepalive.py b/test/benchmark/benchmark_cpu_pinned_alloc_keepalive.py new file mode 100644 index 0000000000..02e7bad583 --- /dev/null +++ b/test/benchmark/benchmark_cpu_pinned_alloc_keepalive.py @@ -0,0 +1,289 @@ +"""Compare torch.empty allocation latency while retaining every tensor. + +Example (select an idle GPU and a CPU on the same NUMA node): + CUDA_VISIBLE_DEVICES=4 python test/benchmark/benchmark_cpu_pinned_alloc_keepalive.py \ + --cpu 48 --malloc-usable-size --output-dir /tmp/cpu_pinned_alloc + +Each mode/repeat runs in a fresh process. Warmup and measured tensors, across +all sizes, remain alive until process exit. No fill, H2D, or release is timed. +Ordinary CPU empty allocations need not fault in all physical pages. This is +an allocation API benchmark, not a comparison of ready-to-use memory bandwidth. +""" + +import argparse +import csv +import ctypes +import gc +import json +import os +from pathlib import Path +import platform +import random +import statistics +import subprocess +import sys +import time + + +MIB = 1024 * 1024 +DEFAULT_SIZES = [ + 64, + 256, + 1024, + 4096, + 4097, + 16384, + 65536, + 262144, + MIB, + MIB + 1, + 4 * MIB, + 16 * MIB, + 64 * MIB, + 64 * MIB + 1, +] + + +def percentile(values, fraction): + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def rss_bytes(): + return int(Path("/proc/self/statm").read_text().split()[1]) * os.sysconf("SC_PAGE_SIZE") + + +def sample_count(size, args): + # Both modes use the same count. Budget using pinned allocator rounding. + rounded = 1 << (size - 1).bit_length() + return max(args.min_count, min(args.max_count, args.case_mib * MIB // rounded)) + + +def run_worker(args): + if args.cpu is not None: + os.sched_setaffinity(0, {args.cpu}) + import torch + + torch.set_num_threads(1) + torch.cuda.init() + pinned = args.mode == "pinned" + # Retain warmup too: initialization is excluded without returning any blocks. + warmup = [torch.empty(4096, dtype=torch.uint8, device="cpu", pin_memory=pinned) for _ in range(32)] + usable_size = None + if args.malloc_usable_size and not pinned: + if platform.libc_ver()[0] != "glibc" or os.environ.get("LD_PRELOAD"): + raise RuntimeError("--malloc-usable-size requires the default glibc CPU allocator without LD_PRELOAD") + # Diagnostic only; never call this on CUDA host allocations. The caller + # must use a PyTorch build with the default system CPU allocator. + usable_size = ctypes.CDLL(None).malloc_usable_size + usable_size.argtypes = [ctypes.c_void_p] + usable_size.restype = ctypes.c_size_t + + sizes = list(args.sizes) + random.Random(args.repeat).shuffle(sizes) + kept = [] + rows = [] + gc.disable() + for size in sizes: + count = sample_count(size, args) + tensors = [None] * count + durations_ns = [0] * count + before_rss = rss_bytes() + before_stats = torch.cuda.memory.host_memory_stats() if pinned else None + for index in range(count): + start = time.perf_counter_ns() + tensor = torch.empty(size, dtype=torch.uint8, device="cpu", pin_memory=pinned, requires_grad=False) + end = time.perf_counter_ns() + tensors[index] = tensor + durations_ns[index] = end - start + kept.append(tensors) + after_stats = torch.cuda.memory.host_memory_stats() if pinned else None + after_rss = rss_bytes() + # Validation and diagnostics are outside the timed region. + assert all(t.nbytes == size and t.untyped_storage().nbytes() == size for t in tensors) + assert len({t.data_ptr() for t in tensors}) == count + if pinned: + actual_allocs = after_stats["num_host_alloc"] - before_stats["num_host_alloc"] + assert actual_allocs == count, (size, count, actual_allocs) + block_bytes = (after_stats["allocated_bytes.current"] - before_stats["allocated_bytes.current"]) // count + assert block_bytes == 1 << (size - 1).bit_length(), (size, block_bytes) + else: + actual_allocs = block_bytes = None + usable = [usable_size(t.data_ptr()) for t in tensors] if usable_size else [] + assert not usable or min(usable) >= size + rows.append( + { + "mode": args.mode, + "repeat": args.repeat, + "requested_bytes": size, + "count": count, + "latency_ns": durations_ns, + "median_us": statistics.median(durations_ns) / 1000, + "mean_us": statistics.mean(durations_ns) / 1000, + "p95_us": percentile(durations_ns, 0.95) / 1000, + "rss_delta_bytes": after_rss - before_rss, + "pinned_block_bytes": block_bytes, + "new_host_allocations": actual_allocs, + "cpu_malloc_usable_min": min(usable) if usable else None, + "cpu_malloc_usable_median": statistics.median(usable) if usable else None, + "cpu_malloc_usable_max": max(usable) if usable else None, + } + ) + + metadata = { + "python": platform.python_version(), + "torch": torch.__version__, + "torch_commit": torch.version.git_version, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(0), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "cpu_affinity": sorted(os.sched_getaffinity(0)), + "libc": platform.libc_ver(), + "cpu_model": next( + line.split(":", 1)[1].strip() + for line in Path("/proc/cpuinfo").read_text().splitlines() + if line.startswith("model name") + ), + "final_rss_bytes": rss_bytes(), + "warmup_count": len(warmup), + "retained_measured_tensors": sum(len(group) for group in kept), + "size_order": sizes, + } + Path(args.worker_output).write_text(json.dumps({"metadata": metadata, "rows": rows}, indent=2) + "\n") + # Exit while warmup/kept still own their storages; never release measured + # tensors inside this process or let local-variable teardown recycle them. + os._exit(0) + + +def summarize(trials, sizes): + summary = [] + for size in sorted(sizes): + row = {"requested_bytes": size} + for mode in ("cpu", "pinned"): + selected = [ + r for trial in trials for r in trial["rows"] if r["mode"] == mode and r["requested_bytes"] == size + ] + times = [value for item in selected for value in item["latency_ns"]] + row[f"{mode}_count_per_repeat"] = selected[0]["count"] + row[f"{mode}_median_us"] = statistics.median(times) / 1000 + row[f"{mode}_mean_us"] = statistics.mean(times) / 1000 + row[f"{mode}_p95_us"] = percentile(times, 0.95) / 1000 + row[f"{mode}_repeat_median_min_us"] = min(r["median_us"] for r in selected) + row[f"{mode}_repeat_median_max_us"] = max(r["median_us"] for r in selected) + if mode == "cpu": + usable = [r for r in selected if r["cpu_malloc_usable_min"] is not None] + row["cpu_malloc_usable_min"] = min(r["cpu_malloc_usable_min"] for r in usable) if usable else None + row["cpu_malloc_usable_max"] = max(r["cpu_malloc_usable_max"] for r in usable) if usable else None + else: + row["pinned_block_bytes"] = selected[0]["pinned_block_bytes"] + row["pinned_over_cpu_median"] = row["pinned_median_us"] / row["cpu_median_us"] + summary.append(row) + return summary + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, default=Path("/tmp/cpu_pinned_alloc_keepalive")) + parser.add_argument("--sizes", nargs="+", type=int, default=DEFAULT_SIZES, help="Allocation sizes in bytes") + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument( + "--case-mib", type=int, default=128, help="Target retained pinned MiB per size; min-count may exceed it" + ) + parser.add_argument("--min-count", type=int, default=8) + parser.add_argument("--max-count", type=int, default=2048) + parser.add_argument("--cpu", type=int, help="Bind workers to this logical CPU") + parser.add_argument( + "--malloc-usable-size", action="store_true", help="Probe CPU malloc sizes; only with default glibc allocator" + ) + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--mode", choices=["cpu", "pinned"], help=argparse.SUPPRESS) + parser.add_argument("--repeat", type=int, default=0, help=argparse.SUPPRESS) + parser.add_argument("--worker-output", help=argparse.SUPPRESS) + args = parser.parse_args() + if min(args.sizes + [args.repeats, args.case_mib, args.min_count]) < 1 or args.max_count < args.min_count: + parser.error("Sizes/counts/repeats must be positive and max-count >= min-count") + if len(set(args.sizes)) != len(args.sizes): + parser.error("Sizes must be unique") + for name in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF"): + if os.environ.get(name): + parser.error(f"Unset {name} to measure default allocator behavior") + if args.worker: + run_worker(args) + + args.output_dir.mkdir(parents=True, exist_ok=False) + trials = [] + estimated_pinned = sum(sample_count(size, args) * (1 << (size - 1).bit_length()) for size in args.sizes) + print(f"Estimated retained pinned payload per worker: {estimated_pinned / MIB:.1f} MiB", flush=True) + for repeat in range(args.repeats): + # Alternate mode order; both modes use the same shuffled size order. + for mode in ("cpu", "pinned") if repeat % 2 == 0 else ("pinned", "cpu"): + path = args.output_dir / f"{mode}_{repeat}.json" + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--worker", + "--mode", + mode, + "--repeat", + str(repeat), + "--worker-output", + str(path), + "--case-mib", + str(args.case_mib), + "--min-count", + str(args.min_count), + "--max-count", + str(args.max_count), + "--sizes", + *map(str, args.sizes), + ] + if args.cpu is not None: + command.extend(["--cpu", str(args.cpu)]) + if args.malloc_usable_size: + command.append("--malloc-usable-size") + subprocess.run(command, check=True) + trial = json.loads(path.read_text()) + trials.append(trial) + print( + f'{mode} repeat {repeat + 1}: retained {trial["metadata"]["retained_measured_tensors"]} tensors; ' + f'RSS {trial["metadata"]["final_rss_bytes"] / MIB:.1f} MiB', + flush=True, + ) + + summary = summarize(trials, args.sizes) + with (args.output_dir / "summary.csv").open("w", newline="") as output: + writer = csv.DictWriter(output, fieldnames=list(summary[0])) + writer.writeheader() + writer.writerows(summary) + (args.output_dir / "results.json").write_text( + json.dumps( + { + "parameters": { + key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items() + }, + "method": "torch.empty(uint8); no fill/copy; " + "all warmup and measured tensors retained until process exit; " + "CUDA initialization excluded; GC disabled; per-allocation timings; fresh process per mode/repeat", + "estimated_retained_pinned_bytes": estimated_pinned, + "summary": summary, + "trials": trials, + }, + indent=2, + ) + + "\n" + ) + print("bytes,count/repeat,cpu median us,pinned median us,ratio,pinned block bytes,cpu malloc min/max") + for row in summary: + print( + f'{row["requested_bytes"]},{row["cpu_count_per_repeat"]},{row["cpu_median_us"]:.3f},' + f'{row["pinned_median_us"]:.3f},{row["pinned_over_cpu_median"]:.2f},' + f'{row["pinned_block_bytes"]},{row["cpu_malloc_usable_min"]}/{row["cpu_malloc_usable_max"]}' + ) + print(f"Results: {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/benchmark_req_token_index_alloc.py b/test/benchmark/benchmark_req_token_index_alloc.py index 55aa202017..f7ccc0cb0b 100644 --- a/test/benchmark/benchmark_req_token_index_alloc.py +++ b/test/benchmark/benchmark_req_token_index_alloc.py @@ -1,4 +1,4 @@ -"""Measure page-cache token index pinned allocation, excluding data writes and H2D. +"""Measure shared token index pinned allocation, excluding data writes and H2D. Example: CUDA_VISIBLE_DEVICES=4 PYTHONPATH=. python test/benchmark/benchmark_req_token_index_alloc.py \ @@ -19,7 +19,7 @@ import torch -from lightllm.server.router.model_infer.infer_req_ext1 import PinMemPageCacheReqTokenIndex +from lightllm.server.router.model_infer.infer_req_ext1 import ReqTokenIndex SLAB_TOKENS = 1024 * 1024 @@ -43,18 +43,15 @@ def summarize(samples): def measure_cached(mode, token_num, repeats, backing): - allocate = allocate_direct if mode == "direct_cached" else PinMemPageCacheReqTokenIndex._alloc_pinned_tensor + allocate = allocate_direct if mode == "direct_cached" else ReqTokenIndex._alloc_return_pinned_tensor samples = [] for _ in range(repeats): if mode == "slab_slice" and 0 < token_num <= SLAB_TOKENS: # 仅测切片路径：当前没有存活的返回值或异步传输，允许在计时外重置测试 backing。 # 生产接口始终向后切分；slab_cached 模式保留其真实换块行为。 - if ( - PinMemPageCacheReqTokenIndex._pinned_buffer is None - or PinMemPageCacheReqTokenIndex._pinned_offset + token_num > SLAB_TOKENS - ): - PinMemPageCacheReqTokenIndex._pinned_buffer = backing - PinMemPageCacheReqTokenIndex._pinned_offset = 0 + buffer, offset = getattr(allocate, "return_buffer_state", (None, 0)) + if buffer is None or offset + token_num > SLAB_TOKENS: + allocate.return_buffer_state = (backing, 0) start = CLOCK() tensor = allocate(token_num) samples.append((CLOCK() - start) / 1000) @@ -64,13 +61,12 @@ def measure_cached(mode, token_num, repeats, backing): def measure_growth(mode, token_num, live_count): # 仅清理 HOST 空闲缓存；用于构造缓存未命中，不用于生产服务。 - PinMemPageCacheReqTokenIndex._pinned_buffer = None - PinMemPageCacheReqTokenIndex._pinned_offset = 0 + ReqTokenIndex._alloc_return_pinned_tensor.return_buffer_state = (None, 0) gc.collect() torch.cuda.synchronize() torch._C._host_emptyCache() before = torch.cuda.memory.host_memory_stats() - allocate = allocate_direct if mode == "direct" else PinMemPageCacheReqTokenIndex._alloc_pinned_tensor + allocate = allocate_direct if mode == "direct" else ReqTokenIndex._alloc_return_pinned_tensor samples, held = [], [] for _ in range(live_count): start = CLOCK() @@ -109,7 +105,7 @@ def main(): del seed torch._C._host_emptyCache() start = CLOCK() - seed = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(1) + seed = ReqTokenIndex._alloc_return_pinned_tensor(1) initial_slab_us = (CLOCK() - start) / 1000 backing = seed._base assert backing.numel() == SLAB_TOKENS and backing.is_pinned() @@ -130,7 +126,7 @@ def main(): "torch_threads": torch.get_num_threads(), "allocator_env": {name: os.environ.get(name) for name in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF")}, "source_sha256": hashlib.sha256( - Path(PinMemPageCacheReqTokenIndex._alloc_pinned_tensor.__code__.co_filename).read_bytes() + Path(ReqTokenIndex._alloc_return_pinned_tensor.__code__.co_filename).read_bytes() ).hexdigest(), "slab_tokens": SLAB_TOKENS, "slab_bytes": SLAB_TOKENS * 4, @@ -182,8 +178,7 @@ def main(): result["growth"][str(token_num)] = data print(f"Uncached growth for {token_num} tokens x {live_count} complete", flush=True) - PinMemPageCacheReqTokenIndex._pinned_buffer = None - PinMemPageCacheReqTokenIndex._pinned_offset = 0 + ReqTokenIndex._alloc_return_pinned_tensor.return_buffer_state = (None, 0) gc.enable() (args.output_dir / "result.json").write_text(json.dumps(result, indent=2) + "\n") with (args.output_dir / "timings.csv").open("w", newline="") as output: diff --git a/unit_tests/common/test_req_token_index.py b/unit_tests/common/test_req_token_index.py index 9dbedfbd78..29117bb4cf 100644 --- a/unit_tests/common/test_req_token_index.py +++ b/unit_tests/common/test_req_token_index.py @@ -57,9 +57,7 @@ def test_factory_uses_request_mode_and_row_view_without_retaining_request(mode, assert not hasattr(InferReq, "move_to_gpu_index_len") assert isinstance(helper, expected_type) - assert hasattr(helper, "_alloc_pinned_tensor") == (expected_type is PinMemPageCacheReqTokenIndex) - assert not hasattr(ReqTokenIndex, "_pinned_buffer") - assert not hasattr(ReqTokenIndex, "_pinned_offset") + assert helper._alloc_return_pinned_tensor is ReqTokenIndex._alloc_return_pinned_tensor assert helper._index_stream is index_stream assert helper.req_idx == req.req_idx assert helper.gpu_indexes.data_ptr() == table[req.req_idx].data_ptr() @@ -90,18 +88,18 @@ def test_start_args_and_cli_select_mode(mode): @requires_cuda @pytest.mark.parametrize("mode", MODES) @pytest.mark.parametrize("pinned", [False, True]) -def test_append_accepts_cpu_input_and_allocates_pinned_storage_when_needed(monkeypatch, mode, pinned): +def test_append_accepts_cpu_input_and_uses_selected_storage_mode(monkeypatch, mode, pinned): table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") helper = create_req_token_index(make_req(table, mode), index_stream=torch.cuda.current_stream()) indexes = torch.arange(4, dtype=torch.int32, pin_memory=pinned) - alloc_pinned = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor + alloc_pinned = PinMemPageCacheReqTokenIndex._alloc_return_pinned_tensor pinned_sizes = [] def record_alloc_pinned(need_size): pinned_sizes.append(need_size) return alloc_pinned(need_size) - monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_alloc_pinned_tensor", staticmethod(record_alloc_pinned)) + monkeypatch.setattr(ReqTokenIndex, "_alloc_return_pinned_tensor", staticmethod(record_alloc_pinned)) helper.append_token_indexes(indexes) # 普通 CPU 输入已经复制到自有 storage，可在 append 返回后立即改写。 if not pinned: @@ -112,17 +110,47 @@ def record_alloc_pinned(need_size): assert helper.get_cpu_token_indexes(0, 4).tolist() == [0, 1, 2, 3] assert helper.get_gpu_token_indexes(0, 4).cpu().tolist() == [0, 1, 2, 3] if mode != "pure_gpu": - assert helper._cpu_buffer.is_pinned() + assert helper._cpu_buffer.is_pinned() == (mode == "pin_mem_page_cache") assert helper._cpu_buffer.untyped_storage().nbytes() == helper._cpu_buffer.nbytes +@requires_cuda +def test_upload_snapshots_own_only_pending_indexes_across_cache_changes(): + table = torch.full((1, 8192), -1, dtype=torch.int32, device="cuda") + helper = create_req_token_index(make_req(table), index_stream=torch.cuda.default_stream()) + helper.append_token_indexes(torch.arange(4, dtype=torch.int32)) + helper.get_gpu_token_indexes(0, 4) + helper.append_token_indexes(torch.arange(4, 8, dtype=torch.int32)) + buffer_ref = weakref.ref(helper._cpu_buffer) + + first = helper.get_upload_indexes() + second = helper.get_upload_indexes() + assert helper.gpu_index_submitted_len == 4 + assert first.is_pinned() and second.is_pinned() + assert first.dtype == second.dtype == table.dtype + assert first.tolist() == second.tolist() == [4, 5, 6, 7] + assert first.nbytes == second.nbytes == 4 * table.element_size() + assert first.data_ptr() != second.data_ptr() + assert not helper._cpu_buffer.is_pinned() + + # 修改常驻缓存、扩容和 reset 均不能改变已经返回的上传快照。 + helper._cpu_buffer[4:8].fill_(-99) + helper.append_token_indexes(torch.arange(8, 4096, dtype=torch.int32)) + assert buffer_ref() is None + assert not helper._cpu_buffer.is_pinned() + helper.reset() + helper.append_token_indexes(torch.arange(100, 104, dtype=torch.int32)) + assert first.tolist() == second.tolist() == [4, 5, 6, 7] + assert helper.get_upload_indexes().tolist() == [100, 101, 102, 103] + + @requires_cuda def test_shared_pinned_allocator_preserves_slices_across_slab_replacement(monkeypatch): - monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_pinned_buffer", None) - monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_pinned_offset", 0) + allocator = ReqTokenIndex._alloc_return_pinned_tensor + monkeypatch.setattr(allocator, "return_buffer_state", (None, 0), raising=False) buffer_token_num = 1024 * 1024 - first = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(5) - second = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(3) + first = PinMemCacheReqTokenIndex._alloc_return_pinned_tensor(5) + second = PinMemPageCacheReqTokenIndex._alloc_return_pinned_tensor(3) first.fill_(11) second.fill_(22) assert second.data_ptr() == first.data_ptr() + first.nbytes @@ -130,28 +158,28 @@ def test_shared_pinned_allocator_preserves_slices_across_slab_replacement(monkey assert first.dtype == second.dtype == torch.int32 assert not first.requires_grad and not second.requires_grad - buffer = PinMemPageCacheReqTokenIndex._pinned_buffer - empty = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(0) - oversized = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(buffer_token_num + 1) + buffer, offset = allocator.return_buffer_state + assert offset == 8 + empty = PinMemPageCacheReqTokenIndex._alloc_return_pinned_tensor(0) + oversized = PinMemPageCacheReqTokenIndex._alloc_return_pinned_tensor(buffer_token_num + 1) assert empty.numel() == 0 and oversized.numel() == buffer_token_num + 1 and oversized.is_pinned() - assert PinMemPageCacheReqTokenIndex._pinned_buffer is buffer - assert PinMemPageCacheReqTokenIndex._pinned_offset == 8 + assert allocator.return_buffer_state[0] is buffer + assert allocator.return_buffer_state[1] == 8 assert oversized.untyped_storage().data_ptr() != first.untyped_storage().data_ptr() - replacement = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(buffer_token_num - 7) + replacement = PinMemPageCacheReqTokenIndex._alloc_return_pinned_tensor(buffer_token_num - 7) replacement.fill_(33) assert replacement.untyped_storage().data_ptr() != first.untyped_storage().data_ptr() assert first.tolist() == [11] * 5 and second.tolist() == [22] * 3 - assert PinMemPageCacheReqTokenIndex._pinned_buffer.numel() == buffer_token_num - assert PinMemPageCacheReqTokenIndex._pinned_offset == buffer_token_num - 7 + assert allocator.return_buffer_state[0].numel() == buffer_token_num + assert allocator.return_buffer_state[1] == buffer_token_num - 7 @requires_cuda def test_shared_pinned_allocator_rollover_keeps_released_h2d_source_valid(monkeypatch): - monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_pinned_buffer", None) - monkeypatch.setattr(PinMemPageCacheReqTokenIndex, "_pinned_offset", 0) + monkeypatch.setattr(ReqTokenIndex._alloc_return_pinned_tensor, "return_buffer_state", (None, 0), raising=False) buffer_token_num = 1024 * 1024 - source = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(buffer_token_num) + source = PinMemPageCacheReqTokenIndex._alloc_return_pinned_tensor(buffer_token_num) source.fill_(11) source_ref = weakref.ref(source) target = torch.empty(buffer_token_num, dtype=torch.int32, device="cuda") @@ -164,7 +192,7 @@ def test_shared_pinned_allocator_rollover_keeps_released_h2d_source_valid(monkey # 连续换掉多个 slab，旧 storage 释放后仍须等在途 H2D 结束才能被重新使用。 for value in range(22, 26): - replacement = PinMemPageCacheReqTokenIndex._alloc_pinned_tensor(buffer_token_num) + replacement = PinMemPageCacheReqTokenIndex._alloc_return_pinned_tensor(buffer_token_num) replacement.fill_(value) del replacement stream.synchronize() @@ -222,8 +250,11 @@ def test_page_cache_rejects_partial_page_appends_before_mutating_state(start, ap @requires_cuda def test_batch_upload_rejects_missing_gpu_submission_without_advancing_progress(monkeypatch): + from lightllm.common import cuda_wrapper from lightllm.server.router.model_infer.mode_backend import base_backend + # 故障注入会禁用原生 batch 路径；测试结束后恢复，避免影响其他上传用例。 + monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", True) table = torch.full((1, 8), -1, dtype=torch.int32, device="cuda") req = make_req(table, "pure_gpu") req.token_index = create_req_token_index(req, index_stream=torch.cuda.current_stream()) @@ -478,7 +509,7 @@ def test_reset_rebuild_keeps_cpu_appends_and_gpu_submission_separate(mode): (1025, 5500, 4096, 5000, 5500), ], ) -def test_pin_buffer_capacity_and_lifecycle( +def test_cpu_buffer_capacity_and_lifecycle( mode, page_size, input_len, max_length, initial_capacity, next_end, grown_capacity ): req = make_req( @@ -500,11 +531,13 @@ def test_pin_buffer_capacity_and_lifecycle( first_index = torch.arange(storage_page_size, dtype=torch.int32, pin_memory=True) helper.append_token_indexes(first_index) assert helper._cpu_buffer.numel() == initial_elements - assert helper._cpu_buffer.is_pinned() and helper._cpu_buffer.data_ptr() != first_index.data_ptr() + assert helper._cpu_buffer.is_pinned() == (mode == "pin_mem_page_cache") + assert helper._cpu_buffer.data_ptr() != first_index.data_ptr() helper.append_token_indexes(torch.arange(storage_page_size, initial_capacity, dtype=torch.int32, pin_memory=True)) old_buffer = helper._cpu_buffer helper.append_token_indexes(torch.arange(initial_capacity, next_end, dtype=torch.int32, pin_memory=True)) - assert helper._cpu_buffer.numel() == grown_elements and helper._cpu_buffer.is_pinned() + assert helper._cpu_buffer.numel() == grown_elements + assert helper._cpu_buffer.is_pinned() == (mode == "pin_mem_page_cache") if mode == "pin_mem_page_cache": assert helper._cpu_buffer is old_buffer else: @@ -528,7 +561,7 @@ def test_pin_buffer_capacity_and_lifecycle( @pytest.mark.parametrize("mode", ["pin_mem_cache", "pin_mem_page_cache"]) @pytest.mark.parametrize("page_size", [1, 64]) @pytest.mark.parametrize("use_default_stream", [False, True]) -def test_pin_buffer_reset_during_upload_keeps_old_data_until_h2d_completes(mode, page_size, use_default_stream): +def test_cpu_buffer_reset_during_upload_keeps_old_data_until_h2d_completes(mode, page_size, use_default_stream): count = 8192 row = torch.full((count,), -1, dtype=torch.int32, device="cuda") index_stream = torch.cuda.default_stream() if use_default_stream else torch.cuda.Stream() @@ -644,7 +677,7 @@ def test_manager_owns_only_gpu_table_and_initializes_hold_page(monkeypatch, page @pytest.mark.parametrize("mode", MODES) @pytest.mark.parametrize("page_size", [1, 64]) @pytest.mark.parametrize("index_stream_kind", ["default", "shared"]) -def test_backend_uploads_pinned_views_without_request_events_or_keepalive_list( +def test_backend_uploads_pinned_snapshots_without_request_events_or_keepalive_list( monkeypatch, mode, page_size, batch_available, index_stream_kind ): from lightllm.common import cuda_wrapper @@ -682,13 +715,15 @@ def test_backend_uploads_pinned_views_without_request_events_or_keepalive_list( event_type = torch.cuda.Event def record_batch(dst_tensors, src_tensors): - # 完整缓存直接借用视图，页缓存先展开；两者都应一次批量提交正确的 token 索引。 + # 完整缓存复制待上传区间，页缓存先展开；两者都应批量提交独立的 token 索引快照。 assert all(tensor.is_pinned() for tensor in src_tensors) assert all(torch.equal(tensor, expected[row]) for row, tensor in enumerate(src_tensors)) if mode == "pin_mem_cache": - assert [tensor.data_ptr() for tensor in src_tensors] == [ - req.token_index.get_upload_indexes().data_ptr() for req in reqs - ] + for req, source in zip(reqs, src_tensors): + assert source.untyped_storage().data_ptr() != req.token_index._cpu_buffer.data_ptr() + # 即使上传尚未开始，改写原缓存也不能改变已经准备好的 pinned 源。 + req.token_index._cpu_buffer[:count].fill_(-99) + assert all(torch.equal(tensor, expected[row]) for row, tensor in enumerate(src_tensors)) batch_calls.append(len(src_tensors)) return copy_batch(dst_tensors, src_tensors)