diff --git a/docs/cache_manager_design.md b/docs/cache_manager_design.md new file mode 100644 index 0000000000..d0610fda4d --- /dev/null +++ b/docs/cache_manager_design.md @@ -0,0 +1,1344 @@ +# Attention Resources 架构设计草案 + +本文定义 attention 资源的架构设计，包括 `AttResourceManager` 的职责、共享 allocator、各类 State 的接口、目录组织及 CPU/DP/PD 传输约定。文中的结构与接口为设计契约，待细化事项集中列于第 9 节。 + +## 1. 设计约定 + +统一在 `lightllm/common/att_resources/` 下组织可复用的资源实现，由 `AttResourceManager` 约定公共资源管理接口，各模型提供具体管理器实现，完成资源初始化、内部状态页的聚合操作及 CPU/DP/PD 数据传输。radix cache、CPU cache 和传输调用方依赖对应模型的管理器，资源层不反向依赖这些调用方。 + +公共层约定三种 allocator 入口，具体实例按模型组合和启用的功能初始化： + +```python +class AttResourceManager: + full_att_allocator: KVPageAllocator | None + bounded_att_big_page_allocator: CheckpointSlotAllocator | None + bounded_att_small_page_allocator: CheckpointSlotAllocator | None +``` + +这三个入口表示资源管理器对分配域的管理责任，聚合操作与传输接口见第 6.2 节。具体初始化必须保证所需 allocator 已创建或绑定，无需的 allocator 为 `None`。例如纯 full-attention 模型只初始化 `full_att_allocator`；启用 linear checkpoint 的组合再按配置初始化大小页 allocator。GPU linear 运行态的存在本身不要求启用 checkpoint 池。 + +- 基类不声明固定的 `main_full_att`、`draft_full_att`、`main_linear_att`、`draft_linear_att`、`main_indexer_tail`、`draft_indexer_tail` 成员。具体模型自行组织内部资源，数量、组合、字段名和实现类型由模型决定。 +- `full_att`、`linear_att`、`indexer_tail` 保留为可复用实现的目录分类。full attention 保存历史 KV；linear attention 由 `LinearAttentionState` 抽象基类约定操作，按算法直接实现 `KDAState`、`GDNState` 等状态对象，包含各自的运行态与可选 checkpoint；indexer tail 保存增量索引所需尾部状态。 +- main/draft 可以采用不同方案；同一共享分配域的参与方使用同一个 allocator 实例，遵守统一的索引和回收约定。资源对象持有 buffer 和数据有效性信息；管理器创建或绑定共享 allocator，调用方直接使用对应 allocator 分配与归还索引。 +- 模型专属 `AttResourceManager` 提供运行态清零、按前缀长度重置 tail、token 复制、状态页的聚合保存、校验、恢复、失效和数据传输操作；缓存功能流程决定匹配、引用、发布与淘汰时机，推理流程决定请求执行阶段。 +- `req_to_token_indexs` 和请求槽位仍由 `ReqManager` 持有。资源管理器接收显式传入的物理索引、请求槽位及恢复长度，提供资源操作接口。 +- 内部 attention 资源的暴露方式作为接口设计问题单独讨论。当前建议按调用方区分统一资源操作、模型推理视图和 State 直接传输操作，详见第 6 节。 + +本文区分已确定的职责边界和接口草案。main/draft 采用组合管理器还是分别装配管理器实例、具体资源的绑定方式及功能接口参数的数据形式仍待细化。 + +## 2. 整体结构与职责边界 + +```mermaid +flowchart TD + Model[模型初始化] --> Concrete[模型专属 AttResourceManager] + Concrete -.->|实现公共接口| API[AttResourceManager] + Context[InferenceContext] --> Req[ReqManager] + Req --> Slots[请求槽位与 KV 映射] + Context -->|前缀复用| Radix[Radix Cache 功能] + Context -->|多级缓存| CPU[CPU Cache 功能] + Radix -->|调用通用状态页操作| Concrete + CPU -->|缓存数据读写| Concrete + Concrete --> Alloc[按需初始化或绑定的共享 allocator] + Concrete --> Internal[模型自行组织的内部资源] + Internal --> Full[full attention 实现] + Internal --> Linear[KDAState / GDNState 等具体状态] + Internal --> Indexer[indexer tail 实现] + Concrete -->|绑定具体视图| Forward[模型与 attention backend] + Transport[DP/PD 调度] -->|调用具体资源操作| Concrete +``` + +图中的资源类别按模型需要选择，可以缺省或包含多个实例。箭头表示调用或组合关系；main/draft 管理器实例的具体装配方式尚未固定。 + +| 层次 | 组件 | 职责 | +| --- | --- | --- | +| 请求流程 | `InferenceContext`、`ReqManager` | 创建、暂停、恢复和完成请求；由 ReqManager 管理请求槽位、KV 索引映射及采样参数 | +| 资源数据操作 | `AttResourceManager` 公共功能接口及模型专属实现 | 管理共享 allocator 与内部 State，实现内部状态页的聚合操作与 CPU/DP/PD 数据传输 | +| 缓存功能 | radix、CPU cache 及其调用流程 | 组织前缀匹配、节点引用、发布与淘汰，管理外部缓存页，并调用具体管理器操作资源数据 | +| 资源实现 | 模型所需的 State、存储布局、checkpoint pool | 持有 buffer、布局与数据有效性，实现读写、复制、重置、快照及自身存储的跨进程导出与绑定 | +| 传输调度 | DP/PD 与多级缓存模块 | 管理任务、队列、通信、事件同步及完成通知 | + +`AttResourceManager` 不持有请求 KV 映射表，也不直接访问全局 `g_infer_context` 或通过 `ReqManager` 查询映射。资源方法接收调用方提供的请求槽位、长度、物理索引和资源引用；将逻辑 token 位置转换为物理索引、写入或清理请求映射属于请求流程。缓存功能调用方提供已选定数据的长度、索引和资源引用，具体管理器执行清零、复制、保存、校验、恢复或失效。请求阶段的组织和缓存节点策略由相应调用流程负责。 + +代码依赖从调用流程向资源层单向传递。radix cache 依赖管理器的通用状态页操作；管理器与 State 不导入或访问 radix 节点、树和策略。资源实现也不反向导入请求调度模块；传输调度使用数据读写与复制接口。 + +## 3. 对象结构与资源归属 + +```text +ReqManager +├── request_slots +└── req_to_token_indexs # 共享请求到 KV 的索引映射 + +AttResourceManager # 公共接口与 allocator 入口 +├── full_att_allocator: KVPageAllocator | None +├── bounded_att_big_page_allocator: CheckpointSlotAllocator | None +└── bounded_att_small_page_allocator: CheckpointSlotAllocator | None + +AttResourceManager # 模型实现，内部结构由其自行定义 +├── 按需创建或绑定所需 allocator +├── 初始化模型所需的资源组合 +├── 实现内部状态页聚合操作与 CPU / DP / PD 数据传输 +└── 绑定推理视图、组织传输所需状态 # 搬运直接调用 State 方法 +``` + +例如纯 full-attention 管理器可以只保存一个 KV 存储对象；混合模型管理器可以组合 KV、多个 linear 运行态及 indexer tail；包含 draft 时按实际方案组织相关资源。可共用初始化逻辑的模型可以复用管理器实现。公共调用方通过接口使用这些组合。 + +以下列出可复用资源组件的结构示例。linear attention 直接使用 `KDAState`、`GDNState` 等具体状态对象，由 `LinearAttentionState` 约定必需操作，不增加资源包装实例。管理器内部的字段名和实例数量由模型决定，checkpoint pool 仅在相应功能启用时存在： + +```text +FullAttentionState +├── storage_spec # 容量、页大小、保留索引范围、device +├── storage # 完整数据区域、固定参数及存储描述 +└── layer_mapping # 本模型的层映射 + +KDAState # 一种具体 attention 状态实现 +├── conv_state: LayerCache # 卷积历史窗口及 MTP 尾部空间 +├── ssm_state: LayerCache # KDA 递推状态及 MTP 状态槽位 +├── mtp_state_index: Tensor | None +├── big_pages_pool # 按需创建，使用共享 allocator 分配的槽位 +│ ├── buffers +│ └── state_token_lens +└── small_pages_pool # 按需创建 + ├── buffers + └── state_token_lens + +GDNState # 独立定义 GDN 状态、布局和快照操作 +└── 按其实现组织 Conv / SSM / MTP 状态及可选大小页 pool + +IndexerTailState # 抽象操作契约，不固定内部存储 +└── KPoolIndexerTailState # GLM K-pool 的具体状态 + └── tail_buffer # 原始 key/gate ring，按层、请求组织 +``` + +同一共享分配域中的 allocator 只初始化一次，由资源管理器负责其完整生命周期。若 main/draft 分别使用管理器实例，装配时应注入共享 allocator 引用，并明确唯一的创建和销毁责任，以及调用方统一失效并向 allocator 单次归还的约定；若采用一个组合管理器，则由该实现协调内部资源。两种装配方式均不能因模型数增加而重复分配或释放同一共享索引。实例装配形式保留为待定项。 + +| 资源 | 持有者 | 分配单位 / 生命周期 | +| --- | --- | --- | +| full-attention 空闲页与容量状态 | `AttResourceManager.full_att_allocator` | 每次为请求统一分配/回收物理页索引，main/draft 共用 | +| 大页 checkpoint 空闲槽位 | `AttResourceManager.bounded_att_big_page_allocator` | 为 main/draft 参与的大页 pool 统一分配槽位 ID | +| 小页 checkpoint 空闲槽位 | `AttResourceManager.bounded_att_small_page_allocator` | 为 main/draft 参与的小页 pool 统一分配槽位 ID，独立于大页分配域 | +| 请求 KV 索引表 | `ReqManager` | 记录逻辑 token 到共享物理索引的映射；负责绑定、更新和清理 | +| HOLD 保留物理页及索引描述 | `AttResourceManager` 及其 full-attention allocator | 从普通分配范围排除保留页；ReqManager 使用其索引初始化 HOLD 行 | +| KV、量化 scale、packed indexer buffer 与层映射 | 具体管理器组织的 full-attention 资源 | 按统一容量建立存储，使用共享索引访问各自数据 | +| KDA/GDN 的 Conv/SSM、MTP 游标 | 具体管理器持有的 `KDAState` / `GDNState` | 使用共同请求槽位域，由各算法状态对象分别维护 | +| 大小页 checkpoint buffer 与 `state_token_lens` | 各具体状态对象中启用的 pool | 使用共同分配的槽位 ID，分别保存数据和长度；不持有独立空闲列表 | +| indexer 原始尾部窗口 | 具体管理器组织的 indexer 资源 | 使用共同请求槽位域，各自维护所需 tail 布局 | +| 请求槽位 | `ReqManager` | 请求创建时分配，完成时回收，暂停时保留 | +| 前缀缓存索引与引用记录 | radix cache 功能 | 维护前缀节点及引用；通过管理器保存、恢复或失效关联资源数据 | + +### 3.1 main/draft 共享 full-attention allocator + +调用方直接通过 `AttResourceManager.full_att_allocator` 分配一次 KV 索引，再通过 ReqManager 将索引写入共享请求表。复用前缀时，radix 功能选定缓存数据，具体管理器校验并恢复关联状态，所需 KV 索引由 ReqManager 绑定到请求映射。main 和 draft 从同一请求映射获取索引，访问各自的存储： + +```text +逻辑 token 位置 t + ↓ ReqManager.req_to_token_indexs +统一物理索引 i + ├── main 对应的 KV 存储位置 + └── draft 对应的 KV 存储位置 +``` + +二者的层数、head 维度、dtype 和量化布局可以不同；同一索引域内需要使用兼容的 page_size，并且各存储都覆盖 allocator 可返回的索引范围。同一分配域统一容量和物理页索引规则，不引入模型间索引重映射。 + +共享索引表示存储位置的对应关系，不表示 main/draft 在每个位置都已经写入有效 KV。MTP 预留区和回退后的有效长度仍由请求执行流程维护。 + +释放时，请求流程与具体管理器确认相关使用方已完成，由 `ReqManager` 清理请求映射。调用方解除相关使用引用后，仅将不再被缓存、请求或传输持有的索引向同一个 `full_att_allocator` 归还一次。任一模型的资源对象都不能自行释放共享索引。容量评估需要计入各模型真实占用的存储；若资源实例只是同一 buffer 的视图，不重复计算该 buffer。 + +### 3.2 main/draft 共享 linear checkpoint allocator + +main/draft 选择各自需要的具体状态对象，例如 main 使用 `KDAState`，draft 使用 `GDNState` 或无需 linear 状态。各状态对象持有自身运行态及启用的大页或小页 pool。对同一前缀的一个大页或小页 checkpoint，调用方直接从管理器持有的对应 allocator 分配一次槽位 ID，再将该 ID 传给管理器的聚合数据操作，供参与的两侧 pool 使用。两侧的状态布局、dtype 和单槽位字节数可以不同，各自的存储必须覆盖共享 allocator 可返回的槽位范围。 + +例如分配大页槽位 `s` 后，main 和 draft 对应的大页 pool 在各自的槽位 `s` 保存状态。具体管理器负责找到参与操作的 pool，请求流程无需读取模型内部字段。两侧各自维护 `state_token_lens[s]`；本次复用所需的两侧状态都对应目标长度并满足同步条件后，该缓存数据才具备发布或恢复条件；发布由缓存调用流程组织。仅一侧具有 linear 资源时，只处理该侧。 + +共享 checkpoint 需要约定共同的恢复边界和 token 位置语义，并确保参与实现都能捕获该位置的状态。MTP 推测中的临时状态需先按执行流程选出对应位置的状态，再保存快照。共享索引不要求 main/draft 每一步同步执行；回收时须等所有参与方结束使用，槽位只归还一次。 + +linear 运行态与两侧 indexer tail 均由 `ReqManager` 提供的 `req_idx` 定位。indexer tail 不单独分配 checkpoint 槽位；其状态布局和恢复规则由各自实现描述。 + +### 3.3 分配状态与数据状态分离 + +allocator 是索引分配的唯一事实来源，资源对象是存储布局与数据有效性的持有者。缓存功能流程通过 radix 维护页/槽位的引用与发布关系；具体管理器校验资源数据的完整性并聚合处理页失效；调用方确认最后一个使用引用已解除后，直接向对应 allocator 归还索引。 + +同一 linear checkpoint 可以包含 main/draft 多个状态资源的快照。槽位分配状态由共享 allocator 维护，数据内容和有效长度由各侧 pool 分别维护。 + +对于使用 K-pool 的实现，已完成池化的压缩 key/scale 保留在对应 full-attention 资源的 packed 存储中，indexer tail 资源管理增量计算需要的原始短窗口；二者的组合关系由具体管理器维护。 + +### 3.4 按具体 attention 实现定义状态 + +定义只约束操作的 `LinearAttentionState` 抽象基类，由 `KDAState`、`GDNState` 等子类负责自己的成员、布局、重置、快照与恢复语义。`AttResourceManager` 初始化并组合所需对象，公共缓存流程通过管理器访问它们。不同模型也可以复用同一种状态实现。 + +`KDAState` 的成员定义如下： + +```python +class KDAState(LinearAttentionState): + conv_state: LayerCache + ssm_state: LayerCache + mtp_state_index: torch.Tensor | None + mtp_step: int # 该实例的 MTP 布局参数 + big_pages_pool: LinearCheckpointPool | None + small_pages_pool: LinearCheckpointPool | None +``` + +`LayerCache` 封装按层组织的 Tensor。该示例中的 Conv、SSM 和 MTP 索引驻留 GPU；checkpoint pool 独立持有快照 buffer 和 `state_token_lens`。设请求槽位容量为 `N`（包含保留槽位），该实例的 `mtp_step` 为 `S`： + +| 成员 | 内容与索引方式 | +| --- | --- | +| `conv_state` | 对应算法的卷积历史窗口，窗口宽度为 `conv_kernel_size - 1 + S`；用 `buffer[layer, req_idx, ...]` 访问 | +| `ssm_state` | KDA 的递推矩阵状态，每请求预留 `S + 1` 个状态槽位；第 `j` 个槽位位于 `buffer[layer, req_idx * (S + 1) + j, ...]` | +| `mtp_state_index` | 形状为 `[N]` 的 `int32` Tensor；记录各请求接受的 MTP 状态索引。`S == 0` 时为 `None` | +| 大小页 pool | 按需创建，保存该算法指定前缀位置的状态，并分别记录 `state_token_lens`；槽位来自管理器的对应 `bounded_att_*_allocator` | + +`GDNState` 按 GDN 算法独立定义 Conv/SSM 和 MTP 状态，可以复用底层 `LayerCache`、兼容的 checkpoint buffer 布局及复制工具。两者共同实现 `LinearAttentionState` 接口；基类不规定 Conv/SSM 等成员，不据此判断两种算法的状态内容可以互换。 + +运行态按 `ReqManager` 的槽位容量预分配，状态对象接收 `req_idx` 执行操作。main/draft 可以持有同一种状态类的两个独立实例，也可以使用不同状态类型；各实例的层数、维度、dtype 和 MTP 配置由对应模型决定。allocator 仍由 `AttResourceManager` 创建或绑定并统一管理。 + +后续 sliding-window attention 可以直接定义 `SlidingWindowState`，按窗口 KV、有效范围和游标等实际需求组织成员。它可以使用适用的 bounded-attention checkpoint 分配域，但其快照布局与恢复语义由自身实现定义。KDA/GDN 等递推状态的必需接口见 6.7 节；sliding-window 可以复用适用的页分配与描述工具，其接口另外按窗口语义细化，不强制继承线性状态基类。 + +## 4. 目录设计 + +目录按资源管理、allocator 与具体 State 组织，传输能力由具体实现类提供。管理器、公共 State、allocator 和 pool 文件按对应类名命名；具体 State 实现使用算法或布局名加 `_state.py`。Triton/CUDA 计算 kernel 位于对应计算模块，由 State 方法调用。 + +```text +lightllm/common/ +├── att_resources/ +│ ├── __init__.py +│ ├── att_resource_manager.py # AttResourceManager：缓存与传输的数据操作接口 +│ ├── types.py # StorageDescriptor / BufferTransferLayout / StateCopyMapping 等 +│ ├── allocators/ # 管理器直接持有，同一分配域共享实例 +│ │ ├── __init__.py +│ │ ├── kv_page_allocator.py # KVPageAllocator：full_att 分配域 +│ │ └── checkpoint_slot_allocator.py # CheckpointSlotAllocator：bounded_att 大小页各一个实例 +│ ├── full_att/ +│ │ ├── __init__.py +│ │ ├── full_attention_state.py # FullAttentionState：公共接口与可复用基础逻辑 +│ │ ├── dense_state.py # DenseFullAttentionState +│ │ ├── group_quantized_state.py # GroupQuantizedFullAttentionState：INT4 / INT8 等 +│ │ ├── static_fp8_state.py # StaticFP8FullAttentionState +│ │ ├── mla_state.py # MLAFullAttentionState +│ │ ├── packed_mla_state.py # PackedMLAFullAttentionState：含 indexer key/scale 的 packed 布局 +│ │ ├── operator/ # State 内部复用的 KV 写入、量化与复制函数 +│ │ └── transfer_ops.py # State 内部复用的 CPU/DP/PD 搬运函数 +│ ├── linear_att/ +│ │ ├── __init__.py +│ │ ├── linear_attention_state.py # LinearAttentionState：抽象操作契约，无固定存储成员 +│ │ ├── kda_state.py # KDAState：Conv / SSM / MTP 与页操作 +│ │ ├── gdn_state.py # GDNState：独立算法状态与页操作 +│ │ ├── linear_checkpoint_pool.py # LinearCheckpointPool：big_pages_pool / small_pages_pool 的存储 +│ │ └── transfer_ops.py # State 内部复用的状态打包与 TP 布局转换函数 +│ └── indexer_tail/ +│ ├── __init__.py +│ ├── indexer_tail_state.py # IndexerTailState：尾部状态抽象操作契约 +│ ├── kpool_state.py # KPoolIndexerTailState：原生 ring 与恢复规则 +│ └── transfer_ops.py # State 内部复用的 tail 搬运函数 +└── req_manager/ # 请求槽位、KV 索引映射、采样与生命周期 + +lightllm/models// +├── model.py # 选择、装配本模型的资源管理器 +└── att_resource_manager.py # 模型资源组合、功能数据操作及具体资源绑定 +``` + +`allocators/` 与三类 State 目录平级。`kv_page_allocator.py` 提供 full-attention 的物理页分配器；`checkpoint_slot_allocator.py` 提供 bounded-attention 的状态槽位分配器，管理器按需创建大小页两个独立实例。后者的分配能力不限定为 linear attention，使用它也不要求继承 `LinearAttentionState`。 + +`full_attention_state.py`、`linear_attention_state.py`、`indexer_tail_state.py` 分别定义三类公共接口；布局或算法差异由对应的 `*_state.py` 实现。full-attention 的 storage 属于具体 State 的内部实现，可与该 State 放在同一文件中；确有复用需要时再拆为内部存储工具，不设置另一套独立的 storage 公共继承体系。 + +`linear_checkpoint_pool.py` 对应 `LinearCheckpointPool` 存储组件，供 `big_pages_pool` 和 `small_pages_pool` 复用；两者通过实例区分。公共页操作仍使用 `save_pages`、`restore_pages` 等名称，pool 不持有 allocator。`KPoolIndexerTailState` 管理请求运行态，不创建独立大小页 pool。 + +`FullAttentionState`、`LinearAttentionState` 与 `IndexerTailState` 的具体子类直接实现 CPU、DP、PD 读写及自身存储的跨进程导出与绑定。各资源目录中的 `transfer_ops.py` 仅放这些方法内部复用的函数，按实际复用需要创建。模型专属 `AttResourceManager` 在自身实现中组织多资源读写、CPU 联合页、长度元数据和位置映射。网络通信、任务队列与完成通知由 server 层管理。 + +模型专属资源组合、缓存与传输数据操作放在 `lightllm/models//att_resource_manager.py`，同类模型可以复用管理器实现。radix 和 CPU cache 功能通过具体管理器操作资源数据，缓存策略与请求执行流程按各自职责组织。公共资源包不反向导入模型实现；模型初始化负责选择和装配。main/draft 的实例数量、内部字段及资源方案由具体管理器决定，共享 allocator 按分配域统一创建或绑定。 + +`types.py` 仅放跨资源使用的纯数据类型，包括 `StorageDescriptor`、`TransferRegion`、`BufferTransferLayout`、`StateCopyMapping`，以及页类别、槽位引用和有效长度等共享描述。各 State 的内部类型放在对应实现文件中，字段格式和状态选择由该实现解释。资源目录不引入统一插件注册机制，也不要求三类 State 继承同一个大接口。 + +## 5. allocator 与 checkpoint pool + +### 5.1 allocator 的公共入口与按需初始化 + +| 公共入口 | 类型 | 初始化条件 | +| --- | --- | --- | +| `full_att_allocator` | `KVPageAllocator \| None` | 组合中存在需要页式 KV 分配的 full-attention 资源 | +| `bounded_att_big_page_allocator` | `CheckpointSlotAllocator \| None` | 启用大页 linear checkpoint | +| `bounded_att_small_page_allocator` | `CheckpointSlotAllocator \| None` | 启用小页 linear checkpoint | + +具体管理器初始化完成时，已启用的分配域必须具有 allocator 实例及容量匹配的资源，未启用的入口保持 `None`。申请未启用的分配域应明确报错；公共缓存流程根据已启用能力选择操作。只有 full attention 的配置无需创建 linear allocator、运行态或 checkpoint pool。 + +allocator 入口的存在只能说明该分配域已启用，具体恢复边界、支持的传输布局等能力仍由模型管理器明确描述。 + +大小页均启用时，使用同一种 allocator 类的两个独立实例，每个实例供参与模型对应类别的 pool 共用。同一模型的单个快照保存该实现所需的完整状态，例如 Conv/SSM；main/draft 的快照布局和大小可以不同。“大页/小页”指 token 边界和缓存策略，不表示快照按 token 数分成不同大小。 + +`KVPageAllocator` 以完整物理页为分配单位，对外提供对应 token 索引；checkpoint allocator 返回状态槽位 ID。GPU 运行态和 indexer tail 使用请求槽位，不使用上述 checkpoint allocator。 + +### 5.2 管理器、allocator 与 pool 分工 + +`CheckpointSlotAllocator` 只维护可用槽位。`LinearCheckpointPool` 只持有状态 buffer 和有效长度，使用外部传入的槽位完成保存、恢复和失效处理。`AttResourceManager` 持有 allocator 并提供跨 State 的聚合数据操作，调用方直接使用 allocator。操作顺序如下： + +1. 调用方直接从 `bounded_att_big_page_allocator` 或 `bounded_att_small_page_allocator` 分配槽位。pool 初始化时有效长度为 None，槽位归还前统一失效，保证空闲槽位不保留有效页记录。 +2. `save_pages` 聚合调用参与 State 保存；异步操作通过 stream/event 保证后续读取顺序。保存失败时使本次目标页在全部参与 pool 中失效。 +3. `has_page` 校验全部参与 State 的页记录；`restore_pages` 在整体校验通过并满足同步条件后执行恢复。 +4. 调用方确认引用已解除且异步使用完成，先调用 `invalidate_pages` 清除全部参与 pool 的有效记录，再直接向同一个 allocator 归还共享槽位一次。 + +CPU load/offload 尾页的两个暂存槽位 ID 由管理器统一保留，排除在大页 allocator 的普通分配范围外；参与操作的 main/draft 大页 pool 分别为这些槽位提供存储。 + +具体管理器声明数据可捕获及可恢复的条件，并实现保存、恢复和失效操作；缓存调用流程选择保存位置、发布时机及淘汰对象，并直接调用 allocator 处理分配与归还。State 和 pool 执行对应的数据操作。 + +### 5.3 页内状态的 token 长度元数据 + +各模型实际创建的 `big_pages_pool` 和 `small_pages_pool` 分别保留一份 `state_token_lens: list[int | None]`，按槽位 ID 索引，记录该槽位保存的状态所对应的 token 前缀长度。长度从序列起点计数；pool 的物理槽位大小由具体状态布局决定。这份元数据用于校验保存和恢复位置。 + +| 槽位记录 | 含义 | +| --- | --- | +| `state_token_lens[slot] is None` | 该槽位没有有效快照记录，包括尚未保存、保存失败或已经释放的情况 | +| `state_token_lens[slot] == L` | 该槽位的快照对应处理完前 `L` 个 token 后的状态 | + +上层的 `length → slot` 映射表达某个前缀关联哪个槽位；pool 中的 `slot → state_token_len` 记录槽位内快照对应的位置。恢复时两者必须一致。只预留槽位但尚未捕获状态的中间页边界保持 `None`。 + +`LinearCheckpointPool` 通过保存、载入和失效接口维护这份记录；`AttResourceManager` 组织操作并校验恢复条件，allocator 只维护槽位的分配状态。CPU offload/load 同步携带并恢复长度元数据。对于包含 linear attention 的模型，每个已发布 CPU 缓存页必须包含全部参与 linear State 在该页前缀终点的有效状态；内部 pool 中尚未捕获的 None 记录不能发布为 CPU 缓存页。长度记录本身不表示异步复制已完成，也不标识具体 token 前缀；同步由 stream/event 保证，前缀身份由请求与缓存关联关系保证。 + +## 6. 对外接口与 attention 资源暴露方式 + +模型专属管理器负责资源初始化、内部组合及缓存和传输功能所需的数据操作。以下接口按数据用途划分；调用方显式提供索引、长度、槽位及来源，具体管理器组织内部 State 完成操作。 + +### 6.1 FullAttentionState：公共接口与不同 storage 的完整处理 + +本节是底层接口草案。`FullAttentionState` 负责一组历史 KV 存储的容量描述、布局访问、写入及复制，包含该布局所需的 scale 和 packed indexer 数据。`AttResourceManager` 决定实例组合、共享索引域及生命周期。allocator 由管理器持有，缓存发布由缓存调用流程组织，请求映射由 ReqManager 负责，通信任务由传输调度负责。 + +#### 继承与内部 storage 的关系 + +建议以 `FullAttentionState` 作为公共接口，由具体布局子类实现差异。子类内部的 storage 持有 buffer 和布局信息，operator 可以复用兼容的写入、量化及复制 kernel。布局选择集中在具体子类及其配套实现，公共调用方使用稳定接口。 + +```text +FullAttentionState +├── DenseFullAttentionState +├── GroupQuantizedFullAttentionState # INT4 / INT8 等 +├── StaticFP8FullAttentionState +└── MLAFullAttentionState + └── PackedMLAFullAttentionState # 示例：附带 indexer 数据 +``` + +这些是目标实现示例，具体继承关系以可复用的行为为准。改变布局的子类必须重新核对存储描述、层视图、写入、完整复制和传输能力；父类的布局假设不能直接用于新的 packed 格式。内部 storage 可以是明确类型的数据对象或可复用的存储组件，不要求建立与 State 平行的公共继承体系。 + +#### 必须提供的公共接口 + +```python +# ViewT / WriteT 为具体布局定义的类型，下面是公共签名草案。 +class FullAttentionState: + @classmethod + def get_required_bytes(cls, layout_config, storage_spec) -> int: + ... + + def __init__(self, layout_config, storage_spec): + ... + + def get_storage_descriptor(self) -> StorageDescriptor: + ... + + def get_layer_view(self, layer_id: int) -> ViewT: + ... + + def write_kv(self, layer_id: int, token_indices: torch.Tensor, kv: WriteT) -> None: + ... + + def copy_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + ... + + # layout 描述该 State 在外部页中的位置及本地轴到页内轴的映射。 + def write_to_cpu( + self, token_indices, cpu_buffer, cpu_page_indices, + *, layout: BufferTransferLayout, + ) -> None: + raise NotImplementedError + + def read_from_cpu( + self, cpu_buffer, cpu_page_indices, token_indices, + *, layout: BufferTransferLayout, + ) -> None: + raise NotImplementedError + + # 一个明确源 State 到当前 State；具体管理器按来源和分片拆分批次。 + def copy_from_dp( + self, src_state, src_token_indices, dst_token_indices, + *, mapping: StateCopyMapping, + ) -> None: + raise NotImplementedError + + # pd_buffer 为已选中的一页，区域偏移相对该页起点。 + def write_to_pd_buffer( + self, token_indices, pd_buffer, + *, layout: BufferTransferLayout, + ) -> None: + raise NotImplementedError + + def read_from_pd_buffer( + self, pd_buffer, token_indices, + *, layout: BufferTransferLayout, + ) -> None: + raise NotImplementedError +``` + +“必须提供”表示各布局满足相同的操作契约，可以继承正确的公共实现。创建或绑定 buffer 属于初始化；请求使用哪些物理 token 槽位由管理器的 allocator 决定。 + +| 接口 | 完整语义 | +| --- | --- | +| `get_required_bytes(layout_config, storage_spec)` | 在创建大块 buffer 前计算常驻 Tensor 存储需求，包括 KV、动态 scale、由本状态持有的固定参数、packed 数据、布局填充及保留区域。与真实初始化使用相同布局规则；管理器汇总资源需求，单独考虑传输空间和框架分配开销 | +| 初始化 | 创建或绑定实际存储，建立层映射并校验容量、dtype、对齐和配置；随后暴露有效视图。共享底层存储的所有权及其负责区域必须明确 | +| `get_storage_descriptor()` | 返回完整存储描述，供管理器做记账、兼容性检查和传输初始化。描述覆盖本 State 负责的全部数据及固定布局信息，其结构在正常推理期间保持稳定 | +| `get_layer_view(layer_id)` | 返回该逻辑层供匹配 backend 使用的具体视图，包含计算所需 Tensor 和参数。视图引用底层存储，不触发 KV 搬运；没有对应 full-attention 存储的层明确报错 | +| `write_kv(layer_id, token_indices, kv)` | 将输入 KV 写入已分配物理位置，完成该布局需要的编码、量化及关联 scale 更新。输入类型与格式由具体布局明确，不能把浮点 KV、已量化数据和 packed 字节隐式混用 | +| `copy_tokens(src_indices, dst_indices)` | 在同一 State 内复制所有所辖存储层的对应 token 数据，覆盖全部随 token 变化的字段；保持存储表示，不重新量化。只操作本 State 负责的区域，遵守该布局声明的复制粒度与对齐条件 | + +容量、device、层映射及布局标识应从同一份稳定配置和描述中取得，避免多处维护互相矛盾的副本。配置相同且只改变容量时，容量计算和真实 buffer 创建应覆盖相同的数据组成。共享底层内存按实际分配去重记账。 + +#### StorageDescriptor 必须描述什么 + +`StorageDescriptor` 是存储结构与语义的描述，具体类型仍待实现时细化。至少包含： + +| 信息 | 用途 | +| --- | --- | +| 容量、保留索引范围、page size、device | 确认存储覆盖管理器给出的物理索引域 | +| 布局标识/版本、资源身份、层映射、量化配置 | 标识模型角色及资源，记录逻辑层到本地物理层的映射，定位实际存储并检查兼容性 | +| 本 State 的 TP 分片描述 | 记录所属 TP 组及 rank、各字段本地轴到全局 head/channel 的对应和副本关系；不同资源可以使用不同分片方案 | +| 本 State 的数据区域及 Tensor 视图 | 列出 KV、动态 scale、packed 辅助字段，说明 token 维、stride、区域偏移和必要复制粒度 | +| 固定参数 | 列出静态 scale 等解释缓存所需的参数，明确如何绑定及验证兼容性 | +| 传输支持与格式 | 声明已实现的路径、可支持的字段表示、页格式及映射约束；组合页的区域位置由管理器统一安排，通过 layout/mapping 参数传入直接方法 | +| 底层分配与视图的关系 | 区分实际占用和别名视图，避免重复计费或重复/越界复制 | + +该描述的结构只读，里面引用的 Tensor 可以由合法计算操作更新。固定布局信息不等于数据就绪信号；stream/event 同步和前缀关联仍由调用流程负责。 + +完整性以本 State 声明的数据范围为准。每次增加辅助 buffer 或 packed 字段，都需要检查描述、容量计算、写入和复制/传输是否一起覆盖它。公共层不根据 `hasattr(scale_buffer)` 或固定属性名推测资源组成。 + +#### 计算视图与完整存储范围 + +| 布局 | 计算视图的主要内容 | 完整复制需要覆盖的内容 | +| --- | --- | --- | +| 普通 K/V | K、V | 对应位置的全部 K/V | +| 分组 INT4/INT8 | K/V 与动态 scale | packed/量化 K/V 加上对应动态 scale | +| 静态 FP8 | 以正确数值类型解释的 KV 及静态 scale | 本地 token 复制只搬 KV；跨资源使用须验证或传递固定量化参数 | +| MLA | 压缩 KV 或其所需分段 | 本 State 负责的完整压缩 KV 表示 | +| 带 indexer 的 packed 布局 | attention KV 子视图、indexer key/scale 子视图 | KV 及全部相关 packed 辅助数据；按实际存储区域处理别名 | + +`get_layer_view()` 可以返回 `DenseKVView`、`QuantizedKVView`、`MLAView` 等具体小对象。视图类型由模型/backend 在初始化绑定时确认。各布局不必返回相同 Tensor 个数，也不必填充包含所有可选字段的大对象。 + +带 indexer 的 packed 布局中，attention 与 indexer 可以分别读取 buffer 的不同区域。因此不能用 attention 的读取视图代替完整复制范围。类似地，静态 FP8 的物理存储 dtype 和计算时解释的 dtype 可以不同，视图必须明确相应参数。 + +`write_kv` 更新其输入所描述的 KV 及关联字段；indexer 池化等专用 kernel 可以通过专用视图写入另一区域。管理器与执行流程在相关写入都满足同步条件后才允许复用或发布。融合写入路径可以直接使用预先绑定的布局视图，同时遵守相同的数据表示与完整性约定。 + +#### 基类与子类的分工 + +| 部分 | 基类可复用逻辑 | 子类必须负责的内容 | +| --- | --- | --- | +| 初始化 | 生命周期、配置与映射校验、稳定视图管理 | `_create_storage(...)`：实际 Tensor/区域和完整描述；`get_required_bytes(...)`：同布局的容量计算 | +| 层访问 | 逻辑层到存储层的映射、视图缓存 | `_make_layer_view(storage_layer_id)`：该布局的计算视图 | +| 写入 | 参数契约、层映射和公共调用入口 | `_write_kv(...)`：编码/量化/写入 kernel 及输入格式校验 | +| 完整复制 | 索引操作契约、公共调用入口 | `_copy_tokens(...)`：布局正确的复制，或显式复用与其描述兼容的公共工具 | +| 传输 | 定义 CPU/DP/PD 直接读写接口，未支持的操作显式报错 | 覆盖所支持的读写方法，直接使用自身 storage 调用相应 kernel | + +受保护方法名为候选。一个子类可以复用满足其布局要求的 kernel；通用复制工具只有在理解 token 维、stride、数据区域和别名关系时才适用。复杂 packed 布局应由专用实现处理，公共层不尝试仅凭 Tensor shape 推导复制语义。 + +#### 索引、异步执行与生命周期 + +`storage_spec` 明确可分配容量、保留区域和 device；HOLD 索引由管理器决定，State 为其提供存储。构造时完成映射和存储配置，正常推理期间保持 buffer 地址稳定。 + +`layer_id` 的命名空间在模型绑定时约定。两个逻辑层若共享同一物理层，计算视图可指向同一存储，完整复制和容量统计按实际区域去重。`token_indices` 为已分配的物理 token 索引，逻辑请求映射仍由 `ReqManager` 管理。 + +写入与复制接受与资源同设备的一维整数索引 Tensor。源和目标索引长度相同，目标唯一；基础复制先约定源目标不重叠，重叠移动另行定义。索引还须满足具体布局的对齐和粒度条件。 + +数据操作在调用方当前 CUDA stream 上提交，方法返回不代表设备操作完成；跨 stream 消费由调用方建立依赖。销毁存储前，管理器等待相关计算和传输结束并解除推理视图及跨进程存储引用。热路径不支持隐式 resize；需要重建存储时必须重新绑定视图及相关执行结构。 + +#### CPU / DP / PD 直接操作 + +三种路径直接作为 `FullAttentionState` 的方法暴露，由子类覆盖实现。接口在数据与物理索引之外显式接收位置映射，描述复杂资源组合及 TP 分片对应关系。`BufferTransferLayout` 和 `StateCopyMapping` 都是只读数据参数，不包含算子、buffer 所有权或执行方法。 + +| 接口 | 数据方向与位置参数 | +| --- | --- | +| `write_to_cpu(..., layout)` | 当前 State → CPU 缓存；layout 描述外部页区域、层/轴映射和写入副本选择 | +| `read_from_cpu(..., layout)` | CPU 缓存 → 当前 State；layout 描述该目标 State 从页中读取的对应范围 | +| `copy_from_dp(src_state, src_token_indices, dst_token_indices, mapping=...)` | 一个明确源 State → 当前 State；mapping 描述两侧字段、物理层及分片范围的对应 | +| `write_to_pd_buffer(..., layout)` | 当前 State → 已选中的 PD 页；按 layout 写入指定区域 | +| `read_from_pd_buffer(..., layout)` | 已选中的 PD 页 → 当前 State；按 layout 恢复对应区域 | + +`write_to_*` 表示将本 State 的数据写到外部区域，`read_from_*` 表示从外部区域读入本 State。DP 的源明确到具体 State，具体管理器按源 DP/TP/资源及对应 token 分组；一个目标分片需要多个来源时，以多次直接调用覆盖，不假定源 rank 与目标 rank 一一对应。 + +基类对未支持的操作抛出 `NotImplementedError`，子类覆盖自身支持的路径，不能用空实现表示成功。`AttResourceManager` 在启用对应功能时校验所需方法和布局映射。具体实现必须声明其支持的组合；不支持的量化、重分片、非连续布局或部分页操作须明确拒绝。 + +每次调用完整处理其映射涉及的数据及关联的动态 scale、packed 字段；一次组合传输可以由若干次调用完成，上层检查所有必需数据区域均已完成。固定参数按约定传递或验证兼容性。位置映射的组成、单位及生成责任见 6.5 节。 + +方法统一在调用方当前 CUDA stream 上提交操作，返回 `None`。若内部使用其他 stream，子类须将完成依赖汇合到调用方 stream，使调用方随后记录的 event 覆盖本次操作。网络收发、任务排队和完成通知由调度层负责。 + +形状相同不能证明不同模型的 KV 语义可互换。模型/资源身份和 token 前缀匹配由管理器与缓存流程校验；State 检查布局及固定参数兼容性。跨布局重编码需要明确实现，本地 `copy_tokens` 不承担这一职责。 + +### 6.2 AttResourceManager：内部资源的聚合操作与数据传输 + +`AttResourceManager` 对内部资源组合提供统一的数据操作。一个模型可以包含多个 FullAttentionState、KDAState、GDNState 或其他状态实现；token 复制覆盖 full-attention 参与集合，清零操作覆盖需要初始运行态的参与集合，状态页操作覆盖指定大小页域的参与集合。外部调用方无需知道具体数量、类型或成员名称。 + +数据操作的依赖方向是 `radix cache / CPU cache / 其他调用方 → AttResourceManager → 各 State`；槽位分配与归还由调用方直接访问管理器持有的共享 allocator。状态页接口不接收 radix 节点、树或前缀匹配对象，也不读取 radix 内部状态；radix 只保存并使用管理器提供的页类别、槽位 ID 和状态长度等数据。状态页在未启用 radix 时同样可以保存、恢复或用于 CPU 缓存。 + +#### 数据操作入口 + +| 操作 | 接口 | 聚合语义 | +| --- | --- | --- | +| full-attention token 复制 | `copy_tokens(src_indices, dst_indices)` | 聚合复制所有参与 FullAttentionState 的对应物理 token 数据，覆盖各自完整存储字段 | +| 空前缀运行态清零 | `reset_state(req_indices)` | 将指定槽位的全部参与运行态恢复到长度 0 的初始状态，覆盖递推状态、控制信息与 indexer tail | +| 无 tail 快照的前缀恢复 | `reset_tail_state(req_indices, state_token_lens)` | 先聚合校验所有 tail 的恢复条件，再按对应长度重置 tail，保留其他已恢复状态 | +| 状态页保存 | `save_pages(...)` | 将指定运行态保存到全部参与 State 的同类 pool、同一组槽位 | +| 状态页校验 | `has_page(page_kind, slot_id, state_token_len)` | 全部参与 State 在该槽位都有匹配长度的有效页时返回 True | +| 状态页恢复 | `restore_pages(...)` | 先整体校验，再将各 State 保存的页恢复到对应运行态位置 | +| 状态页失效 | `invalidate_pages(page_kind, slot_ids)` | 撤销全部参与 pool 中指定槽位的数据有效性，保留分配 | +| PD 写出 / 读回 | `write_to_pd_buffer(...)` / `read_from_pd_buffer(...)` | 组合对应 State 的数据，按指定 KV 或运行态范围打包、还原 | +| DP 复制 | `copy_from_dp(...)` | 按显式来源、字段、层和分片映射，复制内部参与 State 的数据 | +| CPU 写出 / 读回 | `write_to_cpu(...)` / `read_from_cpu(...)` | 组合 KV、状态页与长度元数据，读写各 State 在外部页中的区域 | + +#### 必须实现的聚合接口 + +```python +from abc import ABC, abstractmethod +from collections.abc import Mapping +from typing import Literal + +PageKind = Literal["big", "small"] + + +class AttResourceManager(ABC): + full_att_allocator: KVPageAllocator | None + bounded_att_big_page_allocator: CheckpointSlotAllocator | None + bounded_att_small_page_allocator: CheckpointSlotAllocator | None + + @abstractmethod + def copy_tokens( + self, src_indices: torch.Tensor, dst_indices: torch.Tensor, + ) -> None: ... + + # 空前缀初始化；req_indices 只定位已分配的运行态槽位。 + @abstractmethod + def reset_state(self, req_indices) -> None: ... + + # 无 tail 快照的前缀恢复；只校验并重置 tail。 + @abstractmethod + def reset_tail_state(self, req_indices, state_token_lens) -> None: ... + + @abstractmethod + def save_pages( + self, req_indices, state_token_lens, *, page_kind: PageKind, slot_ids, + ) -> None: ... + + @abstractmethod + def has_page(self, page_kind: PageKind, slot_id, state_token_len) -> bool: ... + + @abstractmethod + def restore_pages( + self, req_indices, expected_token_lens, *, page_kind: PageKind, slot_ids, + ) -> None: ... + + @abstractmethod + def invalidate_pages(self, page_kind: PageKind, slot_ids) -> None: ... + + # CPU 联合页：KV + 可选大小页快照 + 可选 tail。 + @abstractmethod + def write_to_cpu( + self, token_indices, cpu_buffer, cpu_page_indices, + *, layouts: Mapping[str, BufferTransferLayout], + page_kind: PageKind | None = None, slot_ids=None, + state_token_lens=None, req_indices=None, + ) -> None: ... + + # 一次读取同一前缀的有序页面；tail 只恢复最后一页到 req_idx。 + @abstractmethod + def read_from_cpu( + self, cpu_buffer, cpu_page_indices, token_indices, + *, layouts: Mapping[str, BufferTransferLayout], + page_kind: PageKind | None = None, slot_ids=None, + expected_token_lens=None, req_idx: int | None = None, + ) -> None: ... + + # 一个明确来源管理器到当前管理器，按来源与 TP 分片分批调用。 + @abstractmethod + def copy_from_dp( + self, src_manager: "AttResourceManager", + src_token_indices, dst_token_indices, + *, mappings: Mapping[str, StateCopyMapping], + src_req_indices=None, dst_req_indices=None, state_token_lens=None, + ) -> None: ... + + # pd_buffer 始终是已选中的一页，kind 沿用 kv / att_state 两类任务。 + @abstractmethod + def write_to_pd_buffer( + self, pd_buffer, *, kind: Literal["kv", "att_state"], + layouts: Mapping[str, BufferTransferLayout], + token_indices=None, req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: ... + + @abstractmethod + def read_from_pd_buffer( + self, pd_buffer, *, kind: Literal["kv", "att_state"], + layouts: Mapping[str, BufferTransferLayout], + token_indices=None, req_idx: int | None = None, + expected_token_len: int | None = None, + ) -> None: ... +``` + +以上方法均为基类的抽象契约，由具体管理器实现或继承兼容的聚合实现。未支持的 CPU/DP/PD 路径须明确报错，启用相应功能时检查参与 State 的完整支持能力。未启用某种页分配域时，`has_page` 返回 False，保存、恢复和失效该域则明确报错；纯 full-attention 模型无需创建大小页 pool，其 `reset_state` 可以为空操作；没有 indexer tail 的组合，其 `reset_tail_state` 也可以为空操作。 + +`page_kind="big"` 对应 `bounded_att_big_page_allocator` 和参与 State 的 `big_pages_pool`；`"small"` 对应小页 allocator 与 `small_pages_pool`。两者是独立分配域，即使槽位数字相同，也不是同一份资源。 + +调用方根据 `page_kind` 直接选择相应 allocator；该成员为 None 时不能申请该域。分配与归还使用 allocator 自身的接口，管理器不再增加包装方法。allocator 只管理空闲索引，不修改 State 的页数据或有效长度。 + +`req_indices` 只定位运行态存储，不传递请求对象。保存/恢复时，请求状态槽位、页槽位和长度逐项对应；目标槽位唯一且已分配。长度表达实际已接受状态的位置，不能通过传入任意长度把当前运行态改称为历史状态。 + +#### 空前缀运行态初始化 + +`reset_state(req_indices)` 用于没有任何可恢复缓存、从前缀长度 0 开始计算的场景。调用方已通过 ReqManager 取得请求槽位，并确认这些槽位此前的计算和传输使用结束；管理器只按传入索引初始化数据，不查询缓存命中情况或请求对象。 + +- 对全部参与的 KDA/GDN 等状态调用 `reset_requests(req_indices)`，将 Conv/SSM 恢复到零初值，并重置 MTP 游标及影响后续读取的候选信息。未来算法按自身契约恢复空前缀初态。 +- 对全部参与的 IndexerTailState 调用 `reset_requests(req_indices, state_token_lens=zeros)`，其中 `zeros` 与请求槽位等长；K-pool 清零原生 ring，其他实现同时初始化必要控制字段。 +- full-attention 没有需要清零的空前缀递推状态；该操作不遍历或清空历史 KV，不更改请求映射，也不失效大小页快照或归还索引。纯 full-attention 管理器合法地不做数据操作。 +- 目标请求槽位须唯一且在有效范围内，main/draft 所需运行态全部初始化后才可使用。重复执行在没有并发访问时具有相同结果；方法返回仍遵守 stream 完成约定。 + +这是长度 0 的数据初始化接口，不接收非零前缀长度。命中非零前缀时使用 `restore_pages` 或 CPU/DP/PD 读回；没有 tail 快照时通过 `reset_tail_state` 按长度校验并初始化 tail；不能在部分资源恢复后调用此方法，否则会清除已经恢复的运行态。 + +#### 无 tail 快照时按前缀长度重置 + +`reset_tail_state(req_indices, state_token_lens)` 用于已选定恢复前缀、但没有 tail 快照的情况，支持长度 0 和满足各实现恢复条件的非零长度。请求槽位与长度逐项对应，目标槽位唯一且已分配；该接口仅定位数据，不读取请求映射或缓存节点。 + +具体管理器先对全部参与的 IndexerTailState 和全部目标长度调用 `can_restore_without_tail(L)` 校验，全部通过后才逐个调用底层 `reset_requests(req_indices, state_token_lens)`。例如 K-pool 要求 `L % P == 0`；main/draft 各按自己的 P 检查。任意条件不满足时在修改任何 tail 前报错，调用方须使用真实 tail 快照或其他合法恢复点，不能强制清零后继续。 + +该操作只重置 tail 及其必要控制信息，保留已恢复的 Conv/SSM、MTP 游标、KV 和大小页快照。没有 tail 的管理器合法地不做数据操作。调用方保证此前使用已结束，并在相关数据操作完成后继续计算；这不是运行中 MTP 回退的替代操作。 + +前缀页命中时，调用方恢复 KV 和 linear 页后，调用本接口处理未保存的 tail。CPU 最后一页已有 tail 快照，或 DP/PD 已传入完整 tail 时，直接使用该快照，不再调用本接口重置它。恢复条件检查内聚在此方法中，无需调用方访问内部 State。 + +#### full-attention 的聚合 token 复制 + +`copy_tokens(src_indices, dst_indices)` 在本管理器持有的 full-attention 存储内部执行位置复制。参与集合在初始化时按共享 token 索引域确定，可以包含 main/draft 的不同布局，也可以包含同一模型的多个 FullAttentionState。管理器将同一组物理索引交给各参与 State 的 `copy_tokens`；每个 State 在自己的存储中从源位置复制到目标位置，不在不同模型之间互拷数据。 + +- 复制覆盖全部参与存储层的 KV、动态 scale 和 packed indexer key/scale 等随 token 变化的数据，保持各自存储表示；共享物理区域的别名按所有权去重。 +- 源和目标为逐项对应的一维整数索引 Tensor，长度相等、目标唯一，且满足各布局的复制粒度与对齐要求；基础契约要求源目标不重叠。源数据须已就绪并保持有效，目标物理位置须已分配且允许写入。 +- 只搬运 full-attention token 数据，不分配或回收索引、不修改 ReqManager 映射，也不复制 linear 运行态、大小页快照或 IndexerTailState 的原始 tail。 +- 调用前校验全部参与 State 的参数和能力；无 full-attention 参与者时明确报错。方法在当前 stream 提交并返回 None，全部参与复制完成前不能消费目标数据；中途失败时目标视为未就绪。 + +例如 main 和 draft 各有一个 FullAttentionState，复制物理位置 `[10, 11]` 到 `[20, 21]` 时，一次管理器调用会分别搬运 main 和 draft 各自的对应数据。共享的是位置变换，不要求两者具有相同存储布局。 + +#### 多个 State 的参与规则 + +具体管理器在初始化时确定每种页类型的完整参与集合，集合按模型组合固定，不由每次调用临时选择。它可以包含同一模型中的多种 linear attention，也可以包含 main/draft 中应共同保存的状态。内部成员如何命名和组织由模型决定，不增加固定资源成员或公共 State 注册层。 + +例如，一个管理器同时管理两个 KDAState 和一个 GDNState，大页槽位 17 对应： + +```text +bounded_att_big_page_allocator 分配一次 slot_id = 17 + ├── KDAState A.big_pages_pool[17] + ├── KDAState B.big_pages_pool[17] + └── GDNState C.big_pages_pool[17] +``` + +调用 `save_pages(..., page_kind="big", slot_ids=[17])` 保存三个 State；`restore_pages` 恢复这三个 State。各 State 的 buffer 形状、dtype 和单槽字节数可以不同，共享的是槽位索引及对应状态位置。 + +`has_page("big", 17, L)` 是全部参与者的校验结果之合取：每个 State 都必须存在对应页，且 `state_token_lens[17] == L`。任意一项缺失、无效或长度不一致均返回 False；未配置参与 State 时不能因空集合的 `all(...)` 得到 True。应参与的 State 缺少 pool 属于初始化配置错误，不能通过跳过它使校验通过。 + +full-attention 历史 KV 和没有独立 pool 的 IndexerTailState 不属于这组状态页操作。`has_page` 只证明参与 pool 的页记录完整，不证明历史 KV、tail、前缀身份或异步执行全部就绪。需要同时恢复这些数据的调用方，使用相应资源操作并检查共同恢复条件。 + +#### 保存、失效与直接归还槽位的完整性 + +- 调用方只从共享 allocator 分配一次槽位，全部参与 pool 覆盖同一索引范围。分配成功仅表示槽位可用，不表示其中已有有效状态。后续数据操作失败时，由调用方在完成失效与同步后决定重写或归还槽位。 +- `save_pages` 先验证参数、参与集合和目标可写性，再失效本次目标页并调用各 State 保存。只有所有参与者都已捕获对应位置的数据，页记录才完整；一项失败时将本次目标槽位在全部参与 pool 中标为无效，避免留下可误用的部分结果。 +- `restore_pages` 在任何运行态写入前校验全部源页及预期长度。复制阶段失败时，目标状态视为未就绪，源页保持有效；调用方不得消费半恢复的数据。 +- `invalidate_pages` 只撤销全部参与 pool 的逻辑有效性，不回收槽位、不清零全部 buffer，也不改运行态。失效的槽位仍可由持有者再次写入。 +- 归还槽位前，调用方必须解除相关使用引用并满足完成条件，再调用管理器的 `invalidate_pages`，最后直接向对应 allocator 归还一次。只归还指定大小页域的槽位，不操作相同数字的 KV 索引或请求槽位；不能绕过页失效而直接把有效槽位放回空闲集合。 + +保存时全部参与 State 必须实际具备所要求位置的状态；不能从较新的状态补造较早快照。main/draft 可以分别计算，但聚合操作的完整性条件不变；若具体实现分阶段捕获参与者，未全部完成前必须保持聚合页不可用。 + +GPU 数据操作遵守当前 stream 的依赖约定。长度记录或 `has_page=True` 不代表异步复制已完成；调用方在读取、发布或释放前确认完成。数据操作失败后，等待已提交操作结束再复用或释放相关存储。 + +共享域跨多个管理器实例时，调用方使用同一个 allocator 实例分配和归还，统一失效流程必须覆盖所有实际参与 pool。只有所有使用方均已结束且全部相关页失效后，才归还该槽位一次；不能因某个管理器的数据已失效就提前回收。allocator 的创建和销毁责任仍由装配时约定的唯一所有者承担。 + +#### CPU / DP / PD 聚合接口的参数与数据路径 + +`layouts` 是按完整 `resource_id`（包含 main/draft 角色）索引的 `BufferTransferLayout` 映射；`mappings` 按目标 `resource_id` 索引，每个 `StateCopyMapping` 同时标明明确的源资源与源目标布局。具体管理器在传输初始化时生成并缓存这些普通只读数据，调用方选择对应格式/对端的映射集合，无需查找内部 State。管理器查验后将每项传给匹配 State 的直接方法，不增加 Transfer 类或可执行 View。 + +CPU/PD 的映射集合必须覆盖所选页格式的全部参与资源；不能通过遗漏某个布局静默丢弃必需状态。DP 一次调用可以只覆盖一个来源的分片，完整迁移要求全部调用的覆盖范围合起来满足资源恢复要求。源目标语义、字段范围、别名、TP 副本写入者与已有数据的就绪条件均需在写入前校验。 + +**CPU 联合页** + +包含 linear attention 的模型，每个已发布 CPU 缓存页都必须包含全部参与 linear State 对应页终点的有效快照，且长度一致。写出前检查这一条件，读回时校验格式和预期长度；不完整状态不能发布。纯 full-attention 模型仍使用自身的 KV 页格式。接口不增加缺失 linear 快照的 CPU 页语义，也无需为选择有效快照终点新增元数据查询方法。 + +- `token_indices` 是 full-attention 的物理 token 索引，`cpu_page_indices` 是 CPU 页池的页 ID。索引按列表顺序排列，从各页槽位 0 开始；长度与页数遵守 6.5 节的整页/尾页约定。页列表按逻辑前缀顺序排列，不要求物理页 ID 连续或递增。只传状态的格式使用空 token 索引序列。 +- 格式包含 linear 页时，`page_kind` 与 `slot_ids` 必须同时提供，选择已分配 pool 槽位；`slot_ids` 和写出的 `state_token_lens` 或读回的 `expected_token_lens` 均与 CPU 页 ID 逐项对应，每项长度有效且不得为 None。没有 linear 页时，`page_kind/slot_ids` 为 None；纯 KV 传输且无需 tail 恢复时，长度参数也可为 None。 +- 写出带 tail 的格式时，`req_indices` 和 `state_token_lens` 与 CPU 页 ID 逐项对应，各页的 tail 必须在其实际恢复位置仍可捕获时保存。没有 tail 段时，写出方法的 `req_indices=None`；不能用当前 tail 补写历史页。只有 tail 而没有 linear 页时，`page_kind/slot_ids` 为 None。 +- `read_from_cpu` 一次处理同一前缀的有序页面，将全部 KV 写到给定物理位置，将各页 linear 快照写入对应 pool 槽位。linear 运行态随后由调用方使用最后一个槽位及 `expected_token_lens[-1]` 显式调用 `restore_pages` 恢复。 +- 需要 tail 恢复时，调用方传入单个目标 `req_idx`，并提供非空的页列表和逐页 `expected_token_lens`。管理器只使用 `cpu_page_indices[-1]` 和 `expected_token_lens[-1]`，将最后一页的 tail 恢复一次；中间页的 tail 不写入运行态。底层 IndexerTailState 的批量接口使用长度为 1 的请求、页 ID 和长度序列调用，无需改动其签名。 +- CPU 格式没有 tail 段时，管理器针对最后一页的长度调用 `reset_tail_state([req_idx], [expected_token_lens[-1]])`；恢复条件须在任何目标写入前检查。模型没有 tail 时，`req_idx=None`。多个目标前缀按前缀分别调用，不将不同请求的页混成一个列表。 +- 读入 pool 前统一失效本次目标页；任一参与资源读入失败时，所有本次目标 pool 页保持无效，已写入的 KV 或 tail 不得作为就绪数据使用。CPU 写出失败时同样不得发布不完整联合页。 + +**DP 直接复制** + +`src_manager` 明确来源 DP/TP 存储实例，也可以是传输进程中已绑定共享存储的管理器。源和目标 token 索引等长、逐项对应；`src_req_indices`、`dst_req_indices` 与 `state_token_lens` 作为一组提供，长度相等，用来迁移当前已接受的 linear 运行态和完整兼容 tail。token 索引与请求槽位是不同的寻址维度，无需等长，二者的前缀对应关系由调用方保证。 + +只复制 KV 时，运行态三项均为 None；只复制运行态时，两个 token 索引序列均为空；提供运行态参数却缺少所需映射应报错。一个目标分片需要多个来源时，调用方按来源分批调用，具体管理器将各项 mapping 交给对应 State。每次只写所声明区域，不清零其他批次已恢复的数据；全部必需区域及控制信息就绪后才可继续计算。DP 不自动生成历史大小页快照。 + +**PD 页读写** + +| `kind` | 必需参数 | 处理对象 | +| --- | --- | --- | +| `"kv"` | `token_indices`；运行态槽位和长度参数为 None | 按单页容量读写全部参与 full-attention 的 KV、scale 和 packed 字段 | +| `"att_state"` | `req_idx`，以及写出的 `state_token_len` 或读回的 `expected_token_len`；`token_indices=None` | 读写该槽位的已接受 linear 状态及所需完整 tail，不经过大小页 pool | + +PD 参数按上述两种组合校验，不允许混用；`req_idx=0` 和长度 0 都是有效值，不能当成缺失参数。`pd_buffer` 是已选中的一页基址，各资源偏移只通过 layout 应用一次。读回前按协议元数据校验资源身份、格式与预期长度；页内字段或与该页绑定的传输元数据必须能提供实际保存长度，不能仅信任调用方给出的预期值。无运行态资源的管理器不接受 `att_state` 任务。 + +CPU/DP/PD 方法均返回 None，只执行对应数据操作，不申请索引或修改请求映射。GPU 操作在调用方 stream 提交，内部使用其他 stream 时须合并完成依赖；返回不代表全部数据已可用。传输调度负责网络收发和完成事件，CPU 同步读 pinned 数据前还须确认之前的 GPU 写入已完成。 + +#### 调用方与资源实现的边界 + +radix 负责前缀匹配、节点引用及淘汰时机，只把页类别、槽位 ID 和长度关联到自己的节点；保存、检查、恢复和失效状态页时调用上述通用接口，槽位分配与归还直接调用管理器持有的相应 allocator。管理器不生成 radix 节点数据、不返回前缀匹配结果，也不改变请求到 KV 的映射。 + +CPU cache、DP、PD 通过直接数据接口操作同一组内部资源，具体管理器负责组合布局、字段覆盖和分片对应。请求阶段、MTP 接受位置选择及 ReqManager 映射维护由外部流程组织。管理器中的 `req_idx` 仅是运行态数据地址的一部分。 + +### 6.3 模型推理使用具体资源视图 + +建议在模型或 attention backend 初始化时，将所需资源绑定为具体推理视图。视图包含该实现需要的 Tensor、层映射、MTP 游标或资源操作引用，模型 forward 可以直接使用这些引用。 + +视图结构由模型和 backend 的实际需求决定。例如普通 full-attention 层需要 KV 视图，KDA/GDN backend 从对应状态对象获取自身所需视图，GLM indexer 需要 packed KV 与 tail 的关联视图。一个模型可以按层绑定多个不同视图；main/draft 的视图类型也可以不同。 + +模型专属初始化代码直接绑定对应 State 或层视图，绑定方式由具体模型和 backend 决定。管理器基类不规定统一的模型视图获取方法，也不规定固定的 `main_*` / `draft_*` 字段。视图引用的存储由具体管理器组织，请求映射由 ReqManager 持有。 + +模型初始化时为 infer state 和 attention backend 绑定对应 State 或层视图，推理阶段直接读取。每层 forward 无需按字符串检索资源或反复判断模型组合。 + +### 6.4 调用边界与独立传输进程 + +单资源操作直接调用“具体 State 的读写方法 → kernel”；涉及多个资源时，由模型专属 `AttResourceManager` 的具体方法组织调用。该管理器确定参与资源、布局映射和已分配索引，layout/mapping 只作为数据参数传入。推理用的 `get_layer_view()` 保留其计算数据访问用途。 + +具体 State 负责导出重建自身所需的 GPU 句柄、实现标识、布局配置、所属 TP 组/rank 的分片描述和固定参数，并实现导入后的存储绑定。模型专属 `AttResourceManager` 按模型角色及资源身份组织这些描述。独立传输进程按描述使用相同的具体 State 类绑定共享存储，随后直接调用其读写方法；多资源的组合规则由该模型管理器的具体实现提供。 + +这种绑定使用 State 的“创建或绑定存储”初始化语义，不重复分配 KV，也不导出 allocator、整个模型或请求管理器。传输进程中的 State 持有已导入 buffer 的引用，不负责回收原进程的物理槽位；原持有方与调度层共同保持存储有效。跨进程导出仅是初始化过程，具体序列化字段在实施时按实际 State 细化。网络 PD 路径只发送约定数据和必要元信息。 + +### 6.5 复杂 attention 组合与 TP 的位置映射 + +#### 位置映射的范围 + +传输需要分别描述逻辑层到物理层的映射、组合页中的资源区域、各字段的 TP 分片与副本关系，以及本次操作的 token/page 索引。KV、scale、Conv、SSM 和 tail 可以使用不同的 dtype、stride、对齐和轴划分，不能统一简化为一个 `offset` 或一个全局 `tp_rank`。 + +#### 信息由谁维护 + +| 信息 | 维护位置 | 示例 | +| --- | --- | --- | +| 本地存储布局与分片事实 | 具体 State，在初始化时确定并通过 `StorageDescriptor` 描述 | 物理层映射、Tensor stride、packed 字段偏移、字段的全局 head 范围、TP 副本关系 | +| 组合页布局及资源间分区 | 模型专属 `AttResourceManager` 根据参与资源和 CPU/PD 格式组织，结合传输端确认的格式生成 | main/full、draft/full、Conv、SSM、tail 各区域的起点、大小、对齐和总页跨度 | +| 某个 State 到外部页、或两个 State 之间的位置映射 | 管理器在传输初始化时生成并缓存；具体 State 校验是否支持 | `BufferTransferLayout`、`StateCopyMapping`；按资源、路径、对端配置分别保存 | +| 本次实际位置和长度 | 请求流程/传输调度显式传入 | 已分配的 token 索引、CPU 页 ID、所选 PD 页、有效 token 数 | + +资源的全局身份须包括模型角色及资源标识，不能只用层号识别 main/draft。一个模型可以有多个同类 State；管理器组织内部组合，不增加固定的六个成员变量。共享 allocator 只保证索引域，不能推导层、head、字段或传输字节位置。 + +同一 State 可能服务 CPU cache 和多个 PD 对端，不应把组合页的绝对位置固化成唯一的 `self.pd_offset`。本地存储事实固定在 State，外部位置由该次传输选择的 layout/mapping 提供。算子不得用进程全局的 TP rank 替代被操作 State 的分片描述，尤其是在传输进程导入多个 rank 的 State 时。 + +#### BufferTransferLayout：页位置的纯数据参数 + +下面是该描述所需的数据，不要求引入执行器或类型继承体系。可以用冻结的数据类及元组表达；一个对象描述一个具体 State 在某种组合页格式中的位置。 + +| 字段 | 语义与单位 | +| --- | --- | +| `resource_id`、`format_id`、`format_version` | 关联模型角色/资源及双方确认的页格式 | +| `state_layout_id` | 绑定本地 State 的布局配置标识，覆盖本地格式版本、物理层映射和 TP 分片身份；这些配置改变后应重新生成标识及映射 | +| `tokens_per_page` | 含 KV 的页格式为正整数 token 容量；纯递推状态页为 `None`，按请求状态或 checkpoint 槽位处理 | +| `page_stride_bytes` | 相邻物理页起点的距离，以字节为单位；CPU 页池寻址使用此值 | +| `regions` | 本 State 参与传输的物理字段或片段列表，同一字段可分为多个映射项 | +| 每项的 `field`、`byte_offset`、`dtype`、`axes`、`shape`、`strides_bytes` | 标识 State 中的字段；其余描述外部页区域。偏移相对页起点，axes 命名各维度，shape 用元素数量，stride 明确为字节 | +| 每项的 `layer_pairs` | `(local_storage_layer, buffer_layer)`，两者分别属于 State 物理层和该页区域的层坐标；支持不连续层 | +| 每项的 `axis_slices` | `(axis, local_start, buffer_start, length)`，描述 head、channel、dim 等轴上的连续对应；一个 region 内同一轴只出现一次，同轴不连续范围拆成多个 region | +| 每项的 `write_enabled` / `read_enabled` | 指定该副本对该区域是否实际写入/读取；写入应有唯一来源，恢复时可读取到多个目标副本 | + +建议先采用以下两个冻结数据类，均无搬运方法。这里是设计代码，字段名称和类型可随实现细化： + +```python +from dataclasses import dataclass, replace +import torch + + +@dataclass(frozen=True) +class TransferRegion: + field: str # 对应 State 存储描述中的字段，如 k/v/scale/packed + byte_offset: int # 相对外部一页起点的字节偏移 + dtype: torch.dtype # 外部区域的数据类型 + axes: tuple[str, ...] # 如 (token, layer, head, dim) + shape: tuple[int, ...] # 外部区域在各轴上的完整尺寸，不含页池的 page 轴 + strides_bytes: tuple[int, ...] # 与 axes 顺序对应，全部使用字节单位 + layer_pairs: tuple[tuple[int, int], ...] # (本地物理层, 外部区域层) + axis_slices: tuple[tuple[str, int, int, int], ...] # (轴, 本地起点, 外部起点, 长度) + write_enabled: bool = True + read_enabled: bool = True + + +@dataclass(frozen=True) +class BufferTransferLayout: + resource_id: str # 含模型角色，如 main/full_0 + format_id: str # CPU/PD 页格式标识 + format_version: int + state_layout_id: str # 本地 State 的布局及分片配置标识 + tokens_per_page: int | None # KV 页为 token 容量，纯状态页为 None + page_stride_bytes: int # 整个组合页的跨度，包含其他资源及填充 + regions: tuple[TransferRegion, ...] +``` + +`field` 在初始化时与 State 自身描述匹配；本地 Tensor、dtype、stride 和字段起点继续由 State 持有。`axes`、`shape`、`strides_bytes` 均描述外部页区域，长度必须相同且轴名唯一。例如 `axes=("token", "layer", "head", "dim")` 明确 stride 元组各项的含义，避免把 CPU 和 PD 不同的维度顺序混用。 + +`layer_pairs` 定义层选择；存在 layer 轴时必须明确列出对应，不默认源目标层号相同。token 轴由本次索引向量驱动，外部页内从槽位 0 开始，不放进 `axis_slices`。其余实际搬运轴在 `axis_slices` 中明确范围，各轴映射与层对按笛卡尔积组成当前 region 的拷贝范围；某些层需要不同 head 映射时拆分 region，避免产生错误的交叉组合。使用多个连续片段表达非连续范围，不引入额外的通用索引解释器。 + +`shape` 表达该外部字段区域的完整范围，而 `axis_slices` 表达当前 State 实际访问的子范围；两者不能混淆。`byte_offset` 已包含组合资源的起点，不再重复加一个资源 base offset。CPU 按页 ID 使用 `page_stride_bytes`；PD 传入已选中的一页，方法验证该页容纳完整布局。 + +下面将前述结构用于 P 端 TP=2、rank 1：本地 K/V 各有 2 个 heads，外部页使用 4 个全局 heads。假设组合页共 64 KiB，其中前 32 KiB 存 main K/V，后 32 KiB 用于其他资源： + +```python +k_region = TransferRegion( + field="k", + byte_offset=0, + dtype=torch.float16, + axes=("token", "layer", "head", "dim"), + shape=(16, 2, 4, 64), + strides_bytes=(1024, 512, 128, 2), + layer_pairs=((0, 0), (1, 1)), + axis_slices=(("head", 0, 2, 2), ("dim", 0, 0, 64)), +) +v_region = replace(k_region, field="v", byte_offset=16384) + +layout = BufferTransferLayout( + resource_id="main/full_0", + format_id="example_combined_kv", + format_version=1, + state_layout_id="dense_fp16_v1/main_tp/tp_2/rank_1/layers_3_7", + tokens_per_page=16, + page_stride_bytes=65536, + regions=(k_region, v_region), +) +``` + +示例中的配置标识字符串只为说明关联关系，不规定标识编码方式。K/V 每个完整区域为 `16 * 2 * 4 * 64 * 2 = 16384` 字节；当前 rank 通过 head 映射只写各区域的一半，另一半由另一个 TP rank 写入。逻辑层 `[3, 7]` 到本地层 `[0, 1]` 的关系保留在 State，`layer_pairs` 使用本地物理层编号。 + +例如当前 rank 的本地 head 0、token 在页内位置 0、物理层 0、dim 0，对应 K 区域地址偏移为 `0 + 2 * 128 = 256` 字节，V 为 `16384 + 256 = 16640` 字节。下一物理页按 `65536` 字节递进，不能按本 State 占用的 `32768` 字节递进。 + +D 端 TP=4 的 rank 3 可以复用相同外部页格式及 K/V 区域尺寸，但使用自己的 `state_layout_id`，并把 head 映射改成 `("head", 0, 3, 1)`。页格式与本地映射分别校验，不要求 P/D 的完整 layout 对象相同。 + +字段列表按实际物理表示组织。例如 K、V、scale 可分别描述；packed 数据也可整体描述为一个字节区域。别名区域只搬运一次。Conv 的 Q/K 和 V、SSM 可能采用不同 head 映射，应分别描述，不能把一个资源的 `head_start` 套用于所有字段。数据类型、量化组粒度及索引单位由对应格式约束。 + +具体管理器从 State 的存储描述中取得布局需求和分片事实，分配对齐区域并生成映射；子类只需理解自身字段及支持的映射形式，无须知道前后还有哪些资源。简单连续映射可以预先化成 kernel 参数，通用描述不意味着必须使用逐元素的通用复制 kernel。 + +FullAttentionState 的 token 页操作要求 `tokens_per_page` 为正整数。CPU 参数 `cpu_buffer` 指向完整页池基址；PD 参数 `pd_buffer` 指向已选中的一页基址。调用方可以选择 PD 页，但不再额外切到资源区域后重复应用 `byte_offset`。对于常规 K/V，一个外部地址可以写成： + +```text +buffer.data_ptr() ++ page_id * page_stride_bytes # CPU 为实际页 ID；PD 已选页，取 0 ++ region.byte_offset ++ token_in_page * token_stride_bytes ++ buffer_layer * layer_stride_bytes ++ buffer_head * head_stride_bytes ++ dim * dim_stride_bytes +``` + +`data_ptr()` 已指向传入 Tensor 的起点，不重复加 Tensor 的 storage offset。State 内部源/目标地址使用该 State 自身 Tensor 的基址及 stride；PyTorch 元素 stride 进入字节描述前乘以对应 dtype 大小。非连续区域必须保留真实 stride，不能通过会产生拷贝的 reshape/contiguous 假装仍指向原写入目标。 + +CPU/PD 的 token 页操作统一从外部页内槽位 0 开始，公共接口不设置页内 token 起点参数。令 `i` 为本次索引向量位置： + +```text +# CPU：按给定页 ID 顺序，从每页槽位 0 开始 +page_id = cpu_page_indices[i // tokens_per_page] +token_in_page = i % tokens_per_page + +# PD：传入的 buffer 已选到一页 +page_id = 0 +token_in_page = i +``` + +有效 token 数由索引长度给出。普通 CPU 路径要求整页操作；支持尾页的路径允许最后一页只有前部有效，CPU 页 ID 数量须与覆盖页数一致，PD 单页要求 `len(token_indices) <= tokens_per_page`。checkpoint 对应前缀长度继续由 `state_token_lens` 表达。 + +上述固定的是 token 维度起点。组合资源的 `byte_offset`、TP head/channel 范围以及 State 的实际物理 token 索引仍可非零，应继续保留在各自的位置映射中。 + +#### StateCopyMapping：DP 源目标位置对应 + +DP 不经过 CPU/PD 组合页，因此不使用页字节偏移。`StateCopyMapping` 记录一对明确 State 的字段对应、`(src_storage_layer, dst_storage_layer)`、每个轴的 `(src_start, dst_start, length)` 或索引表，以及资源/布局身份。两侧 buffer 基址、容量和 stride 均来自各自 State；固定参数与编码兼容性分别校验。 + +目标分片跨多个源 TP rank 时，具体管理器按全局层/head/channel 的交集拆成多个映射，逐项调用 `copy_from_dp`，各项写入互不重叠。不同 token 来自不同 DP 时，先按来源分组其索引向量；同一批内逐项配对源与目标 token 索引。对源副本只选择一个合法来源。 + +源与目标布局完全一致、分片一一对应时，可以采用同构布局的复制 kernel，但须先确认源目标 stride 等条件成立。重分片或编码变化由具体子类明确声明支持并提供实现。 + +#### main/draft 与不同 TP 的例子 + +假设 main 的 full-attention 逻辑层为 `[3, 7]`，本地压缩存储层为 `[0, 1]`；一页 16 个 token，全局 4 个 KV heads，每 head 64 个 FP16 元素。以 K、V 分区的示例格式组织 `kv` 页： + +```text +byte 0 : main/full K [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB +byte 16384 : main/full V [16 tokens, 2 layers, 4 heads, 64 dims] # 16 KiB +byte 32768 : draft/full 的数据区域，按 draft 自身格式定义 +``` + +在该示例格式下，main 的 P 端 TP=2 时，rank 1 的本地 head `[0, 2)` 映射到页内全局 head `[2, 4)`；D 端 TP=4 时，rank 3 的本地 head `[0, 1)` 从页内全局 head `[3, 4)` 恢复。两端共享页格式，使用不同的局部分片映射。 + +如果 draft 使用 TP 副本形式的 MLA，draft 区域由选定的 P 端副本写一次，D 端各副本读取该区域。它不沿用 main 的 head 偏移公式。main 与 draft 即使共享 token allocator，也分别取得自己的区域映射。 + +CPU 组合页可以包含 full-attention、checkpoint 和元数据；PD 使用 `kv` 与 `att_state` 两类任务，各自拥有布局，不强制把 KV、Conv、SSM 和 tail 放进同一页。采用按 TP rank 分块的格式时，必须明确其 TP 配置约束；跨 TP 配置恢复需要相应格式和转换能力。 + +#### 初始化与提交时的校验 + +初始化时校验模型/资源身份、格式版本、字段覆盖、dtype 与固定参数、字节对齐、页内地址边界、层/轴映射范围，以及 TP 写入覆盖和副本唯一写入者。各资源的实际写入区域不得相互覆盖；packed 别名按实际物理范围去重。涉及多个布局的描述以实际访问地址为准，不能仅比较 shape。 + +每次数据操作校验 token/page 索引、有效长度和对应页数；源数据须已就绪且保持有效。具体管理器校验所有参与 State、TP 分片及长度元数据；调用方确认写入完成后，才允许发布对应缓存或继续计算。布局格式或参与资源改变时，重新生成映射并重新校验，不能复用旧版本的偏移。 + +### 6.6 子类实现与直接调用示例 + +以下为伪代码，layout/mapping 在初始化时由管理器组织并经具体 State 校验，调用过程保持直接。 + +```python +class DenseFullAttentionState(FullAttentionState): + def write_to_pd_buffer( + self, token_indices, pd_buffer, *, layout, + ) -> None: + dense_kv_to_pd_buffer( + kv_buffer=self.kv_buffer, + token_indices=token_indices, + pd_buffer=pd_buffer, + layout=layout, + ) + + def read_from_pd_buffer( + self, pd_buffer, token_indices, *, layout, + ) -> None: + dense_pd_buffer_to_kv( + pd_buffer=pd_buffer, + kv_buffer=self.kv_buffer, + token_indices=token_indices, + layout=layout, + ) +``` + +CPU 操作按路径选取该 State 的位置映射： + +```python +state.write_to_cpu( + src_token_indices, cpu_buffer, cpu_page_indices, + layout=cpu_write_layout, +) +# 当前 stream 上记录 event；所有资源写入完成后发布 CPU 页 ready。 + +state.read_from_cpu( + cpu_buffer, cpu_page_indices, dst_token_indices, + layout=cpu_read_layout, +) +# 后续计算与写入建立 stream/event 依赖。 +``` + +DP 操作按来源和分片交集执行；多来源时重复调用，批次由具体管理器组织： + +```python +dst_state.copy_from_dp( + src_state, src_token_indices, dst_token_indices, + mapping=src_to_dst_mapping, +) +``` + +PD 两端分别调用；buffer 已选到实际发送/接收页，区域与 TP 位置来自各自 layout： + +```python +p_state.write_to_pd_buffer( + src_token_indices, send_page, + layout=p_layout, +) +# 全部参与资源/分片的打包完成后，调度层发送。 + +# D 端接收完成并建立 stream 依赖后执行。 +d_state.read_from_pd_buffer( + recv_page, dst_token_indices, + layout=d_layout, +) +# 所有目标资源/分片恢复完成后允许 decode。 +``` + +调度层负责网络、同步和生命周期。方法返回不等于数据已经就绪，发送页在网络发送完成前不能复用，源槽位在读取完成前不能释放。KDA/GDN 状态的统一操作见 6.7 节，indexer tail 的操作见 6.8 节；它们共用适用的位置描述，不继承 `FullAttentionState`。 + +### 6.7 LinearAttentionState：KDA/GDN 的统一操作接口 + +`LinearAttentionState(ABC)` 约束线性递推状态的操作，`KDAState`、`GDNState` 等具体子类必须实现相应方法；基类不声明 Conv、SSM 或 MTP buffer，不包装另一层状态对象。模型管理器直接持有子类实例，计算 backend 通过层视图访问所需运行态。 + +```text +LinearAttentionState # 抽象操作契约 +├── KDAState # Conv + KDA 递推状态 + 原生 MTP 布局 +├── GDNState # Conv + GDN 递推状态 + 原生 MTP 布局 +└── 其他线性递推算法的 State +``` + +`bounded_att_*_allocator` 的范围仍可覆盖 sliding-window 等有界状态，但使用该 allocator 不要求继承 `LinearAttentionState`。 + +#### 状态操作约定 + +- forward 通过层视图更新 Conv/SSM；MTP 提交选择已接受状态。提交可以只更新游标，保存、复制及传输必须读取该游标对应的数据。 +- 页保存只接受实际捕获的状态位置。跨过中间边界而未保留其状态时，只能预留槽位，不能发布为有效页。 +- 大小页 pool 保存独立快照，默认使用 CPU pinned memory；DP/PD 操作请求运行态。两者使用不同索引和生命周期。 + +#### KDAState 内部布局示例 + +设线性层数为 `L`，请求槽位容量为 `N`（已经包含 HOLD），本地 head 数为 `H`，head 维度为 `D`，Conv 历史宽度 `W = kernel_size - 1`，MTP 额外候选数为 `S`。该 KDA 布局中，Q/K/V 均为 `H * D` 通道： + +| 存储 | 示例形状 | 内容 | +| --- | --- | --- | +| GPU `conv_state` | `[L, N, 3*H*D, W+S]` | 历史窗口及验证时暂存的候选尾部 | +| GPU `ssm_state` | `[L, N*(S+1), H, D, D]` | 每请求 `S+1` 个递推状态槽位 | +| GPU `mtp_state_index` | `[N]`，int32；S=0 时可为 None | 当前已接受状态的原生选择索引 | +| 某个 pool 的 Conv | `[C, L, 3*H*D, W]` | 一个确定前缀的规范化历史窗口 | +| 某个 pool 的 SSM | `[C, L, H, D, D]` | 同一前缀对应的递推矩阵 | +| pool 的 `state_token_lens` | 长度为 C 的 `list[int | None]` | 每个槽位保存的前缀长度 | + +大、小 pool 的 C 分别取管理器给出的容量；两者保存同种状态快照，区别主要在分配域、保存边界和容量，不是按 token 页大小重复展开运行态。pool 不持有空闲列表。默认使用 CPU pinned 快照存储，具体 device 由 `storage_spec` 声明。 + +GDN 的 Q/K 与 V head 数和维度可以不同，其 Conv/SSM 形状由自身实现确定，不把上述 `3*H*D` 或方阵要求放进基类。KDA/GDN 即使 shape 相同，也不能互换状态内容。 + +#### 必须实现的公共接口 + +所有下列方法均为抽象方法。子类可继承已验证兼容的具体实现；对于未支持的 checkpoint、MTP 或传输功能，也必须明确实现为报错并在存储描述中声明不支持，不能依靠基类空方法静默成功。初始化只要求实际启用的功能可用。 + +```python +from abc import ABC, abstractmethod +from typing import Literal + +PageKind = Literal["big", "small"] + + +class LinearAttentionState(ABC): + @classmethod + @abstractmethod + def get_required_bytes(cls, config, storage_spec, *, device) -> int: ... + + @abstractmethod + def __init__(self, config, storage_spec): ... + + @abstractmethod + def get_storage_descriptor(self) -> StorageDescriptor: ... + + @abstractmethod + def get_layer_view(self, layer_id: int): ... + + @abstractmethod + def reset_requests(self, req_indices) -> None: ... + + @abstractmethod + def commit_accepted_state(self, req_indices, accepted_steps) -> None: ... + + @abstractmethod + def copy_requests(self, src_req_indices, dst_req_indices) -> None: ... + + @abstractmethod + def save_pages( + self, req_indices, state_token_lens, *, page_kind: PageKind, slot_ids, + ) -> None: ... + + @abstractmethod + def restore_pages( + self, req_indices, expected_token_lens, *, page_kind: PageKind, slot_ids, + ) -> None: ... + + @abstractmethod + def has_page(self, page_kind: PageKind, slot_id, state_token_len) -> bool: ... + + @abstractmethod + def invalidate_pages(self, page_kind: PageKind, slot_ids) -> None: ... + + @abstractmethod + def write_to_cpu( + self, page_kind: PageKind, slot_ids, cpu_buffer, cpu_page_indices, + *, layout: BufferTransferLayout, + ) -> None: ... + + @abstractmethod + def read_from_cpu( + self, cpu_buffer, cpu_page_indices, expected_token_lens, + *, page_kind: PageKind, slot_ids, layout: BufferTransferLayout, + ) -> None: ... + + @abstractmethod + def copy_from_dp( + self, src_state, src_req_indices, dst_req_indices, state_token_lens, + *, mapping: StateCopyMapping, + ) -> None: ... + + @abstractmethod + def write_to_pd_buffer( + self, req_idx, state_token_len, pd_buffer, *, layout: BufferTransferLayout, + ) -> None: ... + + @abstractmethod + def read_from_pd_buffer( + self, pd_buffer, req_idx, expected_token_len, *, layout: BufferTransferLayout, + ) -> None: ... +``` + +请求/槽位接口按批量处理，单请求也传长度为 1 的序列；PD 状态页按单请求处理。`req_indices` 是 ReqManager 已分配的请求槽位，`slot_ids` 是调用方从共享 allocator 已分配的状态页槽位；两种索引不能混用，`page_kind` 决定使用哪个 pool。 + +#### 运行态与推理访问 + +| 接口 | 必需语义 | +| --- | --- | +| `get_required_bytes(..., device=...)` | 按指定 device 计算所需 Tensor 字节数；分别覆盖 GPU 运行态/MTP 和配置在该 device 上的大小页 pool。CPU pinned 快照不能计入 GPU 常驻需求；共享物理分配去重，管理器另计传输暂存空间 | +| 初始化 | 创建或绑定运行态、层映射和可选 pool，容量覆盖请求域及共享 checkpoint 域；不创建 allocator，不分配请求号或 checkpoint 槽位 | +| `get_storage_descriptor()` | 描述算法/模型资源身份、完整字段、device/dtype/stride、逻辑层与物理层映射、TP 分片、MTP 布局、pool 及保留槽位；声明功能支持和快照可捕获位置，KDA/GDN 的页保存对象为当前已接受状态 | +| `get_layer_view(layer_id)` | 返回该层计算需要的直接引用；KDA 示例包含 Conv、SSM、MTP 游标及必要布局参数。普通计算及融合 kernel 直接更新这些存储，State 不包装 attention forward；访问非本算法层必须报错 | +| `reset_requests(req_indices)` | 恢复算法初始状态，并重置游标及影响后续读取的候选信息；KDA/GDN 初值为零。不改 ReqManager 映射，不释放 checkpoint | +| `commit_accepted_state(req_indices, accepted_steps)` | 在验证完成后选择各请求已接受的运行态，允许通过游标实现；不承担采样、KV 回收或请求长度更新 | +| `copy_requests(src_req_indices, dst_req_indices)` | 复制当前已接受的全部算法状态到本地目标请求，目标恢复为可继续计算的规范化状态；不复制 checkpoint 关联或未接受候选，不改变源状态 | + +`accepted_steps` 为逐请求接受的、已经执行过 forward 的输入步数，包含本轮基准输入步，不是“接受的 draft token 数”或新采样输出数。取值范围为 `[1, 本轮实际验证宽度]`；KDA/GDN 的原生索引为 `accepted_steps - 1`。实现不应把不存在的零步回滚当作有效操作。未启用 MTP 时，合法的一步提交可为无数据移动的操作。 + +KDA 的提交可以只更新 `mtp_state_index`，下一轮 kernel 根据该游标选择输入状态；所有保存、复制和导出方法都必须使用这一选择。设已接受索引为 j，则导出 Conv 的 `[j:j+W]` 窗口和 SSM 的 `req_idx*(S+1)+j` 槽位。导入时放到目标的标准窗口/基准 SSM 槽，并将游标设为 0，候选区只须保证不被误读；可以通过游标和有效范围隔离，不要求每次全量清零。 + +prefill 的最终状态、普通 decode 的结果和已提交验证状态都可以成为“当前已接受状态”；执行流程必须保证所有相关层已完成且选择信息与数据一致。验证结果尚未提交时，不允许将临时候选作为当前状态发布。 + +#### 状态页：只保存真实存在的前缀状态 + +| 接口 | 必需语义 | +| --- | --- | +| `save_pages(...)` | 将每个请求当前已接受状态规范化写到指定 pool 槽位，保存对应 `state_token_lens`；req、slot、长度逐项对齐，目标槽位唯一且已分配 | +| `restore_pages(...)` | 校验槽位长度等于 `expected_token_lens`，再恢复运行态和算法控制信息；源快照保持有效，目标 MTP 从规范化状态继续 | +| `has_page(...)` | 判断 pool 是否存在、槽位有效且长度匹配；没有 pool 或长度为空时返回 False。仅表示数据记录有效，不表示异步搬运已完成，也不校验 token 前缀身份 | +| `invalidate_pages(...)` | 将指定槽位的长度记录设为 None，撤销逻辑有效性；不归还 allocator，也不要求清零全部数据 | + +`state_token_lens` 由调用方提供实际执行后的位置。它不要求 State 再维护一份请求长度副本，也不能仅凭传入长度把当前状态改称为历史状态；KDA/GDN 的页保存契约只覆盖传入请求的当前已接受位置。chunk 跨过中间边界但没有保留对应运行态时，管理器只能预留槽位，不能标记已保存。未来支持中间捕获的算法需明确提供相应数据及接口扩展。 + +保存或 CPU 导入前先使目标槽位无效，数据操作成功提交后记录长度；异步完成前由上层保持 pending、禁止发布或读取。发生失败需失效本次受影响的槽位，并等已提交操作结束后才可回收。恢复前除长度匹配外，还需满足布局兼容、来源就绪及上层前缀身份匹配。 + +同一槽位数字在 big/small pool 中属于不同分配域。main/draft 的相关 State 可以共享分配的槽位 ID，但每侧各自维护数据与长度记录；管理器校验所有必需侧的数据一致性，缓存调用流程在数据就绪后发布；最终释放时，调用方先通过管理器失效各 pool，再直接向 allocator 归还一次。 + +#### CPU、DP、PD 的状态单位 + +| 路径 | 传输单位与接口职责 | +| --- | --- | +| CPU 写出 | `write_to_cpu` 将指定 pool 的有效快照写到 CPU 组合页，并按布局写入/关联本资源的长度元数据；每个待发布页的参与状态必须完整，长度不得为 None；CPU pool 到 CPU 页同样是独立的格式搬运 | +| CPU 读回 | `read_from_cpu` 核对实际页元数据与预期长度，将快照恢复到已分配 pool 槽位并更新其长度记录；随后需要显式 `restore_pages` 才能恢复请求运行态 | +| DP 复制 | `copy_from_dp` 将源请求已接受状态按 mapping 写入目标请求；传入的长度描述该源状态对应位置。不同 TP 的字段/轴分别映射，目标规范化游标，源保持有效 | +| PD 写出 | `write_to_pd_buffer` 将一个请求的已接受状态写入已选中的状态页区域，按协议关联其前缀长度 | +| PD 读回 | `read_from_pd_buffer` 核对协议中的实际长度和 `expected_token_len` 后恢复目标运行态；不自动生成历史 checkpoint | + +CPU 尾页若需要保存当前运行态，管理器先使用保留槽位调用 `save_pages`，再由该槽位写入 CPU 页。保留槽位的借用和复用仍由管理器协调。使用池槽位与运行态的融合搬运属于实现优化，必须满足相同的长度和就绪契约。 + +位置描述继续复用 `BufferTransferLayout` / `StateCopyMapping` 的纯数据部分：Conv 区域可用 `(layer, channel, window)` 轴，SSM 用 `(layer, head, key_dim, value_dim)` 轴。它们没有历史 KV 的 token 轴。含 KV 的 CPU 页仍有 token 容量，checkpoint 组件不按该容量重复展开；纯 PD 状态页的 `tokens_per_page` 为 None。各资源偏移、TP 范围和元数据写入者由管理器组合，State 直接执行方法。 + +DP/PD 导出的是可继续推理的已接受状态，不传未接受的 MTP 候选区。目标可按自己的 S 初始化工作区；仅在算法、模型身份、字段表示和映射兼容且具体实现支持时，才允许两端 MTP/TP 配置不同。不能因为 KDA/GDN 的 Conv/SSM 形状相同就跨算法恢复。 + +多个来源共同恢复目标的不同区域时，每次调用只修改所声明的区域，不清空其他调用已写入的数据；具体管理器完成全部字段、分片和长度检查，调用方确认恢复就绪后再继续计算。只有当前运行态而没有历史 checkpoint 的 DP/PD 恢复，不能据此发布以前边界的快照。 + +#### 依赖、同步与使用示例 + +资源方法只接受请求槽位、checkpoint 槽位、长度和明确的资源引用，不读取 `InferReq`、ReqManager 映射或全局 infer context。KDAState 不持有 IndexerTailState；模型管理器单独协调 tail 在恢复边界上的重置或迁移。 + +GPU 操作在调用方当前 stream 上提交，返回 None；CPU 同步复制在返回前完成。方法若直接由 CPU 读取 pinned buffer，调用方须先确认此前 GPU 写入已经完成，不能仅靠 GPU stream 排队保证 CPU 可见。其余异步路径遵守事件依赖，释放请求/槽位前保持相关引用有效。正常执行不替换已有 buffer 地址。 + +```python +# ReqManager 已分配请求槽位；调用方已从管理器持有的 allocator 分配状态页槽位。 +state.reset_requests(req_indices) +# backend 经层视图完成 prefill；state_token_lens 为实际 chunk 终点。 +state.save_pages( + req_indices, state_token_lens, + page_kind="big", slot_ids=slot_ids, +) +# 等待所有参与状态保存完成且元数据校验通过后，缓存调用方发布联合页。 + +state.restore_pages( + dst_req_indices, expected_token_lens, + page_kind="big", slot_ids=slot_ids, +) +# target 请求在相关 KV、linear、tail 均就绪后继续执行。 + +# 验证 forward 完成后，将接受步数按请求提交。 +state.commit_accepted_state(req_indices, accepted_steps) +# 后续保存/传输由 State 自行选取已接受的 Conv 窗口和 SSM 槽位。 +``` + +各具体实现必须满足数据完整性、接受位置选择和跨布局复制约束。 + +### 6.8 IndexerTailState：原始尾部窗口的操作接口 + +`IndexerTailState(ABC)` 约束 indexer 尾部运行态的操作，具体实现负责自身窗口、位置规则和搬运 kernel。例如 `KPoolIndexerTailState` 定义 GLM K-pool 的尾部布局。它不继承 `FullAttentionState` 或 `LinearAttentionState`，也不持有 allocator、请求 KV 索引表或独立大小页 pool。 + +#### KPoolIndexerTailState 布局与状态边界 + +设本实例层数为 `L`、请求容量为 `N`（包含 HOLD）、池化步长为 `P`、MTP 额外容量为 `S`、index head 维度为 `D`，原生 ring 布局为 `[L, N, P+S, 2*D]`，末维分别存 raw key 与 gate score。GLM K-pool 示例取 `P=4`、`D=128`，具体实现须声明其 kernel 支持的配置范围。 + +计算先用此前 tail 和本轮 raw 输入完成跨 chunk 池化，再更新 tail。绝对位置 `t` 写入 `t % (P+S)`。完成池化的压缩 key/scale 写入 full-attention packed buffer，由 `FullAttentionState` 管理；tail 保留增量计算所需的原始数据。 + +请求复制和传输保留完整原生 ring。容量 `P+S` 还用于保留拒绝候选或 draft 重访位置时需要的历史；不能仅按 `state_token_len % P` 裁剪，也不能认为长度对齐时所有运行态传输都可以省略 tail。前缀恢复的清零条件单独由 `can_restore_without_tail` 定义。 + +#### 必须实现的接口 + +接口直接执行数据操作。所有方法均为抽象契约，子类可继承兼容的实现；未支持的可选 CPU/DP/PD 路径必须在描述中声明并明确报错，启用相应功能时由初始化检查。CPU 接口为可选能力，需要联合页格式提供 tail 区域；不包含该区域的缓存格式仅支持满足无 tail 恢复条件的前缀。 + +```python +from abc import ABC, abstractmethod + + +class IndexerTailState(ABC): + @classmethod + @abstractmethod + def get_required_bytes(cls, config, storage_spec, *, device) -> int: ... + + @abstractmethod + def __init__(self, config, storage_spec): ... + + @abstractmethod + def get_storage_descriptor(self) -> StorageDescriptor: ... + + @abstractmethod + def get_layer_view(self, layer_id: int): ... + + @abstractmethod + def can_restore_without_tail(self, state_token_len: int) -> bool: ... + + @abstractmethod + def reset_requests(self, req_indices, state_token_lens) -> None: ... + + @abstractmethod + def copy_requests( + self, src_req_indices, dst_req_indices, state_token_lens, + ) -> None: ... + + @abstractmethod + def write_to_cpu( + self, req_indices, state_token_lens, cpu_buffer, cpu_page_indices, + *, layout: BufferTransferLayout, + ) -> None: ... + + @abstractmethod + def read_from_cpu( + self, cpu_buffer, cpu_page_indices, req_indices, expected_token_lens, + *, layout: BufferTransferLayout, + ) -> None: ... + + @abstractmethod + def copy_from_dp( + self, src_state, src_req_indices, dst_req_indices, state_token_lens, + *, mapping: StateCopyMapping, + ) -> None: ... + + @abstractmethod + def write_to_pd_buffer( + self, req_idx, state_token_len, pd_buffer, *, layout: BufferTransferLayout, + ) -> None: ... + + @abstractmethod + def read_from_pd_buffer( + self, pd_buffer, req_idx, expected_token_len, *, layout: BufferTransferLayout, + ) -> None: ... +``` + +除单请求 PD 和单长度能力查询外，本 State 的接口均按批量处理，请求、长度和外部页索引逐项对应。管理器的 CPU 读回方法只选择该前缀最后一页，以长度为 1 的批次调用此 State，不要求为中间页提供目标请求槽位。`state_token_len=L` 表示已经接受且执行过 forward 的前缀 `[0, L)`，不是本轮输入数，也不是 sampled token 数。请求槽位由 ReqManager 分配；长度由执行流程显式传入，不要求 State 再维护一份请求长度表。 + +#### 容量、访问与恢复 + +| 接口 | 语义 | +| --- | --- | +| `get_required_bytes(..., device=...)` | 计算该 device 上全部 tail 字段及必要控制信息的 Tensor 字节数；K-pool 原生 ring 为 `L*N*(P+S)*2*D*element_size`。共享 buffer 不重复计数，外部传输页另由具体管理器计入 | +| 初始化 | 创建或绑定 tail 存储、层映射及配置，覆盖请求槽位域与 HOLD；不分配请求号、KV 索引或 page 槽位 | +| `get_storage_descriptor()` | 描述资源身份、算法/格式、字段与 dtype/shape/stride、层映射、请求容量、TP 分片或副本关系、恢复规则及传输能力；K-pool 额外声明 P、S、ring 寻址方式和原生格式版本 | +| `get_layer_view(layer_id)` | 返回该层计算直接使用的存储引用及必要布局信息。K-pool 可以直接返回 `[N, P+S, 2*D]` Tensor；逻辑层到本地物理层由子类映射，非所属层报错 | +| `can_restore_without_tail(L)` | 判断在“从已恢复完整历史前缀继续计算”的流程中，是否无需保存的 tail 即可恢复。K-pool 对合法非负 L 检查 `L % P == 0`；该查询不检查 KV 就绪性、前缀身份或其他资源 | +| `reset_requests(req_indices, state_token_lens)` | 为全新请求或无 tail 的前缀恢复初始化状态。全新请求传长度 0；其他长度必须逐项满足 `can_restore_without_tail`，不满足时在修改目标前拒绝。K-pool 清零 ring；包含位置控制信息的其他实现还需初始化对应字段 | +| `copy_requests(src, dst, state_token_lens)` | 在本 State 内完整复制指定请求的 tail 恢复状态，保持原始 ring 槽位对应关系及必要控制信息，不改源数据、不复制请求映射或释放索引 | + +`get_layer_view` 供计算 backend 直接使用；池化与写 tail 仍在模型 kernel 中完成，不增加通用的 `append_token` 或 attention forward 包装。执行顺序必须满足“先池化、再更新 tail”。 + +`reset_requests` 是新请求/前缀恢复入口，不是 MTP 回滚入口。即使已接受长度恰好对齐 P，运行中的 draft 重访仍可能需要 ring 内历史，不能用该条件代替 MTP 状态处理。main/draft 分别检查自己的恢复条件；管理器还需确认对应 full-attention、linear 等资源全部恢复到约定位置。 + +批量复制要求目标请求唯一，禁止会覆盖其他源请求的交叉复制；源目标完全相同可以直接跳过。调用者保证源处于已接受边界且数据不再被并发更新；传入任意历史长度不会使当前 ring 自动成为该历史位置的有效状态。 + +#### CPU、DP、PD：直接搬运完整 tail 恢复状态 + +| 接口 | 数据路径及约束 | +| --- | --- | +| `write_to_cpu(...)` | 当前请求 tail → 外部 CPU 页中的 tail 区域。与 linear 的 CPU 方法不同，此处源是请求运行态，没有 `page_kind/slot_ids`；只能捕获当前实际可恢复位置 | +| `read_from_cpu(...)` | CPU 页的 tail 区域 → 目标请求运行态，恢复 ring 和必要控制信息。管理器只将本次前缀最后一页传给此接口，核对保存长度与 `expected_token_lens`；不能把缺失 tail 当成全零快照 | +| `copy_from_dp(...)` | 一个显式来源 State → 目标请求；`mapping` 指定字段、层和分片对应关系。GLM K-pool 在 TP 间为副本，选一个有效来源并恢复每个需要的目标副本；多来源由上层组织 | +| `write_to_pd_buffer(...)` | 指定请求 tail → 已选中的 PD 状态页区域，使用 layout 的资源偏移、字段布局与 TP 写入规则 | +| `read_from_pd_buffer(...)` | 已接收的 PD 状态页区域 → 目标请求 tail；具体管理器先核对资源身份、格式及预期长度，并建立接收完成依赖 | + +这些方法传输 tail 自身的恢复状态，不承担历史 KV、池化 key/scale、linear 状态或 MTP proposal 的迁移。原生 ring 可以包含未接受候选留下的字节，完整复制不表示把这些候选标记为已接受；请求执行流程仍按接受位置和既定重算规则继续。其他实现如有游标、有效范围或辅助字段，必须一并恢复，不能只复制一个看似主要的 Tensor。 + +`BufferTransferLayout` 和 `StateCopyMapping` 继续只描述位置。K-pool 的外部 region 可使用 `("layer", "ring_slot", "channel")` 三个轴；`ring_slot` 保留原生环形槽位语义，不解释为历史 KV 的页内 token 轴。独立状态页的 `tokens_per_page=None`，tail 的字节偏移由具体管理器安排；混合 CPU 页仍按该组合页格式描述。外部页位置由 layout 描述，搬运由 State 方法执行。 + +原生 ring 传输要求算法语义、P、S、ring 寻址与字段格式兼容，层/TP 位置差异通过显式 mapping 处理。不同 S 导致 ring 容量和取模基数变化，不能直接复制或按较小形状截断。未来若增加规范化 tail 格式，需要单独定义可恢复历史范围、最近写入位置及绝对位置到目标 ring 的映射；仅有已接受长度不足以描述所有推测写入历史，该能力不能由当前接口名称推定已支持。 + +长度、前缀身份和格式来自保存/传输控制元数据，不能从原始 tail 字节反推。调用层负责核对元数据并维持 main/draft 的一致恢复点；State 负责布局和数据操作。所有 GPU 操作遵循当前 stream，返回不表示完成；源请求与目标页在读写完成前不能复用，CPU host 消费 pinned 数据前必须确认 GPU 写入已完成。 + +#### 与 linear 页接口和 MTP 提交的区别 + +- 不提供 `save_pages`、`restore_pages`、`has_page`、`invalidate_pages`：tail 没有独立大小页 pool。对齐前缀通过受校验的 reset 恢复；将来缓存非对齐位置时，可把 tail 段放入联合 CPU 页，并记录真实保存长度。 +- 不要求 `commit_accepted_state`：K-pool 没有类似 KDA 的 SSM 选择游标，按绝对位置访问 ring，MTP 拒绝后由执行流程选择后续位置。增加空的提交方法不能替代窗口容量与重访范围的正确性；未来出现确实需要提交控制状态的实现时，再扩展对应契约。 +- tail CPU 快照必须在目标恢复位置仍可捕获时保存，不能在请求推进后给旧 CPU 页补上一份当前 tail。没有 tail 快照的缓存格式仅允许满足无 tail 恢复条件的恢复点。 + +## 7. 资源初始化与调用约定 + +### 7.1 初始化与绑定 + +1. 模型选择具体 `AttResourceManager` 实现，确定内部 State 组合、请求状态槽位容量、共享分配域和存储布局。 +2. 具体管理器创建或绑定所需 allocator，初始化 State、可选大小页 pool 和保留区域；同一共享分配域只创建一个 allocator 实例。 +3. 模型初始化代码将具体 State 或层视图绑定到 backend。ReqManager 持有请求槽位及 KV 映射，使用管理器提供的 HOLD 物理索引初始化保留行。 +4. 为启用的大小页域确定完整参与 State 集合，并配置 CPU/DP/PD 数据格式、资源位置与分片映射，检查相关 State 的数据操作能力。独立传输进程使用具体实现导入并绑定所需存储。 + +资源容量确定后，buffer 地址及已绑定视图保持稳定。资源申请和使用遵守共享 allocator 的索引规则；初始化过程不要求额外的公共请求生命周期接口。 + +### 7.2 缓存保存与恢复 + +没有任何可恢复缓存、从长度 0 开始时，调用方在获得请求槽位后调用 `reset_state(req_indices)`，统一初始化全部参与运行态，再开始计算。命中非零前缀时使用对应恢复路径；完整恢复前不得读取目标状态。 + +调用方直接使用指定大小页域的 allocator 取得槽位，在状态仍可捕获时调用管理器的 `save_pages`。具体管理器将同一组槽位传给全部参与 State 保存，并记录长度；调用方确认数据就绪后，可以把 `(page_kind, slot_id, state_token_len)` 关联到 radix 节点，也可以用于其他缓存操作。 + +radix 匹配得到节点引用后，调用方使用 `has_page` 检查状态页记录，并调用 `restore_pages` 恢复全部参与 State。KV 索引由调用方交给 ReqManager 绑定；没有 tail 快照时，调用方使用管理器的 `reset_tail_state(req_indices, expected_token_lens)` 校验并初始化 tail，保留已恢复的 linear 状态。状态页校验只覆盖对应 pool，调用方还需确认历史 KV、tail 和数据就绪性。 + +使用 K-pool 的 tail 在完整池化边界上通过 `reset_tail_state` 重置；非对齐恢复需要原始 tail 数据。CPU 读回只从该前缀最后一页恢复 tail；格式没有 tail 段时，由管理器按最后一页长度调用 `reset_tail_state`。DP/PD 运行态复制按第 6.8 节保留完整兼容的 ring，不能把前缀恢复的清零规则用于 MTP 回退。 + +MTP 接受位置由模型执行流程选择，再保存或传输已接受状态。KDA/GDN 只保存实际捕获的状态；跨越但未保存的中间位置不可作为有效页，也不能发布为 CPU 缓存页；执行流程需在待缓存页边界安排真实状态捕获。管理器的数据操作不承担一次 forward 的提交或请求长度更新。 + +### 7.3 同步、引用与回收 + +请求流程负责请求暂停、结束及请求槽位回收；缓存功能流程负责节点引用和淘汰时机。归还资源索引前，调用方应停止相关使用、满足异步完成条件并解除相应引用。ReqManager 清理请求映射；对于状态页，调用方先通过管理器统一失效，再直接向对应 allocator 归还一次。 + +有缓存、请求或传输引用的数据必须保持有效。释放一份缓存数据不等于结束某个请求；请求结束也不表示其已共享缓存应被释放。管理器对明确传入的状态页执行聚合失效，调用方负责将具备回收条件的索引归还 allocator。 + +保存、导入或复制失败时，不发布本次不完整数据；调用方在相关数据统一失效且操作完成后，直接回收本次不再使用的分配，已完成的共享缓存保持有效。CPU/DP/PD 操作遵守 stream/event 依赖，调用方在继续计算、复用 buffer 或回收索引前确认完成。 + +## 8. CPU、DP 与 PD 传输设计 + +### 8.1 具体 State 实现传输，模型管理器组织组合 + +外部通过 6.2 节的五个管理器抽象方法执行 CPU/DP/PD 操作。具体管理器将布局及显式索引分发给各 State；各 State 继续使用 6.1、6.7、6.8 节的直接接口。 + +- full-attention 资源直接提供 CPU、DP、PD 读写方法，覆盖 KV、scale 或 packed 数据；请求流程从 `ReqManager` 获取请求映射，向 State 方法显式传入源、目标索引。 +- `KDAState`、`GDNState` 实现 `LinearAttentionState` 的直接传输方法：CPU 路径读写 checkpoint，DP/PD 路径迁移已接受请求状态；管理器组织参与对象及位置映射。 +- `IndexerTailState` 提供请求复制与 CPU/DP/PD 直接读写方法，详见 6.8 节；CPU 能力按联合页格式启用。GLM K-pool indexer 在 TP ranks 间为副本，其他实现自行描述其布局。 +- 各具体 State 实现自身 GPU 存储句柄与布局元数据的导出、导入和绑定；模型专属 `AttResourceManager` 按角色及资源身份组织所需 State，并实现组合传输流程。传输进程使用这些具体实现，CPU pinned checkpoint pool 和 allocator 留在推理进程。main/draft 共享索引时，仍需逐项检查资源身份和格式兼容性。 + +### 8.2 PD + +PD 任务分为 `kv` 和 `att_state` 两类；涉及 main/draft 时按模型角色组合所需资源： + +| 任务 | 组成 | 处理方式 | +| --- | --- | --- | +| `kv` | full-attention 历史 KV，包含布局所需的 packed indexer 数据 | 按 token 页复制 | +| `att_state` | linear Conv/SSM + 可选 indexer tail | 模型专属 `AttResourceManager` 按请求恢复位置调用对应 State 方法 | + +P 端在 prefill 完成时发送状态页。linear state 按全局 head 布局打包，indexer tail 从来源 TP 组取一份；D 端按目标 TP 布局恢复 linear state，并将 tail 写入目标 TP 各 rank 的本地请求槽位。 + +模型专属 `AttResourceManager` 实现状态页的布局组织和容量校验，调用 linear State 与 indexer tail State 的直接传输方法。各具体 State 负责自身区域的读写和格式校验，PD 调度负责任务队列、网络传输与完成通知。收到所有 KV/state 数据并完成 GPU 写入后，请求才能继续 decode。 + +双方需要明确兼容的状态页格式，包括模型角色、资源实现、dtype、层数、全局 head 布局和 tail 布局。不同实现的 main/draft 组合需要明确各资源段布局、标识和版本，不能仅凭共享 allocator 推定传输兼容。linear 状态按 6.7 节导出已接受前缀的规范化状态，不传 MTP 候选区；不同 MTP 配置能否互通还须校验状态格式和实现能力。indexer tail 按 6.8 节复制完整、兼容的原生 ring；不同 MTP 配置的 ring 转换需要另行定义格式和可恢复历史范围，暂不支持直接互通。 + +### 8.3 CPU/Disk 与 DP fetch + +模型专属 `AttResourceManager` 实现 CPU 联合页的布局与读写组织，调用所需 main/draft State 的 CPU 方法，组合 full-attention KV、linear 页及各参与 pool 的 `state_token_lens`。CPU 缓存调度提供外部页和任务生命周期；每个已发布 CPU 页都具备该模型组合所需的有效 linear 状态。一次读回处理同一前缀的页面，仅从最后一页恢复 indexer tail；没有 tail 段时通过 `reset_tail_state` 按最终长度重置。非对齐恢复需要最后一页包含真实 tail 快照，并由其具体 State 实现读回。 + +CPU load/offload 的尾页暂存槽位由管理器保留和协调，参与的大页 pool 分别提供存储。各参与方的暂存数据消费完成后，后续请求才能复用该槽位。磁盘存储使用 CPU 页格式；main/draft 异构组合需定义各自资源段的布局、格式版本和兼容规则。 + +DP fetch 同时恢复需要的历史 KV 和请求运行态。只有当前运行态、缺少历史 checkpoint 的请求，不能据此发布未经保存的历史状态。传输能力由资源布局显式声明，不支持的组合在初始化阶段拒绝。 + +## 9. 待细化的设计问题 + +| 问题 | 设计约定 | 待确定内容 | +| --- | --- | --- | +| main/draft 管理器装配 | 模型实现自己的管理器；同一分配域共享唯一 allocator 实例 | 使用组合管理器还是多个模型管理器实例，allocator 的唯一创建/销毁者；调用方覆盖全部参与 pool 的失效顺序及单次归还责任 | +| AttResourceManager 聚合接口 | 按 6.2 节提供 full-attention token 复制、空前缀清零、按长度重置 tail、大小页聚合操作及完整 CPU/DP/PD 接口；radix 是调用方 | 各模型参与集合、跨实例共享域装配、分阶段捕获及索引/元数据的具体承载类型 | +| 内部 attention 资源的暴露 | 具体模型在初始化阶段绑定 State 或层视图，不要求统一模型视图获取接口 | 各模型/backend 的具体绑定方式与按层访问结构 | +| FullAttentionState 接口与继承 | 必需语义：容量计算、初始化、完整存储描述、层视图、KV 写入、完整 token 复制；CPU/DP/PD 由子类按支持路径直接实现读写方法 | StorageDescriptor 的最小字段、BufferTransferLayout / StateCopyMapping 的具体数据类型，以及各布局的复制/传输约束 | +| LinearAttentionState 接口与继承 | KDAState、GDNState 实现统一抽象操作；已接受运行态、checkpoint 和传输分别定义，allocator 仍归管理器 | 按 6.7 节细化各子类的描述字段和 kernel，核对 MTP 状态选择、跨 TP/CPU 格式兼容性 | +| IndexerTailState 接口与继承 | 按 6.8 节实现层访问、受校验的重置、请求复制及直接传输；原生 ring 完整迁移，无独立页 pool | 各子类的恢复规则、带 tail 段的 CPU 格式及跨 ring 布局转换能力 | +| allocator 能力描述 | 公共入口按需初始化，不需要的 allocator 为 None | 恢复边界和传输能力使用哪些最小描述，避免重复配置导致不一致 | +| checkpoint 引用形式 | 明确大小页类别、共享槽位 ID、预期恢复长度及参与模型资源；各侧 pool 保留 `state_token_lens` | 引用对象或参数组合的具体类型，以及参与资源集合的表达 | +| 配置类型拆分 | main/draft 分别产生实现和布局参数，共享 allocator 接收统一容量与索引配置 | 各类 State 的配置字段、跨资源共享字段与校验责任 | +| PD tail 的协议布局 | 完整复制兼容原生 ring，组件描述字段，具体管理器安排位置与 TP 副本读写 | 不同 MTP 配置下的可恢复历史范围、位置元数据和格式转换规则 | +| CPU 缓存页的组合方式 | 每页包含全部有效 linear 状态；管理器组织读写，并只从最后一页恢复 tail | 不同 main/draft 实现的资源段布局、非对齐 tail 快照及版本兼容 | + +以上待细化事项须遵守本文的资源归属、共享分配域和恢复语义约定。 diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 2cde561013..270812836a 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -139,7 +139,7 @@ def prefill_att( layer_num = linear_att_dict["layer_num"] backend: LinearAttBackend = self.backend - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.infer_state.linear_state_manager.get_layer_buffers(layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: @@ -227,7 +227,7 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int): ) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index, + req_to_mtp_state_index=self.infer_state.linear_state_manager.mtp_state_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) @@ -249,7 +249,7 @@ def _init_fixed_mtp_decode_state(self, draft_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + self.b_num_accepted_tokens = self.infer_state.linear_state_manager.mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): @@ -281,7 +281,7 @@ def decode_att( backend: LinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.infer_state.linear_state_manager.get_layer_buffers(layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index 3d0f0cf5f4..4385c6efc3 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -83,7 +83,7 @@ def prefill_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.infer_state.linear_state_manager.get_layer_buffers(layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] mixed_qkv = causal_conv1d_fn( mixed_qkv.transpose(0, 1), @@ -151,7 +151,7 @@ def _init_dynamic_mtp_decode_state(self, mtp_size: int): ) = build_dynamic_mtp_linear_att_state_params( b_req_idx=self.infer_state.b_req_idx, b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=self.infer_state.req_manager.req_to_mtp_state_index, + req_to_mtp_state_index=self.infer_state.linear_state_manager.mtp_state_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) @@ -173,7 +173,7 @@ def _init_fixed_mtp_decode_state(self, mtp_step: int): device=self.infer_state.b_req_idx.device, ) self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = self.infer_state.req_manager.req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 + self.b_num_accepted_tokens = self.infer_state.linear_state_manager.mtp_state_index[self.b_conv_buffer_idx] + 1 self._init_mtp_ssm_buffer_idx(mtp_size) def _init_mtp_ssm_buffer_idx(self, mtp_size: int): @@ -204,7 +204,7 @@ def decode_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - conv_states, ssm_states = self.infer_state.req_manager.get_mamba_cache(layer_num) + conv_states, ssm_states = self.infer_state.linear_state_manager.get_layer_buffers(layer_num) conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) conv_update = causal_conv1d_update if backend.mtp_step > 0: diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 34a1c9c71c..9f6ada6334 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -16,6 +16,7 @@ from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class from lightllm.common.req_manager import ReqManager +from lightllm.common.cache_coordinator import CacheCoordinator 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, @@ -114,12 +115,14 @@ def __init__(self, kvargs): self._init_weights() with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): self._init_req_manager() + if getattr(self, "is_mtp_draft_model", False): + self.cache_coordinator = self.main_model.cache_coordinator + else: + self._init_cache_coordinator() self._init_mem_manager() - # Qwen3.5 等 linear attention 模型会在 req_manager 中保存大量运行时 state。先初始化 - # req_manager、再初始化 mem_manager，可以让 KV cache 显存评估包含这些 state 的实际占用； - # 因此 req_manager 创建时暂不传入 mem_manager，需要在两者初始化完成后再进行绑定。 self.req_manager.bind_mem_manager(self.mem_manager) + self.cache_coordinator.bind_mem_manager(self.mem_manager) self._check_mem_size() self._init_infer_layer() self._init_some_value() @@ -196,9 +199,13 @@ def load_weights(self, weight_dict: dict): [weight.verify_load() for weight in self.trans_layers_weight] return + @classmethod + def get_mem_manager_class(cls): + return select_mem_manager_class() + def _init_mem_manager(self): assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 - self.mem_manager: MemoryManager = select_mem_manager_class()( + self.mem_manager: MemoryManager = self.get_mem_manager_class()( self.max_total_token_num, dtype=self.data_type, head_num=self.config["num_attention_heads"] // self.tp_world_size_, @@ -229,6 +236,9 @@ def _check_mem_size(self): return + def _init_cache_coordinator(self): + self.cache_coordinator = CacheCoordinator() + def _init_req_manager(self): create_max_seq_len = 0 @@ -376,6 +386,9 @@ 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.state_manager = self.cache_coordinator.state_manager + infer_state.linear_state_manager = self.cache_coordinator.linear_state_manager + infer_state.indexer_state_manager = self.cache_coordinator.indexer_state_manager infer_state.mem_index = self._select_mem_indexes(model_input) infer_state.microbatch_index = microbatch_index @@ -631,7 +644,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" diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index bf25149450..c9f5f7a53c 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -49,6 +49,8 @@ def __init__(self): self.mem_manager: MemoryManager = None self.req_manager: ReqManager = None + self.linear_state_manager = None + self.indexer_state_manager = None self.mem_index: torch.Tensor = None diff --git a/lightllm/models/deepseek3_2/triton_kernel/destindex_copy_kv_flashmla_fp8.py b/lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py similarity index 100% rename from lightllm/models/deepseek3_2/triton_kernel/destindex_copy_kv_flashmla_fp8.py rename to lightllm/common/basemodel/triton_kernel/kv_copy/destindex_copy_kv_flashmla_fp8.py diff --git a/lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py b/lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py similarity index 100% rename from lightllm/models/deepseek3_2/triton_kernel/prefill_compact_kv_flashmla_fp8.py rename to lightllm/common/basemodel/triton_kernel/kv_copy/prefill_compact_kv_flashmla_fp8.py diff --git a/lightllm/common/cache_coordinator.py b/lightllm/common/cache_coordinator.py new file mode 100644 index 0000000000..45e5df2990 --- /dev/null +++ b/lightllm/common/cache_coordinator.py @@ -0,0 +1,479 @@ +import torch +import triton +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.multimodal_chunk_utils import clamp_mm_cache_len + + +class CacheCoordinator: + has_state_cache = False + state_manager = None + linear_state_manager = None + indexer_state_manager = None + + def __init__(self): + self.mem_manager = None + + def bind_mem_manager(self, mem_manager): + if self.mem_manager is not None and self.mem_manager is not mem_manager: + raise ValueError("CacheCoordinator is already bound to a different memory manager") + self.mem_manager = mem_manager + + def validate_radix_cache(self, radix_cache): + if radix_cache is not None and radix_cache.mem_manager is not self.mem_manager: + raise ValueError("Radix cache and coordinator must share the memory manager") + + def create_state_transfer(self): + return None + + def restore_prefix(self, req, req_manager, radix_cache): + enabled = not req.sampling_param.disable_prompt_cache and radix_cache is not None + if enabled and req.get_cur_total_len() > 1 and req.cur_kv_len == 0: + tokens = req.shm_req.shm_prompt_ids.arr[: req.get_cur_total_len()] + key = torch.tensor(tokens, dtype=torch.int64, device="cpu")[:-1] + node, _, indexes = radix_cache.match_prefix(key, update_refs=True, mm_input_spans=req.mm_input_spans) + if node is not None: + length = node.node_prefix_total_len + req_manager.req_to_token_indexs[req.req_idx, :length] = indexes + req.shared_kv_node = node + req.cur_kv_len = req.hold_kv_len = int(length) + assert req.hold_kv_len % req.args.page_size == 0 + req.shm_req.prompt_cache_len = req.cur_kv_len + req.shm_req.shm_cur_kv_len = req.cur_kv_len + + def save_state_cache(self, b_req_idx, reqs, radix_cache, args): + """Snapshot after forward, before advancing req.cur_kv_len from the chunk start.""" + return + + def get_cacheable_prefix_len(self, req, args): + return req.cur_kv_len + + def release_request_cache(self, req, req_manager, radix_cache, args, publish_cache) -> list[torch.Tensor]: + """Publish reusable cache and release unpublished checkpoints and held radix refs. + + Returns borrowed request-index slices for the caller to batch-free. The caller + resets request lengths and must consume these slices before reusing the request + row. Call only after pending CPU offload reads have completed. Published KV and + checkpoints belong to radix; returned KV indexes still need allocator release. + """ + indexes = req_manager.req_to_token_indexs[req.req_idx] + shared_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len + if publish_cache and radix_cache is not None: + cache_len = req.cur_kv_len // args.page_size * args.page_size + key = torch.tensor(req.get_input_token_ids()[:cache_len], dtype=torch.int64, device="cpu") + prefix_len, _ = radix_cache.insert(key, indexes[:cache_len].detach().cpu()) + assert shared_len <= prefix_len + free_indexes = [indexes[cache_len : req.hold_kv_len], indexes[shared_len:prefix_len]] + else: + free_indexes = [indexes[shared_len : req.hold_kv_len]] + if req.shared_kv_node is not None: + radix_cache.dec_node_ref_counter(req.shared_kv_node) + req.shared_kv_node = None + return free_indexes + + def restore_peer_state(self, req_idx, restored_len, peer_state, peer_req_idx): + return + + def get_cpu_cache_match_page_num(self, page_indexes, page_lens, cpu_cache_client): + return len(page_indexes) + + def load_cpu_cache_to_gpu(self, mem_indexes, page_indexes, cpu_cache_client, req, page_indexes_cpu=None) -> None: + """Enqueue KV and any live request-state restoration on the current stream.""" + self.mem_manager.operator.load_cpu_cache_to_gpu(mem_indexes, page_indexes, cpu_cache_client, req) + + def offload_gpu_kv_to_cpu_cache( + self, + mem_indexes, + page_indexes, + page_readies, + cpu_cache_client, + req, + radix_cache=None, + page_indexes_cpu=None, + page_readies_cpu=None, + ): + return self.mem_manager.operator.offload_gpu_kv_to_cpu_cache( + mem_indexes, page_indexes, page_readies, cpu_cache_client, req + ) + + +class HybridCacheCoordinator(CacheCoordinator): + """Coordinate KV prefix reuse and hybrid attention state checkpoints.""" + + def __init__(self, linear_state_manager=None, indexer_state_manager=None, *, state_manager=None): + self.state_manager = state_manager if state_manager is not None else linear_state_manager + self.linear_state_manager = linear_state_manager + self.indexer_state_manager = indexer_state_manager + self.mem_manager = None + self.small_page_buffers = None + self.big_page_buffers = None + + @property + def has_state_cache(self): + return self.state_manager is not None + + def bind_mem_manager(self, mem_manager): + if self.mem_manager is mem_manager: + return + if self.mem_manager is not None: + raise ValueError("HybridCacheCoordinator is already bound to a different memory manager") + if self.has_state_cache: + args = get_env_start_args() + big_page_tokens = args.linear_att_page_block_num * args.linear_att_hash_page_size + # The last two slots are reserved for CPU-cache tail staging, never allocated to requests. + self.big_page_buffers = self.state_manager.create_checkpoint_pool( + triton.cdiv(mem_manager.size, big_page_tokens) + 2, + keep_num=2, + ) + self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 + self.cpu_cache_transfer = self.state_manager.create_cpu_transfer(mem_manager, self.big_page_buffers) + self.mem_manager = mem_manager + + def create_state_transfer(self): + if not self.has_state_cache: + return None + return self.state_manager.create_state_transfer(self.indexer_state_manager) + + def validate_radix_cache(self, radix_cache): + super().validate_radix_cache(radix_cache) + if radix_cache is not None and self.has_state_cache: + if ( + radix_cache.big_page_buffers is not self.big_page_buffers + or radix_cache.small_page_buffers is not self.small_page_buffers + ): + raise ValueError("Radix cache and coordinator must share checkpoint pools") + + def create_small_page_cache_manager(self, size): + if self.small_page_buffers is not None: + if self.small_page_buffers.size != size: + raise ValueError("Cannot replace an initialized checkpoint pool") + return self.small_page_buffers + self.small_page_buffers = self.state_manager.create_checkpoint_pool(size) + return self.small_page_buffers + + def reset(self, req_idx): + if self.state_manager is not None: + self.state_manager.reset(req_idx) + if self.indexer_state_manager is not None: + self.indexer_state_manager.reset(req_idx) + + def restore_state(self, req_idx, state_cache_manager, buffer_idx, restored_len): + assert state_cache_manager.has_checkpoint(buffer_idx, restored_len) + self.state_manager.restore_checkpoint(req_idx, restored_len, state_cache_manager, buffer_idx) + if self.indexer_state_manager is not None: + self.indexer_state_manager.reset(req_idx) + + def restore_peer_state(self, req_idx, restored_len, peer_state, peer_req_idx): + self.state_manager.restore_peer_state(req_idx, restored_len, peer_state, peer_req_idx) + if self.indexer_state_manager is not None: + self.indexer_state_manager.state.buffer[:, req_idx].copy_( + peer_state.req_to_indexer_tail.buffer[:, peer_req_idx], non_blocking=True + ) + + def save_state(self, req_idx, buffer_idx, state_cache_manager, checkpoint_len): + self.state_manager.save_state(req_idx, buffer_idx, state_cache_manager, checkpoint_len) + state_cache_manager.checkpoint_lens[buffer_idx] = checkpoint_len + + def save_big_page_states(self, b_req_idx, req_indexes, buffer_indexes, checkpoint_lens): + self.state_manager.save_checkpoints( + b_req_idx, req_indexes, buffer_indexes, self.big_page_buffers, checkpoint_lens + ) + for slot, length in zip(buffer_indexes, checkpoint_lens): + self.big_page_buffers.checkpoint_lens[slot] = length + + def _restore_shared_prefix( + self, req_idx, token_indexes, cached_indexes, restored_len, checkpoint_pool, checkpoint_slot + ): + token_indexes[:restored_len] = cached_indexes[:restored_len] + self.restore_state(req_idx, checkpoint_pool, checkpoint_slot, restored_len) + + def restore_prefix(self, req, req_manager, radix_cache): + enable_prompt_cache = (not req.sampling_param.disable_prompt_cache) and radix_cache is not None + block_hashs = req.shm_req.hybrid_token_hash_list.get_all() + hash_page_size = req.args.linear_att_hash_page_size + match_tokens = max(0, min(len(block_hashs) * hash_page_size, req.get_cur_total_len() - 1)) + match_tokens = (match_tokens // hash_page_size) * hash_page_size + block_hashs = block_hashs[: match_tokens // hash_page_size] + assert len(block_hashs) == req.shm_req.hybrid_token_hash_list.size + big_page_token_num = hash_page_size * req.args.linear_att_page_block_num + big_page_is_disable = big_page_token_num > req.args.max_req_total_len + restored_kv_len = req.cur_kv_len + retained_node = req.shared_kv_node + if enable_prompt_cache and match_tokens > 1 and block_hashs and req.cur_kv_len == 0: + input_token_ids = req.shm_req.shm_prompt_ids.arr[: req.get_cur_total_len()] + key = torch.tensor(input_token_ids[:match_tokens], dtype=torch.int64, device="cpu") + share_node, _, value_tensor = radix_cache.match_prefix( + key, block_hashs=block_hashs, update_refs=True, mm_input_spans=req.mm_input_spans + ) + if share_node is not None: + assert req.tail_small_page_buffer_id is None + token_indexes = req_manager.req_to_token_indexs[req.req_idx] + if not share_node.is_big_page_node() and not big_page_is_disable: + matched_len = share_node.node_prefix_total_len + big_page_len = (matched_len // big_page_token_num) * big_page_token_num + need_tokens = matched_len - big_page_len + available_tokens = ( + self.mem_manager.allocator.can_use_mem_size + + radix_cache.get_tree_total_tokens_num() + - radix_cache.get_refed_tokens_num() + ) + if available_tokens > need_tokens: + assert need_tokens % req.args.page_size == 0 + radix_cache.free_radix_cache_to_get_enough_token(need_token_num=need_tokens) + tail_mems = self.mem_manager.alloc(need_size=need_tokens) + token_indexes[:big_page_len] = value_tensor[:big_page_len] + token_indexes[big_page_len:matched_len] = tail_mems + self.mem_manager.operator.copy_mem_to_mem(value_tensor[big_page_len:matched_len], tail_mems) + req.prompt_selected_logprobs.copy_capture_slots_if_needed( + source_indexes=value_tensor[big_page_len:matched_len], destination_indexes=tail_mems + ) + self.restore_state( + req.req_idx, self.small_page_buffers, share_node.small_page_buffer_idx, matched_len + ) + restored_kv_len = matched_len + retained_node = radix_cache.deref_to_first_big_page_node(node=share_node, require_state=False) + share_node = None + else: + share_node = radix_cache.deref_to_first_big_page_node( + node=share_node, mm_input_spans=req.mm_input_spans + ) + if share_node is not None: + restored_kv_len = share_node.node_prefix_total_len + retained_node = share_node + if share_node.is_big_page_node(): + pool, slot = self.big_page_buffers, share_node.big_page_buffer_idx + else: + pool, slot = self.small_page_buffers, share_node.small_page_buffer_idx + self._restore_shared_prefix(req.req_idx, token_indexes, value_tensor, restored_kv_len, pool, slot) + if restored_kv_len > 0: + req.shm_req.prompt_cache_len = int(restored_kv_len) + + req.shared_kv_node = retained_node + req.cur_kv_len = int(restored_kv_len) + req.shm_req.shm_cur_kv_len = req.cur_kv_len + req.hold_kv_len = req.cur_kv_len + assert req.hold_kv_len % req.args.page_size == 0 + if req.cur_kv_len == 0: + self.reset(req.req_idx) + + def save_state_cache(self, b_req_idx, reqs, radix_cache, args): + """Snapshot after forward, before advancing req.cur_kv_len from the chunk start.""" + if radix_cache is None: + return + # Request-state snapshot at a big-page boundary. + big_page_token_num = args.linear_att_hash_page_size * args.linear_att_page_block_num + big_page_buffer_ids, checkpoint_lens, req_rows = [], [], [] + for row, req in enumerate(reqs): + chunk_end = req.get_chuncked_input_token_len() + first_boundary = (req.cur_kv_len // big_page_token_num + 1) * big_page_token_num + for length in range(first_boundary, min(chunk_end, req.hybrid_cache_len) + 1, big_page_token_num): + slot = self.big_page_buffers.alloc_one_state_cache() + assert slot is not None and length not in req.hybrid_len_to_big_page_id + req.hybrid_len_to_big_page_id[length] = slot + if self.state_manager.capture_intermediate_checkpoints or length == chunk_end: + big_page_buffer_ids.append(slot) + checkpoint_lens.append(length) + req_rows.append(row) + if big_page_buffer_ids: + selected_rows = torch.tensor(req_rows, dtype=torch.int64, device=b_req_idx.device) + self.save_big_page_states( + b_req_idx=b_req_idx[selected_rows], + req_indexes=[reqs[row].req_idx for row in req_rows], + buffer_indexes=big_page_buffer_ids, + checkpoint_lens=checkpoint_lens, + ) + + assert ( + not args.disable_chunked_prefill or self.state_manager.capture_intermediate_checkpoints + ), "chunked prefill must be enabled when only chunk-end state can be captured" + + # Request-state snapshot at the final small-page boundary. + for req in reqs: + # 判断本次prefill 完以后 kv 的长度是否到达 hybrid checkpoint 的存储边界。 + if req.cur_kv_len < req.hybrid_cache_len <= req.get_chuncked_input_token_len() and ( + self.state_manager.capture_intermediate_checkpoints + or req.get_chuncked_input_token_len() == req.hybrid_cache_len + ): + assert req.tail_small_page_buffer_id is None + if req.hybrid_cache_len % big_page_token_num != 0: + radix_cache.free_one_small_page_buffer() + req.tail_small_page_buffer_id = self.small_page_buffers.alloc_one_state_cache() + if req.tail_small_page_buffer_id is not None: + dst_buffer_idx = req.tail_small_page_buffer_id + self.save_state( + req_idx=req.req_idx, + buffer_idx=dst_buffer_idx, + state_cache_manager=self.small_page_buffers, + checkpoint_len=req.hybrid_cache_len, + ) + return + + def get_cacheable_prefix_len(self, req, args): + """Peer runtime restores permit continuation, but do not supply historical checkpoints.""" + page_size = args.linear_att_hash_page_size * args.linear_att_page_block_num + shared_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len + end = min(req.cur_kv_len, req.hybrid_cache_len) + cacheable_len = 0 + node = req.shared_kv_node + while node is not None and node.node_prefix_total_len: + length = node.node_prefix_total_len + pool, slot = ( + (self.big_page_buffers, node.big_page_buffer_idx) + if length % page_size == 0 + else (self.small_page_buffers, node.small_page_buffer_idx) + ) + if pool.has_checkpoint(slot, length): + cacheable_len = length + break + node = node.parent + for boundary in range((shared_len // page_size + 1) * page_size, end + 1, page_size): + if boundary not in req.hybrid_len_to_big_page_id: + return cacheable_len + slot = req.hybrid_len_to_big_page_id[boundary] + if self.big_page_buffers.has_checkpoint(slot, boundary): + cacheable_len = boundary + if self.small_page_buffers.has_checkpoint(req.tail_small_page_buffer_id, end): + cacheable_len = end + return cacheable_len + + def release_request_cache(self, req, req_manager, radix_cache, args, publish_cache) -> list[torch.Tensor]: + """Publish reusable cache and release unpublished checkpoints and held radix refs. + + Returns borrowed request-index slices for the caller to batch-free. The caller + resets request lengths and must consume these slices before reusing the request + row. Call only after pending CPU offload reads have completed. Published KV and + checkpoints belong to radix; returned KV indexes still need allocator release. + """ + token_indexes = req_manager.req_to_token_indexs[req.req_idx] + shared_len = 0 if req.shared_kv_node is None else req.shared_kv_node.node_prefix_total_len + publish_len = None + if publish_cache and radix_cache is not None: + hash_page_size = args.linear_att_hash_page_size + big_page_tokens = hash_page_size * args.linear_att_page_block_num + last_big_page = req.hybrid_cache_len // big_page_tokens * big_page_tokens + assert req.hybrid_cache_len >= shared_len + if req.tail_small_page_buffer_id is not None: + assert req.hybrid_cache_len <= req.cur_kv_len + assert req.hybrid_cache_len % big_page_tokens != 0 + publish_len = req.hybrid_cache_len + elif shared_len < last_big_page <= req.cur_kv_len: + publish_len = last_big_page + else: + assert shared_len <= req.cur_kv_len + + if publish_len is not None: + publish_len = min(publish_len, self.get_cacheable_prefix_len(req, args)) + if publish_len <= shared_len: + publish_len = None + + if publish_len is not None: + discarded = [length for length in req.hybrid_len_to_big_page_id if length > publish_len] + if discarded: + self.big_page_buffers.free_state_cache( + [req.hybrid_len_to_big_page_id.pop(length) for length in discarded] + ) + if req.tail_small_page_buffer_id is not None and publish_len < req.hybrid_cache_len: + self.small_page_buffers.free_state_cache([req.tail_small_page_buffer_id]) + req.tail_small_page_buffer_id = None + page_num = publish_len // hash_page_size + block_hashs = req.shm_req.hybrid_token_hash_list.get_all()[:page_num] + state_idxs = [None] * page_num + if req.tail_small_page_buffer_id is not None: + state_idxs[-1] = req.tail_small_page_buffer_id + key = torch.tensor(req.get_input_token_ids()[:publish_len], dtype=torch.int64, device="cpu") + value = token_indexes[:publish_len].detach().cpu() + prefix_len, _ = 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, + ) + assert not req.hybrid_len_to_big_page_id + req.tail_small_page_buffer_id = None + assert shared_len <= prefix_len + free_indexes = [token_indexes[publish_len : req.hold_kv_len], token_indexes[shared_len:prefix_len]] + else: + free_indexes = [token_indexes[shared_len : req.hold_kv_len]] + if req.tail_small_page_buffer_id is not None: + self.small_page_buffers.free_state_cache([req.tail_small_page_buffer_id]) + req.tail_small_page_buffer_id = None + if req.hybrid_len_to_big_page_id: + self.big_page_buffers.free_state_cache(list(req.hybrid_len_to_big_page_id.values())) + req.hybrid_len_to_big_page_id.clear() + + if req.shared_kv_node is not None: + radix_cache.dec_node_ref_counter(req.shared_kv_node) + req.shared_kv_node = None + return free_indexes + + def get_cpu_cache_match_page_num(self, page_indexes, page_lens, cpu_cache_client): + lengths = self.cpu_cache_transfer.get_checkpoint_lens(page_indexes, cpu_cache_client) + for index in range(len(lengths) - 1, -1, -1): + if lengths[index] == page_lens[index]: + return index + 1 + return 0 + + def load_cpu_cache_to_gpu(self, mem_indexes, page_indexes, cpu_cache_client, req, page_indexes_cpu=None) -> None: + """Enqueue KV and any live request-state restoration on the current stream.""" + args = get_env_start_args() + page_size = args.cpu_cache_token_page_size + assert page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num + assert len(mem_indexes) % args.linear_att_hash_page_size == 0 + assert triton.cdiv(len(mem_indexes), page_size) == len(page_indexes) + assert clamp_mm_cache_len(req.cur_kv_len, req.mm_input_spans) == req.cur_kv_len + if page_indexes_cpu is None: + page_indexes_cpu = page_indexes.cpu().tolist() + checkpoint_lens = self.cpu_cache_transfer.get_checkpoint_lens(page_indexes_cpu, cpu_cache_client) + assert checkpoint_lens[-1] == req.cur_kv_len + page_count = len(mem_indexes) // page_size + end = req.cur_kv_len // page_size * page_size + slots = self.big_page_buffers.alloc_state_cache(page_count) + assert slots is not None, "Insufficient checkpoint slots for CPU cache restoration" + first_end = end - (page_count - 1) * page_size + for index, slot in enumerate(slots): + length = first_end + index * page_size + req.hybrid_len_to_big_page_id[length] = slot + if checkpoint_lens[index] == length: + self.big_page_buffers.checkpoint_lens[slot] = length + if len(mem_indexes) % page_size: + slots.append(self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) + self.big_page_buffers.checkpoint_lens[slots[-1]] = req.cur_kv_len + self.cpu_cache_transfer.load(mem_indexes, page_indexes, cpu_cache_client, slots) + self.restore_state(req.req_idx, self.big_page_buffers, slots[-1], req.cur_kv_len) + + def offload_gpu_kv_to_cpu_cache( + self, + mem_indexes, + page_indexes, + page_readies, + cpu_cache_client, + req, + radix_cache=None, + page_indexes_cpu=None, + page_readies_cpu=None, + ): + page_size = get_env_start_args().cpu_cache_token_page_size + end = len(mem_indexes) // page_size * page_size + slots = [] if radix_cache is None else radix_cache.get_big_page_ids_by_node(req.shared_kv_node) + slots = slots[: end // page_size] + for length in range((len(slots) + 1) * page_size, end + 1, page_size): + slots.append(req.hybrid_len_to_big_page_id[length]) + tail_state = None + if len(mem_indexes) % page_size: + assert req.tail_small_page_buffer_id is not None + tail_state = self.small_page_buffers.get_state_cache(req.tail_small_page_buffer_id) + slots.append(self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID) + self.big_page_buffers.checkpoint_lens[slots[-1]] = self.small_page_buffers.checkpoint_lens[ + req.tail_small_page_buffer_id + ] + self.cpu_cache_transfer.offload( + mem_indexes, + page_indexes, + page_readies, + cpu_cache_client, + slots, + tail_state, + page_indexes_cpu=page_indexes_cpu, + page_readies_cpu=page_readies_cpu, + ) diff --git a/lightllm/common/hybrid_cache_transfer.py b/lightllm/common/hybrid_cache_transfer.py new file mode 100644 index 0000000000..d3a094b5f8 --- /dev/null +++ b/lightllm/common/hybrid_cache_transfer.py @@ -0,0 +1,142 @@ +import torch +import triton + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size + + +class HybridCpuCacheTransfer: + def __init__(self, kv_buffer, linear_config, big_page_buffers): + self.kv_buffer = kv_buffer + self.linear_config = linear_config + self.big_page_buffers = big_page_buffers + + def get_checkpoint_lens(self, page_indexes, cpu_cache_client): + cache = cpu_cache_client.cpu_kv_cache_tensor + metadata = cache.view(cache.shape[0], -1).view(torch.uint8)[:, -16:].view(torch.int64).numpy() + return [int(length) if version == 1 and length >= 0 else None for length, version in metadata[page_indexes]] + + def load(self, mem_indexes, page_indexes, cpu_cache_client, big_page_buffer_ids_cpu): + args = get_env_start_args() + assert mem_indexes.is_cuda and page_indexes.is_cuda + # 碎页情况的处理 + has_tail_page = len(mem_indexes) % args.cpu_cache_token_page_size != 0 + if has_tail_page: + padded_token_num = triton.cdiv( + len(mem_indexes), args.cpu_cache_token_page_size + ) * args.cpu_cache_token_page_size - len(mem_indexes) + mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), mode="constant", value=-1) + + big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( + non_blocking=True + ) + + assert len(big_page_buffer_ids_gpu) == len(page_indexes) + + from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( + copy_cpu_cache_to_kv_buffer, + ) + + copy_cpu_cache_to_kv_buffer( + mem_indexes=mem_indexes, + big_page_buffer_ids=big_page_buffer_ids_gpu, + page_indexes=page_indexes, + gpu_full_att_kv_state=self.kv_buffer, + cpu_kv_conv_state=self.big_page_buffers.conv_state_cache.buffer, + cpu_kv_ssm_state=self.big_page_buffers.ssm_state_cache.buffer, + cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, + tp_rank=get_current_rank_in_dp(), + tp_world_size=get_dp_world_size(), + big_page_token_num=args.cpu_cache_token_page_size, + linear_config=self.linear_config, + ) + + def offload( + self, + mem_indexes, + page_indexes, + page_readies, + cpu_cache_client, + big_page_buffer_ids_cpu, + tail_state, + page_indexes_cpu=None, + page_readies_cpu=None, + ): + args = get_env_start_args() + if not hasattr(self, "big_page_ids_buffer_store"): + self.big_page_ids_buffer_store = torch.empty((1024 * 1024 * 4,), dtype=torch.int64, device="cuda") + # 多申请3个cpu cache token page size，用于处理碎页情况，碎页情况需要将对应的大页数据拷贝到临时的大页上， + # 再从大页上拷贝到对应的运行态页面上 + self.mem_indexes_buffer = torch.empty( + (args.max_req_total_len + 3 * args.cpu_cache_token_page_size,), dtype=torch.int32, device="cuda" + ) + + assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda + + assert len(mem_indexes) % args.linear_att_hash_page_size == 0 + assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) + + if len(mem_indexes) % args.cpu_cache_token_page_size != 0: + # 存在不满大页的碎页的页面存在需要复制的情况 + dst_len = triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size + assert dst_len <= self.mem_indexes_buffer.shape[0] + dst_mem_indexes = self.mem_indexes_buffer[0:dst_len].fill_(-1) + dst_mem_indexes[0 : len(mem_indexes)].copy_(mem_indexes, non_blocking=True) + mem_indexes = dst_mem_indexes + from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( + copy_linear_att_state_to_linear_att_state, + ) + + src_conv_state, src_ssm_state = tail_state + dst_conv_state, dst_ssm_state = self.big_page_buffers.get_state_cache( + buffer_idx=big_page_buffer_ids_cpu[-1], + ) + copy_linear_att_state_to_linear_att_state( + src_conv_state=src_conv_state, + src_ssm_state=src_ssm_state, + dst_conv_state=dst_conv_state, + dst_ssm_state=dst_ssm_state, + ) + + assert len(big_page_buffer_ids_cpu) == len(page_indexes) == len(page_readies) + + if get_current_rank_in_dp() == 0: + if page_indexes_cpu is None: + page_indexes_cpu = page_indexes.cpu().tolist() + if page_readies_cpu is None: + page_readies_cpu = page_readies.cpu().tolist() + cache = cpu_cache_client.cpu_kv_cache_tensor + metadata = cache.view(cache.shape[0], -1).view(torch.uint8)[:, -16:].view(torch.int64).numpy() + for page_index, ready, slot in zip( + page_indexes_cpu, page_readies_cpu, big_page_buffer_ids_cpu, strict=True + ): + if page_index != -1 and not ready: + length = self.big_page_buffers.checkpoint_lens[slot] + metadata[page_index] = (-1 if length is None else length, 1) + + big_page_buffer_ids_cpu = torch.tensor( + big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu", pin_memory=True + ) + assert len(big_page_buffer_ids_cpu) <= self.big_page_ids_buffer_store.shape[0] + big_page_buffer_ids_gpu = self.big_page_ids_buffer_store[0 : len(big_page_buffer_ids_cpu)] + big_page_buffer_ids_gpu.copy_(big_page_buffer_ids_cpu, non_blocking=True) + + from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( + copy_kv_buffer_to_cpu_cache, + ) + + copy_kv_buffer_to_cpu_cache( + mem_indexes=mem_indexes, + page_indexes=page_indexes, + page_readies=page_readies, + big_page_buffer_ids=big_page_buffer_ids_gpu, + gpu_kv_full_att_state=self.kv_buffer, + cpu_kv_conv_state=self.big_page_buffers.conv_state_cache.buffer, + cpu_kv_ssm_state=self.big_page_buffers.ssm_state_cache.buffer, + cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, + tp_rank=get_current_rank_in_dp(), + tp_world_size=get_dp_world_size(), + big_page_token_num=args.cpu_cache_token_page_size, + linear_config=self.linear_config, + ) + return diff --git a/lightllm/common/indexer_state.py b/lightllm/common/indexer_state.py new file mode 100644 index 0000000000..06b8c72d2f --- /dev/null +++ b/lightllm/common/indexer_state.py @@ -0,0 +1,19 @@ +from lightllm.common.state_cache_manager import LayerCache + + +class IndexerStateManager: + def __init__(self, max_request_num, linear_config, mtp_step, device="cuda"): + self.linear_config = linear_config + self.state = LayerCache( + size=max_request_num + 1, + dtype=linear_config.full_att_dtype, + shape=(linear_config.index_kpool + mtp_step, 2 * linear_config.index_head_dim), + layer_num=linear_config.get_full_att_kv_layer_num_with_draft_model(), + device=device, + ) + + def get_layer_buffer(self, layer_index): + return self.state.buffer[self.linear_config.get_full_att_kv_layer_index(layer_index)] + + def reset(self, req_idx): + self.state.buffer[:, req_idx].zero_() diff --git a/lightllm/common/kv_cache_mem_manager/__init__.py b/lightllm/common/kv_cache_mem_manager/__init__.py index b3597b4bc8..211768a428 100644 --- a/lightllm/common/kv_cache_mem_manager/__init__.py +++ b/lightllm/common/kv_cache_mem_manager/__init__.py @@ -1,13 +1,8 @@ from .allocator import KvCacheAllocator from .mem_manager import MemoryManager, ReadOnlyStaticsMemoryManager -from .ppl_int8kv_mem_manager import PPLINT8KVMemoryManager -from .ppl_int4kv_mem_manager import PPLINT4KVMemoryManager -from .deepseek2_mem_manager import Deepseek2MemoryManager -from .deepseek3_2mem_manager import Deepseek3_2MemoryManager -from .fp8_per_token_group_quant_deepseek3_2mem_manager import FP8PerTokenGroupQuantDeepseek3_2MemoryManager -from .fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager -from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager -from .qwen3next_mem_manager import Qwen3NextMemManager +from .group_quantized_mem_manager import PPLINT4KVMemoryManager, PPLINT8KVMemoryManager +from .mla_mem_manager import MlaMemoryManager, IndexedMlaMemoryManager, FP8IndexedMlaMemoryManager +from .static_fp8_mem_manager import FP8StaticPerHeadQuantMemManager, FP8StaticPerTensorQuantMemManager from .glm5_next_mem_manager import Glm5NextMemManager __all__ = [ @@ -16,11 +11,10 @@ "ReadOnlyStaticsMemoryManager", "PPLINT4KVMemoryManager", "PPLINT8KVMemoryManager", - "Deepseek2MemoryManager", - "Deepseek3_2MemoryManager", - "FP8PerTokenGroupQuantDeepseek3_2MemoryManager", + "MlaMemoryManager", + "IndexedMlaMemoryManager", + "FP8IndexedMlaMemoryManager", "FP8StaticPerHeadQuantMemManager", "FP8StaticPerTensorQuantMemManager", - "Qwen3NextMemManager", "Glm5NextMemManager", ] diff --git a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py deleted file mode 100644 index 4be820b753..0000000000 --- a/lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py +++ /dev/null @@ -1,83 +0,0 @@ -import torch -import os -import torch.distributed as dist -from .mem_manager import MemoryManager -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 - - -logger = init_logger(__name__) - - -class Deepseek2MemoryManager(MemoryManager): - - operator_class = Deepseek2MemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Any: - kv = self.kv_buffer[layer_index] - return kv - - def get_cell_size(self): - return self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" - self.kv_buffer = torch.empty((layer_num, size + self.page_size, head_num, head_dim), dtype=dtype, device="cuda") - - def get_paged_kv_move_buffer_shape(self, page_num, page_size): - # DeepSeek MLA's single compressed KV latent is replicated across TP ranks, - # so the transfer buffer only needs one copy and does not multiply head_num by TP size. - return (page_num, page_size, self.layer_num, self.head_num, self.head_dim) - - def write_mem_to_page_kv_move_buffer( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - mla_page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=dp_mems[0].kv_buffer, - mode="write", - ) - return - - def read_page_kv_move_buffer_to_mem( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - for mem in dp_mems: - mla_page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=mem.kv_buffer, - mode="read", - ) diff --git a/lightllm/common/kv_cache_mem_manager/deepseek3_2mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek3_2mem_manager.py deleted file mode 100644 index 4034bcc8fa..0000000000 --- a/lightllm/common/kv_cache_mem_manager/deepseek3_2mem_manager.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -from typing import Any -from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager -from .operator import Deepseek3_2MemOperator - - -class Deepseek3_2MemoryManager(Deepseek2MemoryManager): - - operator_class = Deepseek3_2MemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - assert dtype in [torch.bfloat16, torch.float16] - # 因为V3.2 使用了NSA 稀疏的缘故，所以其head_dim 会比原始的kv 多 128 + 4 = 132 个字节 (128 fp8 + 4byte float32 scale)， - # 但是为了让整个数组具备16字节对齐，满足一些算子的约束，修改为添加 128 + 16 = 144 个字节, 这 144个字节中，后面132个字节用于 - # 存储真实数据，剩下12个，浪费了，只是占位。 - # 所以在子类中定制为其pad上，对外使用的接口，需要进行重载区别。 - super().__init__(size, dtype, head_num, head_dim + (144 // 2), layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Any: - kv = self.kv_buffer[layer_index][:, :, : (self.head_dim - (144 // 2))] - return kv - - def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor: - return self.kv_buffer[layer_index].view(dtype=torch.uint8)[:, :, -132:] diff --git a/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py deleted file mode 100644 index b72587545f..0000000000 --- a/lightllm/common/kv_cache_mem_manager/fp8_per_token_group_quant_deepseek3_2mem_manager.py +++ /dev/null @@ -1,66 +0,0 @@ -import torch -from typing import Any - -from .deepseek2_mem_manager import Deepseek2MemoryManager -from .operator import FP8PerTokenGroupQuantDeepseek3_2MemOperator - - -class FP8PerTokenGroupQuantDeepseek3_2MemoryManager(Deepseek2MemoryManager): - - operator_class = FP8PerTokenGroupQuantDeepseek3_2MemOperator - - kv_nope_dim = 512 - kv_rope_dim = 64 - # 576 = 512 + 64 - kv_head_dim = kv_nope_dim + kv_rope_dim - - quant_group_size = 128 - # 4 = 512 / 128 - quant_group_num = kv_nope_dim // quant_group_size - # 4 * 4 = quant_group_num * fp32 - # 64 * 2 = kv_rope_dim * bfloat16 - # 656 bytes = 512 + (4 * 4) + (64 * 2) - flashmla_bytes_per_token = kv_nope_dim + quant_group_num * 4 + kv_rope_dim * 2 - - indexer_head_dim = 128 - # 128 + 4 = indexer_head_dim + fp32 - # 132 bytes = 128 + 4 - indexer_bytes_per_token = indexer_head_dim + 4 - - # 16-byte 对齐，满足FlashMLA的对齐要求 - alignment = 16 - total_bytes_per_token = ( - (flashmla_bytes_per_token + indexer_bytes_per_token + alignment - 1) // alignment * alignment - ) - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - assert head_num == 1, "DeepSeek-V3.2 DSA FP8 path expects MQA-style head_num == 1" - self.prefill_dtype = dtype - super().__init__(size, torch.uint8, head_num, self.total_bytes_per_token, layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Any: - return self.kv_buffer[layer_index][:, :, : self.flashmla_bytes_per_token] - - def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor: - begin = self.flashmla_bytes_per_token - end = begin + self.indexer_bytes_per_token - return self.kv_buffer[layer_index][:, :, begin:end] - - def get_prefill_kv_cache_and_remap_indices( - self, - packed_kv: torch.Tensor, - topk_indices: torch.Tensor, - prefill_mem_index: torch.Tensor, - prefill_cache_kv: torch.Tensor, - ): - from lightllm.models.deepseek3_2.triton_kernel.prefill_compact_kv_flashmla_fp8 import ( - get_prefill_kv_cache_and_remap_indices_triton, - ) - - return get_prefill_kv_cache_and_remap_indices_triton( - packed_kv=packed_kv, - topk_mem_indices=topk_indices, - prefill_mem_index=prefill_mem_index, - prefill_cache_kv=prefill_cache_kv, - prefill_dtype=self.prefill_dtype, - ) diff --git a/lightllm/common/kv_cache_mem_manager/fp8_static_per_tensor_quant_mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_static_per_tensor_quant_mem_manager.py deleted file mode 100755 index 06338808bb..0000000000 --- a/lightllm/common/kv_cache_mem_manager/fp8_static_per_tensor_quant_mem_manager.py +++ /dev/null @@ -1,70 +0,0 @@ -import os -import json -import torch -import torch.distributed as dist -from typing import Tuple, Any -from lightllm.utils.config_utils import get_model_architectures -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from .mem_manager import MemoryManager -from .operator import FP8StaticPerTensorQuantMemOperator - -logger = init_logger(__name__) - - -class FP8StaticPerTensorQuantMemManager(MemoryManager): - - operator_class = FP8StaticPerTensorQuantMemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - # 这里用uint8存储量化后的kv，方便兼容各种torch算子。fp8量化目前采用离线方案，kv_buffer不存储scale - super().__init__(size, torch.uint8, head_num, head_dim, layer_num, always_copy, mem_fraction) - - self.qmax = torch.finfo(torch.float8_e4m3fn).max - self.qmin = torch.finfo(torch.float8_e4m3fn).min - self.scales = None - - if get_env_start_args().kv_quant_calibration_config_path is not None: - logger.info( - f"kv_quant_calibration_config_path {get_env_start_args().kv_quant_calibration_config_path} is set, " - "will load kv quant calibration config" - ) - cfg = self._load_and_check_config() - self.scales = torch.tensor(cfg["scales"], dtype=torch.float32, device="cuda").view(cfg["scales_shape"]) - else: - self.scales = torch.ones((self.kv_buffer.shape[0], 2), dtype=torch.float32, device="cuda") - - self.cpu_scales = self.scales.detach().cpu().numpy() - return - - def _load_and_check_config(self): - if os.path.exists(get_env_start_args().kv_quant_calibration_config_path): - with open(get_env_start_args().kv_quant_calibration_config_path, "r") as f: - cfg = json.load(f) - - if cfg["qmin"] != self.qmin: - raise ValueError(f"qmin {cfg['qmin']} in config not match torch.float8_e4m3fn.min {self.qmin}") - if cfg["qmax"] != self.qmax: - raise ValueError(f"qmax {cfg['qmax']} in config not match torch.float8_e4m3fn.max {self.qmax}") - model_arch = get_model_architectures(get_env_start_args().model_dir) - if cfg["architectures"] != model_arch: - raise ValueError( - f"architectures {cfg['architectures']} in config " f"not match current model_arch {model_arch}" - ) - if cfg["num_layers"] != self.layer_num: - raise ValueError( - f"num_layers {cfg['num_layers']} in config " f"not match current layer_num {self.layer_num}" - ) - assert ( - cfg["quant_type"] == "per_tensor" - ), f"quant type {cfg['quant_type']} in config not match per-tensor backend" - return cfg - else: - raise FileNotFoundError( - f"kv_quant_calibration_config {get_env_start_args().kv_quant_calibration_config_path} not found" - ) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - return k, v diff --git a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py index 45a3be667c..2101fec0e8 100644 --- a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py @@ -1,122 +1,41 @@ import torch -from lightllm.common.kv_cache_mem_manager.operator import LinearAttMemOperator -from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( - Qwen3NextMemManager, - Qwen3NextLinearAttPageHelper, -) +from .mla_mem_manager import MlaMemoryManager +from .operator import NormalMemOperator from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv -from lightllm.common.kv_trans_kernel.nixl_kv_trans import mla_page_io -class Glm5NextMemOperator(LinearAttMemOperator): +class Glm5NextMemOperator(NormalMemOperator): def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): output = self.mem_manager.get_att_input_params(layer_index)[:, :, : kv.shape[-1]] destindex_copy_kv(kv, mem_index, output) -class Glm5NextMemManager(Qwen3NextMemManager): - """One packed token buffer; KDA uses the standard big/small page pools.""" +class Glm5NextMemManager(MlaMemoryManager): + """Packed MLA and indexer KV storage, replicated across TP ranks.""" operator_class = Glm5NextMemOperator - def __init__(self, *args, mla_head_dim=512, **kwargs): - self.mla_head_dim = mla_head_dim - super().__init__(*args, **kwargs) - - def get_cell_size(self): - return self.head_dim * self.dtype.itemsize * self.layer_num - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): + def __init__( + self, + size, + dtype, + head_num, + head_dim, + layer_num, + always_copy=False, + mem_fraction=0.9, + mla_head_dim=512, + ): assert head_num == 1 - self.kv_buffer = torch.empty((layer_num, size + 1, 1, head_dim), dtype=dtype, device="cuda") - self._init_linear_att_buffers() + self.mla_head_dim = mla_head_dim + super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy, mem_fraction) def _layer_buffer(self, layer_index): - return self.kv_buffer[self.linear_config.get_full_att_kv_layer_index(layer_index)] + return self.kv_buffer[self.get_kv_layer_index(layer_index)] def get_att_input_params(self, layer_index): return self._layer_buffer(layer_index)[:, :, : self.mla_head_dim] def get_indexer_k_buffer(self, layer_index): return self._layer_buffer(layer_index).view(torch.uint8)[:, :, -132:] - - def write_to_shm(self, req_manager): - self.req_to_indexer_tail = req_manager.req_to_indexer_tail - return super().write_to_shm(req_manager) - - def _create_att_state_page_helper(self): - return Glm5NextAttStatePageHelper(self) - - def get_paged_kv_move_buffer_shape(self, page_num, page_size): - # Packed MLA/index KV is replicated across TP ranks. - return (page_num, page_size, self.layer_num, self.head_num, self.head_dim) - - def write_mem_to_page_kv_move_buffer( - self, mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind="kv", req_idx=None - ): - if page_kind != "kv": - return super().write_mem_to_page_kv_move_buffer( - mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind, req_idx - ) - pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] - pin_indexes.numpy()[:] = mem_indexes - mla_page_io( - mem_indexes=pin_indexes.cuda(non_blocking=True), - page_tensor=self.kv_move_buffer[page_index], - kv_buffer=mem_managers[dp_index * dp_world_size].kv_buffer, - mode="write", - ) - - def read_page_kv_move_buffer_to_mem( - self, mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind="kv", req_idx=None - ): - if page_kind != "kv": - return super().read_page_kv_move_buffer_to_mem( - mem_indexes, page_index, dp_index, mem_managers, dp_world_size, page_kind, req_idx - ) - pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] - pin_indexes.numpy()[:] = mem_indexes - indexes = pin_indexes.cuda(non_blocking=True) - for mem in mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size]: - mla_page_io( - mem_indexes=indexes, - page_tensor=self.kv_move_buffer[page_index], - kv_buffer=mem.kv_buffer, - mode="read", - ) - - -class Glm5NextAttStatePageHelper(Qwen3NextLinearAttPageHelper): - """Append the replicated K-pool tail to the global Conv/SSM state page.""" - - def __init__(self, mem_manager): - super().__init__(mem_manager) - self.tail_dtype = mem_manager.req_to_indexer_tail.buffer.dtype - self.tail_shape = ( - self.linear_config.get_full_att_kv_layer_num_with_draft_model(), - *mem_manager.req_to_indexer_tail.buffer.shape[2:], - ) - self.tail_offset = ((self.state_nbytes + 15) // 16) * 16 - self.tail_nbytes = self.tail_shape[0] * self.tail_shape[1] * self.tail_shape[2] * self.tail_dtype.itemsize - self.state_nbytes = self.tail_offset + self.tail_nbytes - - def _view_page_to_tail(self, page_index): - page_bytes = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) - return ( - page_bytes[self.tail_offset : self.tail_offset + self.tail_nbytes] - .view(self.tail_dtype) - .view(self.tail_shape) - ) - - def write_req_to_page(self, page_index, req_idx, dp_mems): - super().write_req_to_page(page_index, req_idx, dp_mems) - # The ring uses absolute token positions; sequence length already travels - # with the request, so restoring it needs no additional PD metadata. - self._view_page_to_tail(page_index).copy_(dp_mems[0].req_to_indexer_tail.buffer[:, req_idx], non_blocking=True) - - def read_page_to_req(self, page_index, req_idx, dp_mems): - super().read_page_to_req(page_index, req_idx, dp_mems) - tail = self._view_page_to_tail(page_index) - for mem in dp_mems: - mem.req_to_indexer_tail.buffer[:, req_idx].copy_(tail, non_blocking=True) diff --git a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/group_quantized_mem_manager.py old mode 100755 new mode 100644 similarity index 76% rename from lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py rename to lightllm/common/kv_cache_mem_manager/group_quantized_mem_manager.py index 6b02cbed63..7c25797378 --- a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/group_quantized_mem_manager.py @@ -1,12 +1,11 @@ import torch from typing import Tuple, Any from .mem_manager import MemoryManager -from .operator import PPLInt4KVMemOperator +from .operator import PPLInt4KVMemOperator, PPLInt8KVMemOperator -class PPLINT4KVMemoryManager(MemoryManager): - - operator_class = PPLInt4KVMemOperator +class GroupQuantizedMemoryManager(MemoryManager): + packing_factor = 1 def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=0.9): self.kv_dtype = torch.int8 @@ -14,6 +13,7 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=always_copy, mem_fraction=mem_fraction) def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: + layer_index = self.get_kv_layer_index(layer_index) k = self.kv_buffer[layer_index][:, : self.head_num, :] k_scale = self.scale_buffer[layer_index][:, : self.head_num, :] v = self.kv_buffer[layer_index][:, self.head_num :, :] @@ -21,7 +21,7 @@ def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: return (k, k_scale), (v, v_scale) def get_cell_size(self): - return 2 * self.head_num * self.head_dim // 2 * self.layer_num * torch._utils._element_size( + return 2 * self.head_num * (self.head_dim // self.packing_factor) * self.layer_num * torch._utils._element_size( self.kv_dtype ) + 2 * self.head_num * self.head_dim // self.group_quant_size * self.layer_num * torch._utils._element_size( self.dtype @@ -30,7 +30,9 @@ def get_cell_size(self): def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" self.kv_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim // 2), dtype=torch.int8, device="cuda" + (layer_num, size + self.page_size, 2 * head_num, head_dim // self.packing_factor), + dtype=torch.int8, + device="cuda", ) self.scale_buffer = torch.empty( (layer_num, size + self.page_size, 2 * head_num, head_dim // self.group_quant_size), @@ -48,3 +50,12 @@ def get_index_kv_buffer(self, index): def load_index_kv_buffer(self, index, load_tensor_dict): self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"]) self.scale_buffer[:, index].copy_(load_tensor_dict["scale_buffer"]) + + +class PPLINT8KVMemoryManager(GroupQuantizedMemoryManager): + operator_class = PPLInt8KVMemOperator + + +class PPLINT4KVMemoryManager(GroupQuantizedMemoryManager): + operator_class = PPLInt4KVMemOperator + packing_factor = 2 diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 38c870ce4e..05bb68fe74 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -1,35 +1,23 @@ import math -import os -import re import torch import torch.distributed as dist -import torch.multiprocessing as mp from typing import List, Tuple, Any, Union from lightllm.utils.log_utils import init_logger from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt from .allocator import KvCacheAllocator from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory -from lightllm.utils.dist_utils import ( - get_current_device_id, - get_current_rank_in_node, - get_node_world_size, -) +from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.shm_utils import get_service_shm_name from lightllm.utils.config_utils import get_num_key_value_heads -from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io -from lightllm.utils.device_utils import kv_trans_use_p2p -from lightllm.utils.shm_utils import create_or_link_shm -from multiprocessing.reduction import ForkingPickler -from filelock import FileLock from .operator import BaseMemManagerOperator, NormalMemOperator logger = init_logger(__name__) class MemoryManager: - operator_class = NormalMemOperator + kv_layer_mapping = None + kv_is_replicated = False def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): self.size = size @@ -61,7 +49,20 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False # 构建对外的操作类接口 self.operator: BaseMemManagerOperator = self.operator_class(self) + def set_kv_layer_mapping(self, layer_mapping): + layer_mapping = tuple(layer_mapping) + if not layer_mapping or any(index is not None and not 0 <= index < self.layer_num for index in layer_mapping): + raise ValueError("KV layer mapping must refer to allocated layers") + self.kv_layer_mapping = layer_mapping + + def get_kv_layer_index(self, layer_index): + kv_layer_index = layer_index if self.kv_layer_mapping is None else self.kv_layer_mapping[layer_index] + if kv_layer_index is None: + raise ValueError(f"layer {layer_index} has no KV storage") + return kv_layer_index + def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: + layer_index = self.get_kv_layer_index(layer_index) k = self.kv_buffer[layer_index][:, : self.head_num, :] v = self.kv_buffer[layer_index][:, self.head_num :, :] return k, v @@ -69,6 +70,11 @@ def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: def get_cell_size(self): return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype) + def get_transfer_view(self, *, include_page_layout=False): + from lightllm.common.kv_transfer_resources import KVTransferView + + return KVTransferView.from_mem_manager(self, include_page_layout=include_page_layout) + def get_paged_kv_move_buffer_shape(self, page_num, page_size): num_kv_head = get_num_key_value_heads(get_env_start_args().model_dir) return ( @@ -98,7 +104,7 @@ def profile_size(self, mem_fraction): available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) cell_size = self.get_cell_size() pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size + available_memory_bytes = available_memory * 1024**3 - pd_kv_move_buffer_size self.size = int(available_memory_bytes / cell_size) if world_size > 1: tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") @@ -121,79 +127,6 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): (layer_num, size + self.page_size, 2 * head_num, head_dim), dtype=dtype, device="cuda" ) - def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: - self.kv_move_buffer = torch.empty( - self.get_paged_kv_move_buffer_shape(page_num, page_size), dtype=self.dtype, device="cuda" - ) - self._buffer_mem_indexes_tensors = [ - torch.empty((page_size,), dtype=torch.int64, device="cpu", pin_memory=True) for _ in range(page_num) - ] - return self.kv_move_buffer - - def write_mem_to_page_kv_move_buffer( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - repeat_count = dp_world_size * self.kv_buffer.shape[2] // self.kv_move_buffer.shape[3] - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - for tp_index in range(dp_world_size): - if tp_index % repeat_count == 0: - page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=dp_mems[tp_index].kv_buffer, - tp_index=tp_index, - tp_world_size=dp_world_size, - mode="write", - ) - # keep for debug - # logger.info(f"src token tensor {self.kv_buffer[:, mem_indexes[0], 0, 0]}") - # logger.info(f"src page token tensor {cur_page[0, :, 0, 0]}") - return - - def read_page_kv_move_buffer_to_mem( - self, - mem_indexes: List[int], - page_index: int, - dp_index: int, - mem_managers: List["MemoryManager"], - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" - cur_page = self.kv_move_buffer[page_index] - pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] - pin_mem_indexes.numpy()[:] = mem_indexes - mem_indexes_gpu = pin_mem_indexes.cuda(non_blocking=True) - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - mem_indexes_gpu = torch.tensor(mem_indexes, dtype=torch.int64, device="cpu", pin_memory=True).cuda( - non_blocking=True - ) - for tp_index in range(dp_world_size): - page_io( - mem_indexes=mem_indexes_gpu, - page_tensor=cur_page, - kv_buffer=dp_mems[tp_index].kv_buffer, - tp_index=tp_index, - tp_world_size=dp_world_size, - mode="read", - ) - # keep for debug - # 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 @@ -232,57 +165,6 @@ def get_index_kv_buffer(self, index): def load_index_kv_buffer(self, index, load_tensor_dict): self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"]) - def write_to_shm(self, req_manager): - """ - 将 mem manager 写入到 shm中，方便pd分离等特性直接从中读取，不依赖进程间队列。 - """ - if kv_trans_use_p2p(): - from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor - - mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__ - - from lightllm.common.req_manager import ReqManager - - req_manager: ReqManager = 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 - - lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") - with lock: - node_world_size = get_node_world_size() - shm_name = f"mem_manager_{get_current_rank_in_node()}" - obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] - obj_size = len(obj_bytes_array[0]) - shm = create_or_link_shm( - name=shm_name, expected_size=obj_size * (node_world_size * 2) + 4 + 4, force_mode="create" - ) - logger.info(f"create shm {shm.name} size {shm.size} for mem manger shared buffer") - shm.buf[0:4] = (node_world_size * 2).to_bytes(4, "little") - shm.buf[4:8] = obj_size.to_bytes(4, "little") - start_index = 8 - for obj_bytes in obj_bytes_array: - shm.buf[start_index : start_index + obj_size] = obj_bytes - start_index += obj_size - - @staticmethod - def loads_from_shm(rank_in_node: int) -> "MemoryManager": - shm_name = f"mem_manager_{rank_in_node}" - lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") - logger.info(f"get memmanager from shm {shm_name}") - with lock: - shm = create_or_link_shm(name=shm_name, expected_size=-1, force_mode="link") - left_num = int.from_bytes(shm.buf[0:4], "little") - obj_size = int.from_bytes(shm.buf[4:8], "little") - assert left_num > 0 - end_index = 8 + left_num * obj_size - start_index = 8 + (left_num - 1) * obj_size - obj_bytes = shm.buf[start_index:end_index].tobytes() - shm.buf[0:4] = (left_num - 1).to_bytes(4, byteorder="little") - shm.close() - return ForkingPickler.loads(obj_bytes) - class ReadOnlyStaticsMemoryManager: """ diff --git a/lightllm/common/kv_cache_mem_manager/mem_utils.py b/lightllm/common/kv_cache_mem_manager/mem_utils.py index 79ea448794..efb8b524fb 100644 --- a/lightllm/common/kv_cache_mem_manager/mem_utils.py +++ b/lightllm/common/kv_cache_mem_manager/mem_utils.py @@ -2,14 +2,11 @@ MemoryManager, PPLINT8KVMemoryManager, PPLINT4KVMemoryManager, - Deepseek2MemoryManager, - Deepseek3_2MemoryManager, FP8StaticPerHeadQuantMemManager, FP8StaticPerTensorQuantMemManager, ) from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.llm_utils import get_llm_model_class from functools import lru_cache logger = init_logger(__name__) @@ -17,52 +14,17 @@ @lru_cache(maxsize=None) def select_mem_manager_class(): - # case 1 - # 先判断是否是 deepseek 系列的模型 - model_class = get_llm_model_class() - - from lightllm.models import Deepseek3_2TpPartModel - - if issubclass(model_class, Deepseek3_2TpPartModel): - if get_env_start_args().llm_kv_type == "fp8kv_dsa": - from . import FP8PerTokenGroupQuantDeepseek3_2MemoryManager - - mem_class = FP8PerTokenGroupQuantDeepseek3_2MemoryManager - else: - mem_class = Deepseek3_2MemoryManager - logger.info(f"Model kv cache using default, mem_manager class: {mem_class}") - return mem_class - - from lightllm.models import Deepseek2TpPartModel - - if issubclass(model_class, Deepseek2TpPartModel): - mem_class = Deepseek2MemoryManager - logger.info(f"Model kv cache using default, mem_manager class: {mem_class}") - return mem_class - - # case normal - logger.info(f"mode setting params: {get_env_start_args().llm_kv_type}") - if get_env_start_args().llm_kv_type == "int8kv": - memory_manager_class = PPLINT8KVMemoryManager - elif get_env_start_args().llm_kv_type == "int4kv": - memory_manager_class = PPLINT4KVMemoryManager - elif get_env_start_args().llm_kv_type == "fp8kv_sph": - memory_manager_class = FP8StaticPerHeadQuantMemManager - elif get_env_start_args().llm_kv_type == "fp8kv_spt": - memory_manager_class = FP8StaticPerTensorQuantMemManager - elif get_env_start_args().llm_kv_type == "None": - memory_manager_class = MemoryManager - + kv_type = get_env_start_args().llm_kv_type + managers = { + "int8kv": PPLINT8KVMemoryManager, + "int4kv": PPLINT4KVMemoryManager, + "fp8kv_sph": FP8StaticPerHeadQuantMemManager, + "fp8kv_spt": FP8StaticPerTensorQuantMemManager, + "None": MemoryManager, + } + try: + memory_manager_class = managers[kv_type] + except KeyError: + raise ValueError(f"Unsupported KV cache type: {kv_type}") from None logger.info(f"Model kv cache using mem_manager class: {memory_manager_class}") return memory_manager_class - - -@lru_cache(maxsize=None) -def used_mem_manager_has_scale() -> bool: - mem_class = select_mem_manager_class() - return mem_class in [ - PPLINT8KVMemoryManager, - PPLINT4KVMemoryManager, - FP8StaticPerHeadQuantMemManager, - FP8StaticPerTensorQuantMemManager, - ] diff --git a/lightllm/common/kv_cache_mem_manager/mla_mem_manager.py b/lightllm/common/kv_cache_mem_manager/mla_mem_manager.py new file mode 100644 index 0000000000..ff74ba164d --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/mla_mem_manager.py @@ -0,0 +1,116 @@ +import torch +from .mem_manager import MemoryManager +from typing import Any +from lightllm.utils.log_utils import init_logger +from .operator import MlaMemOperator, FP8IndexedMlaMemOperator + + +logger = init_logger(__name__) + + +class MlaMemoryManager(MemoryManager): + + kv_is_replicated = True + operator_class = MlaMemOperator + + def get_att_input_params(self, layer_index: int) -> Any: + layer_index = self.get_kv_layer_index(layer_index) + kv = self.kv_buffer[layer_index] + return kv + + def get_cell_size(self): + return self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype) + + def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): + assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" + self.kv_buffer = torch.empty((layer_num, size + self.page_size, head_num, head_dim), dtype=dtype, device="cuda") + + def get_paged_kv_move_buffer_shape(self, page_num, page_size): + # DeepSeek MLA's single compressed KV latent is replicated across TP ranks, + # so the transfer buffer only needs one copy and does not multiply head_num by TP size. + return (page_num, page_size, self.layer_num, self.head_num, self.head_dim) + + +class IndexedMlaMemoryManager(MlaMemoryManager): + + operator_class = MlaMemOperator + + def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): + assert dtype in [torch.bfloat16, torch.float16] + # 因为V3.2 使用了NSA 稀疏的缘故，所以其head_dim 会比原始的kv 多 128 + 4 = 132 个字节 (128 fp8 + 4byte float32 scale)， + # 但是为了让整个数组具备16字节对齐，满足一些算子的约束，修改为添加 128 + 16 = 144 个字节, 这 144个字节中，后面132个字节用于 + # 存储真实数据，剩下12个，浪费了，只是占位。 + # 所以在子类中定制为其pad上，对外使用的接口，需要进行重载区别。 + super().__init__(size, dtype, head_num, head_dim + (144 // 2), layer_num, always_copy, mem_fraction) + + def get_att_input_params(self, layer_index: int) -> Any: + layer_index = self.get_kv_layer_index(layer_index) + kv = self.kv_buffer[layer_index][:, :, : (self.head_dim - (144 // 2))] + return kv + + def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor: + layer_index = self.get_kv_layer_index(layer_index) + return self.kv_buffer[layer_index].view(dtype=torch.uint8)[:, :, -132:] + + +class FP8IndexedMlaMemoryManager(MlaMemoryManager): + + operator_class = FP8IndexedMlaMemOperator + + kv_nope_dim = 512 + kv_rope_dim = 64 + # 576 = 512 + 64 + kv_head_dim = kv_nope_dim + kv_rope_dim + + quant_group_size = 128 + # 4 = 512 / 128 + quant_group_num = kv_nope_dim // quant_group_size + # 4 * 4 = quant_group_num * fp32 + # 64 * 2 = kv_rope_dim * bfloat16 + # 656 bytes = 512 + (4 * 4) + (64 * 2) + flashmla_bytes_per_token = kv_nope_dim + quant_group_num * 4 + kv_rope_dim * 2 + + indexer_head_dim = 128 + # 128 + 4 = indexer_head_dim + fp32 + # 132 bytes = 128 + 4 + indexer_bytes_per_token = indexer_head_dim + 4 + + # 16-byte 对齐，满足FlashMLA的对齐要求 + alignment = 16 + total_bytes_per_token = ( + (flashmla_bytes_per_token + indexer_bytes_per_token + alignment - 1) // alignment * alignment + ) + + def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): + assert head_num == 1, "DeepSeek-V3.2 DSA FP8 path expects MQA-style head_num == 1" + self.prefill_dtype = dtype + super().__init__(size, torch.uint8, head_num, self.total_bytes_per_token, layer_num, always_copy, mem_fraction) + + def get_att_input_params(self, layer_index: int) -> Any: + layer_index = self.get_kv_layer_index(layer_index) + return self.kv_buffer[layer_index][:, :, : self.flashmla_bytes_per_token] + + def get_indexer_k_buffer(self, layer_index: int) -> torch.Tensor: + layer_index = self.get_kv_layer_index(layer_index) + begin = self.flashmla_bytes_per_token + end = begin + self.indexer_bytes_per_token + return self.kv_buffer[layer_index][:, :, begin:end] + + def get_prefill_kv_cache_and_remap_indices( + self, + packed_kv: torch.Tensor, + topk_indices: torch.Tensor, + prefill_mem_index: torch.Tensor, + prefill_cache_kv: torch.Tensor, + ): + from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( + get_prefill_kv_cache_and_remap_indices_triton, + ) + + return get_prefill_kv_cache_and_remap_indices_triton( + packed_kv=packed_kv, + topk_mem_indices=topk_indices, + prefill_mem_index=prefill_mem_index, + prefill_cache_kv=prefill_cache_kv, + prefill_dtype=self.prefill_dtype, + ) diff --git a/lightllm/common/kv_cache_mem_manager/operator/__init__.py b/lightllm/common/kv_cache_mem_manager/operator/__init__.py index 85c37ad39b..84c3ca1c09 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/__init__.py +++ b/lightllm/common/kv_cache_mem_manager/operator/__init__.py @@ -1,11 +1,9 @@ from .base import BaseMemManagerOperator from .normal import NormalMemOperator from .quant import QuantScaleMemOperator, PPLInt4KVMemOperator, PPLInt8KVMemOperator -from .linear_att import LinearAttMemOperator -from .deepseek import ( - Deepseek2MemOperator, - Deepseek3_2MemOperator, - FP8PerTokenGroupQuantDeepseek3_2MemOperator, +from .mla import ( + MlaMemOperator, + FP8IndexedMlaMemOperator, ) from .fp8_quant import ( FP8StaticPerHeadQuantMemOperator, diff --git a/lightllm/common/kv_cache_mem_manager/operator/base.py b/lightllm/common/kv_cache_mem_manager/operator/base.py index 682b8a5d60..14e5b57786 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/base.py +++ b/lightllm/common/kv_cache_mem_manager/operator/base.py @@ -6,8 +6,11 @@ from ..mem_manager import MemoryManager from lightllm.server.router.model_infer.infer_batch import InferReq + # 定义一个抽象基类 class BaseMemManagerOperator(ABC): + supports_dp_cache_fetch = False + def __init__(self, mem_manager: "MemoryManager") -> None: super().__init__() self.mem_manager = mem_manager @@ -17,12 +20,18 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: pass def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): - raise NotImplementedError() + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + src_mem_index = src_mem_index.cuda(non_blocking=True) + dst_mem_index = dst_mem_index.cuda(non_blocking=True) + copy_kv_buffer_to_kv_buffer(src_mem_index, dst_mem_index, self.mem_manager.kv_buffer) + if hasattr(self.mem_manager, "scale_buffer"): + copy_kv_buffer_to_kv_buffer(src_mem_index, dst_mem_index, self.mem_manager.scale_buffer) # cpu cache 的相关操作接口 def load_cpu_cache_to_gpu( self, mem_indexes: torch.Tensor, page_indexes: torch.Tensor, cpu_cache_client, req: "InferReq" - ): + ) -> None: raise NotImplementedError() def offload_gpu_kv_to_cpu_cache( @@ -38,7 +47,7 @@ def offload_gpu_kv_to_cpu_cache( # dp 间共享 kv 的操作接口, 提升dp 模式下的kv 共享效率，降低调度的难度 def copy_kv_from_other_dp_ranks( self, - mem_managers: List["MemoryManager"], + kv_views: List, move_token_indexes: torch.Tensor, token_dp_indexes: torch.Tensor, mem_indexes: torch.Tensor, diff --git a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py b/lightllm/common/kv_cache_mem_manager/operator/deepseek.py deleted file mode 100644 index 6e05b96e10..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py +++ /dev/null @@ -1,80 +0,0 @@ -import torch -from .normal import NormalMemOperator -from .base import BaseMemManagerOperator -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -class Deepseek2MemOperator(NormalMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager - - mem_manager: Deepseek2MemoryManager = self.mem_manager - - from ...basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv - - rope_dim = 64 - kv_lora_rank = kv.shape[2] - rope_dim - assert kv_lora_rank + rope_dim == mem_manager.kv_buffer.shape[-1] - - destindex_copy_kv( - kv[:, :, :kv_lora_rank], - kv[:, :, kv_lora_rank:], - mem_index, - mem_manager.kv_buffer[layer_index][:, :, :kv_lora_rank], - mem_manager.kv_buffer[layer_index][:, :, kv_lora_rank:], - ) - return - - -class Deepseek3_2MemOperator(Deepseek2MemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.deepseek3_2mem_manager import Deepseek3_2MemoryManager - - mem_manager: Deepseek3_2MemoryManager = self.mem_manager - from ...basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv - - rope_dim = 64 - kv_lora_rank = kv.shape[2] - rope_dim - assert kv_lora_rank + rope_dim == mem_manager.kv_buffer.shape[-1] - (144 // 2) - - destindex_copy_kv( - kv[:, :, :kv_lora_rank], - kv[:, :, kv_lora_rank:], - mem_index, - mem_manager.kv_buffer[layer_index][:, :, :kv_lora_rank], - mem_manager.kv_buffer[layer_index][:, :, kv_lora_rank : (kv_lora_rank + rope_dim)], - ) - return - - -class FP8PerTokenGroupQuantDeepseek3_2MemOperator(BaseMemManagerOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.fp8_per_token_group_quant_deepseek3_2mem_manager import ( - FP8PerTokenGroupQuantDeepseek3_2MemoryManager, - ) - - mem_manager: FP8PerTokenGroupQuantDeepseek3_2MemoryManager = self.mem_manager - from lightllm.models.deepseek3_2.triton_kernel.destindex_copy_kv_flashmla_fp8 import ( - destindex_copy_kv_flashmla_fp8, - ) - - rope_dim = 64 - kv_lora_rank = kv.shape[2] - rope_dim - assert kv_lora_rank == 512, f"Expected kv_lora_rank=512, got {kv_lora_rank}" - - flashmla_bytes_per_token = mem_manager.flashmla_bytes_per_token - - o_nope = mem_manager.kv_buffer[layer_index][:, :, :512].view(torch.float8_e4m3fn) - o_scale = mem_manager.kv_buffer[layer_index][:, :, 512:528].view(torch.float32) - o_rope = mem_manager.kv_buffer[layer_index][:, :, 528:flashmla_bytes_per_token].view(torch.bfloat16) - destindex_copy_kv_flashmla_fp8( - kv[:, :, :kv_lora_rank], - kv[:, :, kv_lora_rank:], - mem_index, - o_nope, - o_scale, - o_rope, - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py b/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py index 4b1f0ac207..057d28f311 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py +++ b/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py @@ -10,7 +10,8 @@ class FP8StaticPerHeadQuantMemOperator(BaseMemManagerOperator): def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.fp8_static_per_head_quant_mem_manager import ( + layer_index = self.mem_manager.get_kv_layer_index(layer_index) + from lightllm.common.kv_cache_mem_manager.static_fp8_mem_manager import ( FP8StaticPerHeadQuantMemManager, ) @@ -31,6 +32,7 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: class FP8StaticPerTensorQuantMemOperator(BaseMemManagerOperator): def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): + layer_index = self.mem_manager.get_kv_layer_index(layer_index) from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager mem_manager: MemoryManager = self.mem_manager diff --git a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py deleted file mode 100644 index e935539e37..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py +++ /dev/null @@ -1,210 +0,0 @@ -import torch -import triton -from typing import List -from typing import TYPE_CHECKING -from .base import BaseMemManagerOperator -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size -from lightllm.utils.log_utils import init_logger - -if TYPE_CHECKING: - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient - from lightllm.server.router.model_infer.infer_batch import InferReq - -logger = init_logger(__name__) - - -class LinearAttMemOperator(BaseMemManagerOperator): - """ - 只用于非量化的linear att 混合 full att的模型，列入 qwen3.5 - """ - - def __init__(self, mem_manager): - super().__init__(mem_manager) - self.linear_config = mem_manager.linear_config - - def load_cpu_cache_to_gpu( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda - args = get_env_start_args() - assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager - - mem_manager: Qwen3NextMemManager = self.mem_manager - - big_page_num = len(mem_indexes) // args.cpu_cache_token_page_size - max_kv_len = (req.cur_kv_len // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - assert max_kv_len % args.cpu_cache_token_page_size == 0 - - big_page_buffer_ids_cpu = [] - for i in range(big_page_num): - page_id = mem_manager.big_page_buffers.alloc_one_state_cache() - assert page_id is not None - req.hybrid_len_to_big_page_id[max_kv_len] = page_id - big_page_buffer_ids_cpu.append(page_id) - max_kv_len -= args.cpu_cache_token_page_size - assert max_kv_len % args.cpu_cache_token_page_size == 0 - - big_page_buffer_ids_cpu.reverse() - - # 碎页情况的处理 - has_tail_page = len(mem_indexes) % args.cpu_cache_token_page_size != 0 - if has_tail_page: - padded_token_num = triton.cdiv( - len(mem_indexes), args.cpu_cache_token_page_size - ) * args.cpu_cache_token_page_size - len(mem_indexes) - mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), mode="constant", value=-1) - - # 将对应的小叶数据拷贝到临时的大页上，再从大页上拷贝到对应的运行态页面上 - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) - - big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( - non_blocking=True - ) - - assert len(big_page_buffer_ids_gpu) == len(page_indexes) - - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_cpu_cache_to_kv_buffer, - ) - - copy_cpu_cache_to_kv_buffer( - mem_indexes=mem_indexes, - big_page_buffer_ids=big_page_buffer_ids_gpu, - page_indexes=page_indexes, - gpu_full_att_kv_state=mem_manager.kv_buffer, - cpu_kv_conv_state=mem_manager.big_page_buffers.conv_state_cache.buffer, - cpu_kv_ssm_state=mem_manager.big_page_buffers.ssm_state_cache.buffer, - cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, - tp_rank=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - big_page_token_num=args.cpu_cache_token_page_size, - linear_config=self.linear_config, - ) - - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=big_page_buffer_ids_cpu[-1], - req=req, - ) - - return - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - args = get_env_start_args() - if not hasattr(self, "big_page_ids_buffer_store"): - self.big_page_ids_buffer_store = torch.empty((1024 * 1024 * 4,), dtype=torch.int64, device="cuda") - # 多申请3个cpu cache token page size，用于处理碎页情况，碎页情况需要将对应的大页数据拷贝到临时的大页上， - # 再从大页上拷贝到对应的运行态页面上 - self.mem_indexes_buffer = torch.empty( - (args.max_req_total_len + 3 * args.cpu_cache_token_page_size,), dtype=torch.int32, device="cuda" - ) - - assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda - - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) - - from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager - - mem_manager: Qwen3NextMemManager = self.mem_manager - - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - big_page_buffer_ids_cpu = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) - max_kv_len = (len(mem_indexes) // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - start_kv_len = (len(big_page_buffer_ids_cpu) + 1) * args.cpu_cache_token_page_size - for seq_len in range(start_kv_len, max_kv_len + 1, args.cpu_cache_token_page_size): - page_id = req.hybrid_len_to_big_page_id[seq_len] - big_page_buffer_ids_cpu.append(page_id) - - if len(mem_indexes) % args.cpu_cache_token_page_size != 0: - # 存在不满大页的碎页的页面存在需要复制的情况 - dst_len = triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - assert dst_len <= self.mem_indexes_buffer.shape[0] - dst_mem_indexes = self.mem_indexes_buffer[0:dst_len].fill_(-1) - dst_mem_indexes[0 : len(mem_indexes)].copy_(mem_indexes, non_blocking=True) - mem_indexes = dst_mem_indexes - assert req.tail_small_page_buffer_id is not None - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_linear_att_state_to_linear_att_state, - ) - - src_conv_state, src_ssm_state = g_infer_context.radix_cache.small_page_buffers.get_state_cache( - buffer_idx=req.tail_small_page_buffer_id - ) - dst_conv_state, dst_ssm_state = mem_manager.big_page_buffers.get_state_cache( - buffer_idx=mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID, - ) - copy_linear_att_state_to_linear_att_state( - src_conv_state=src_conv_state, - src_ssm_state=src_ssm_state, - dst_conv_state=dst_conv_state, - dst_ssm_state=dst_ssm_state, - ) - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID) - - assert len(big_page_buffer_ids_cpu) == len(page_indexes) == len(page_readies) - - big_page_buffer_ids_cpu = torch.tensor( - big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu", pin_memory=True - ) - assert len(big_page_buffer_ids_cpu) <= self.big_page_ids_buffer_store.shape[0] - big_page_buffer_ids_gpu = self.big_page_ids_buffer_store[0 : len(big_page_buffer_ids_cpu)] - big_page_buffer_ids_gpu.copy_(big_page_buffer_ids_cpu, non_blocking=True) - - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_kv_buffer_to_cpu_cache, - ) - - copy_kv_buffer_to_cpu_cache( - mem_indexes=mem_indexes, - page_indexes=page_indexes, - page_readies=page_readies, - big_page_buffer_ids=big_page_buffer_ids_gpu, - gpu_kv_full_att_state=mem_manager.kv_buffer, - cpu_kv_conv_state=mem_manager.big_page_buffers.conv_state_cache.buffer, - cpu_kv_ssm_state=mem_manager.big_page_buffers.ssm_state_cache.buffer, - cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, - tp_rank=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - big_page_token_num=args.cpu_cache_token_page_size, - linear_config=self.linear_config, - ) - return - - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - # Qwen3Next 需要调整 layer_index - layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import ( - destindex_copy_kv, - ) - - destindex_copy_kv(kv, mem_index, mem_manager.kv_buffer[layer_index]) - return - - def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer - - copy_kv_buffer_to_kv_buffer( - src_mem_index.cuda(non_blocking=True), dst_mem_index.cuda(non_blocking=True), self.mem_manager.kv_buffer - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/operator/mla.py b/lightllm/common/kv_cache_mem_manager/operator/mla.py new file mode 100644 index 0000000000..9e101f9c97 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/operator/mla.py @@ -0,0 +1,55 @@ +import torch +from .normal import NormalMemOperator +from .base import BaseMemManagerOperator +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +class MlaMemOperator(NormalMemOperator): + def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): + from ...basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv + + output = self.mem_manager.get_att_input_params(layer_index) + rope_dim = 64 + kv_lora_rank = kv.shape[2] - rope_dim + assert kv.shape[2] == output.shape[-1] + destindex_copy_kv( + kv[:, :, :kv_lora_rank], + kv[:, :, kv_lora_rank:], + mem_index, + output[:, :, :kv_lora_rank], + output[:, :, kv_lora_rank:], + ) + + +class FP8IndexedMlaMemOperator(BaseMemManagerOperator): + def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): + layer_index = self.mem_manager.get_kv_layer_index(layer_index) + from lightllm.common.kv_cache_mem_manager.mla_mem_manager import ( + FP8IndexedMlaMemoryManager, + ) + + mem_manager: FP8IndexedMlaMemoryManager = self.mem_manager + from lightllm.common.basemodel.triton_kernel.kv_copy.destindex_copy_kv_flashmla_fp8 import ( + destindex_copy_kv_flashmla_fp8, + ) + + rope_dim = 64 + kv_lora_rank = kv.shape[2] - rope_dim + assert kv_lora_rank == 512, f"Expected kv_lora_rank=512, got {kv_lora_rank}" + + flashmla_bytes_per_token = mem_manager.flashmla_bytes_per_token + + o_nope = mem_manager.kv_buffer[layer_index][:, :, :512].view(torch.float8_e4m3fn) + o_scale = mem_manager.kv_buffer[layer_index][:, :, 512:528].view(torch.float32) + o_rope = mem_manager.kv_buffer[layer_index][:, :, 528:flashmla_bytes_per_token].view(torch.bfloat16) + destindex_copy_kv_flashmla_fp8( + kv[:, :, :kv_lora_rank], + kv[:, :, kv_lora_rank:], + mem_index, + o_nope, + o_scale, + o_rope, + ) + return diff --git a/lightllm/common/kv_cache_mem_manager/operator/normal.py b/lightllm/common/kv_cache_mem_manager/operator/normal.py index 3c53ace079..e361ced90d 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/normal.py +++ b/lightllm/common/kv_cache_mem_manager/operator/normal.py @@ -13,7 +13,10 @@ class NormalMemOperator(BaseMemManagerOperator): + supports_dp_cache_fetch = True + def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): + layer_index = self.mem_manager.get_kv_layer_index(layer_index) from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager mem_manager: MemoryManager = self.mem_manager @@ -86,7 +89,7 @@ def offload_gpu_kv_to_cpu_cache( def copy_kv_from_other_dp_ranks( self, - mem_managers: List, + kv_views: List, move_token_indexes: torch.Tensor, token_dp_indexes: torch.Tensor, mem_indexes: torch.Tensor, @@ -96,8 +99,8 @@ def copy_kv_from_other_dp_ranks( if not hasattr(self, "mem_ptrs_tensor"): # 构建一个2D tensor，shape为(layer_num, mem_num) mems_ptr_list = [] - for i in range(0, len(mem_managers)): - mems_ptr_list.append(mem_managers[i].kv_buffer.data_ptr()) + for i in range(0, len(kv_views)): + mems_ptr_list.append(kv_views[i].kv_buffer.data_ptr()) self.mem_ptrs_tensor = torch.tensor(mems_ptr_list, dtype=torch.uint64, device="cpu", pin_memory=True) from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_for_dp diff --git a/lightllm/common/kv_cache_mem_manager/operator/quant.py b/lightllm/common/kv_cache_mem_manager/operator/quant.py index a3a1c1d01c..d37dcbc4c2 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/quant.py +++ b/lightllm/common/kv_cache_mem_manager/operator/quant.py @@ -93,7 +93,8 @@ def offload_gpu_kv_to_cpu_cache( class PPLInt4KVMemOperator(QuantScaleMemOperator): def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager + layer_index = self.mem_manager.get_kv_layer_index(layer_index) + from lightllm.common.kv_cache_mem_manager.group_quantized_mem_manager import PPLINT4KVMemoryManager mem_manager: PPLINT4KVMemoryManager = self.mem_manager from ...basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import ( @@ -112,7 +113,8 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: class PPLInt8KVMemOperator(QuantScaleMemOperator): def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager + layer_index = self.mem_manager.get_kv_layer_index(layer_index) + from lightllm.common.kv_cache_mem_manager.group_quantized_mem_manager import PPLINT8KVMemoryManager mem_manager: PPLINT8KVMemoryManager = self.mem_manager from ...basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import ( 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 deleted file mode 100755 index 0f5e8e9370..0000000000 --- a/lightllm/common/kv_cache_mem_manager/ppl_int8kv_mem_manager.py +++ /dev/null @@ -1,49 +0,0 @@ -import torch -from typing import Tuple, Any -from .mem_manager import MemoryManager -from .operator import PPLInt8KVMemOperator - - -class PPLINT8KVMemoryManager(MemoryManager): - operator_class = PPLInt8KVMemOperator - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=0.9): - self.kv_dtype = torch.int8 - self.group_quant_size = 8 - super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=always_copy, mem_fraction=mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - k_scale = self.scale_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - v_scale = self.scale_buffer[layer_index][:, self.head_num :, :] - return (k, k_scale), (v, v_scale) - - def get_cell_size(self): - return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size( - self.kv_dtype - ) + 2 * self.head_num * self.head_dim // self.group_quant_size * self.layer_num * torch._utils._element_size( - self.dtype - ) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - assert size % self.page_size == 0, f"KV cache size {size} must be a multiple of page_size {self.page_size}" - self.kv_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim), dtype=torch.int8, device="cuda" - ) - self.scale_buffer = torch.empty( - (layer_num, size + self.page_size, 2 * head_num, head_dim // self.group_quant_size), - dtype=dtype, - 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]} - - def load_index_kv_buffer(self, index, load_tensor_dict): - self.kv_buffer[:, index].copy_(load_tensor_dict["kv_buffer"]) - self.scale_buffer[:, index].copy_(load_tensor_dict["scale_buffer"]) diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py deleted file mode 100644 index 1962fa71b0..0000000000 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ /dev/null @@ -1,464 +0,0 @@ -import torch -import triton -from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager -from .fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager -from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager -from .operator import ( - FP8StaticPerHeadQuantMemOperator, - FP8StaticPerTensorQuantMemOperator, - LinearAttMemOperator, -) -from typing import Tuple, Any, List - -logger = init_logger(__name__) - - -class Qwen3NextMemManager(MemoryManager): - operator_class = LinearAttMemOperator - - def __init__( - self, - size, - dtype, - num_kv_heads, - head_dim, - full_att_layer_num, - linear_config: LinearAttCacheConfig, - always_copy=False, - mem_fraction=0.9, - ): - self.linear_config = linear_config - - super().__init__(size, dtype, num_kv_heads, head_dim, full_att_layer_num, always_copy, mem_fraction) - - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - return super().get_att_input_params(layer_index) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - super()._init_buffers(size, dtype, head_num, head_dim, layer_num) - # TODO 初始化线性 att 对应的部分 buffer. - self._init_linear_att_buffers() - return - - def _init_linear_att_buffers(self): - big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) - # 申请大页可能需要对应的资源, 多申请了两个linear att的状态，理论上这个状态 - # 永远不会被 alloc 申请到，只会在 cpu cache中，用于过渡和存储碎页情况下的 - # cpu cache 的页面拷贝。 - self.big_page_buffers = LinearAttCacheManager( - size=triton.cdiv(self.size, big_page_token_num) + 2, - linear_config=self.linear_config, - keep_num=2, - ) - - self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 - self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 - 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 - # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 - # CPU tensor 的 storage 原地迁到共享内存，使本进程大页 state cache 原本 - # pinned(cudaHostAlloc) 的内存退化为普通 shm mmap，之后 Triton kernel 携带该指针 - # 启动会报 "Pointer argument cannot be accessed from Triton (cpu tensor?)"。 - # 跨进程消费方并不使用 cpu 侧大页 state cache，序列化期间临时剔除以保住 pinned。 - big_page_buffers = self.big_page_buffers - self.big_page_buffers = None - try: - return super().write_to_shm(req_manager) - finally: - self.big_page_buffers = big_page_buffers - - def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: - kv_move_buffer = super().alloc_paged_kv_move_buffer(page_num, page_size) - self.att_state_page_helper = self._create_att_state_page_helper() - self.att_state_page_helper.assert_page_size() - return kv_move_buffer - - def _create_att_state_page_helper(self): - return Qwen3NextLinearAttPageHelper(self) - - def write_mem_to_page_kv_move_buffer( - self, - mem_indexes, - page_index: int, - dp_index: int, - mem_managers, - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - if page_kind == "kv": - return super().write_mem_to_page_kv_move_buffer( - mem_indexes=mem_indexes, - page_index=page_index, - dp_index=dp_index, - mem_managers=mem_managers, - dp_world_size=dp_world_size, - page_kind=page_kind, - req_idx=req_idx, - ) - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - assert req_idx is not None - helper = self.att_state_page_helper - dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) - helper.write_req_to_page(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) - return - - def read_page_kv_move_buffer_to_mem( - self, - mem_indexes, - page_index: int, - dp_index: int, - mem_managers, - dp_world_size: int, - page_kind: str = "kv", - req_idx: int = None, - ): - if page_kind == "kv": - return super().read_page_kv_move_buffer_to_mem( - mem_indexes=mem_indexes, - page_index=page_index, - dp_index=dp_index, - mem_managers=mem_managers, - dp_world_size=dp_world_size, - page_kind=page_kind, - req_idx=req_idx, - ) - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - assert req_idx is not None - helper = self.att_state_page_helper - dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) - helper.read_page_to_req(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) - return - - -class _FP8StaticPerHeadQuantLinearAttMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - full_att_layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - FP8StaticPerHeadQuantMemOperator.copy_kv_to_mem_manager(self, full_att_layer_index, mem_index, kv) - - -class _FP8StaticPerTensorQuantLinearAttMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - full_att_layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) - FP8StaticPerTensorQuantMemOperator.copy_kv_to_mem_manager(self, full_att_layer_index, mem_index, kv) - - -class FP8StaticPerHeadQuantQwen3NextMemManager(Qwen3NextMemManager, FP8StaticPerHeadQuantMemManager): - operator_class = _FP8StaticPerHeadQuantLinearAttMemOperator - - -class FP8StaticPerTensorQuantQwen3NextMemManager(Qwen3NextMemManager, FP8StaticPerTensorQuantMemManager): - operator_class = _FP8StaticPerTensorQuantLinearAttMemOperator - - -class Qwen3NextLinearAttPageHelper: - def __init__(self, mem_manager: "Qwen3NextMemManager"): - self.mem_manager = mem_manager - self.linear_config = mem_manager.linear_config - self.req_to_conv_state = mem_manager.req_to_conv_state - self.req_to_ssm_state = mem_manager.req_to_ssm_state - self.global_linear_k_heads = self.linear_config.global_linear_k_heads - self.global_linear_v_heads = self.linear_config.global_linear_v_heads - - self.global_q_dim = self.global_linear_k_heads * self.linear_config.head_linear_k_dim - self.global_k_dim = self.global_q_dim - self.global_v_heads = self.global_linear_v_heads - self.global_v_dim = self.global_v_heads * self.linear_config.head_linear_v_dim - # conv state follows mixed_qkv layout: [q, k, v], each as a flat channel block. - self.conv_shape = ( - self.linear_config.linear_layer_num, - self.global_q_dim + self.global_k_dim + self.global_v_dim, - self.linear_config.conv_kernel_size - 1, - ) - self.ssm_shape = ( - self.linear_config.linear_layer_num, - self.global_v_heads, - self.linear_config.head_linear_k_dim, - self.linear_config.head_linear_v_dim, - ) - - self.conv_nbytes = ( - self.conv_shape[0] * self.conv_shape[1] * self.conv_shape[2] * self.req_to_conv_state.buffer.element_size() - ) - ssm_alignment = self.req_to_ssm_state.buffer.element_size() - # 做一下字节对齐，防止切出来的不对齐，导致一些操作的性能下降。 - self.ssm_offset = ((self.conv_nbytes + ssm_alignment - 1) // ssm_alignment) * ssm_alignment - self.ssm_nbytes = ( - self.ssm_shape[0] - * self.ssm_shape[1] - * self.ssm_shape[2] - * self.ssm_shape[3] - * self.req_to_ssm_state.buffer.element_size() - ) - self.state_nbytes = self.ssm_offset + self.ssm_nbytes - - def assert_page_size(self): - kv_move_buffer = self.mem_manager.kv_move_buffer - page_nbytes = kv_move_buffer[0].numel() * kv_move_buffer.element_size() - assert ( - page_nbytes >= self.state_nbytes - ), f"nixl kv move page bytes {page_nbytes} is smaller than global linear att state bytes {self.state_nbytes}" - return - - def get_dp_mems(self, mem_managers: List["Qwen3NextMemManager"], dp_index: int, dp_world_size: int): - dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] - assert len(dp_mems) == dp_world_size - for mem in dp_mems: - assert hasattr(mem, "req_to_conv_state") and hasattr(mem, "req_to_ssm_state") - assert mem.linear_config.linear_layer_num == self.linear_config.linear_layer_num - assert mem.linear_config.conv_kernel_size == self.linear_config.conv_kernel_size - assert mem.linear_config.head_linear_k_dim == self.linear_config.head_linear_k_dim - assert mem.linear_config.head_linear_v_dim == self.linear_config.head_linear_v_dim - assert mem.linear_config.num_linear_k_heads == self.linear_config.num_linear_k_heads - assert mem.linear_config.num_linear_v_heads == self.linear_config.num_linear_v_heads - return dp_mems - - def view_page_to_linear_att_state(self, page_index: int): - page_bytes = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) - conv_page = page_bytes[0 : self.conv_nbytes].view(self.req_to_conv_state.buffer.dtype).view(self.conv_shape) - ssm_page = ( - page_bytes[self.ssm_offset : self.ssm_offset + self.ssm_nbytes] - .view(self.req_to_ssm_state.buffer.dtype) - .view(self.ssm_shape) - ) - return conv_page, ssm_page - - def write_req_to_page( - self, - page_index: int, - req_idx: int, - dp_mems: List["Qwen3NextMemManager"], - ): - conv_page, ssm_page = self.view_page_to_linear_att_state(page_index) - for tp_index, mem in enumerate(dp_mems): - self._write_one_rank(mem, tp_index, req_idx, conv_page, ssm_page) - return - - def read_page_to_req( - self, - page_index: int, - req_idx: int, - dp_mems: List["Qwen3NextMemManager"], - ): - conv_page, ssm_page = self.view_page_to_linear_att_state(page_index) - for tp_index, mem in enumerate(dp_mems): - self._read_one_rank(mem, tp_index, req_idx, conv_page, ssm_page) - return - - def _get_req_state_indexes(self, req_idx: int): - mtp_size = get_env_start_args().mtp_step + 1 - # Conv is one widened slot per request; SSM keeps the historical S+1 block layout. - return req_idx, req_idx * mtp_size - - def _write_one_rank( - self, - mem: "Qwen3NextMemManager", - tp_index: int, - req_idx: int, - conv_page: torch.Tensor, - ssm_page: torch.Tensor, - ): - conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx) - conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]] - ssm_state = mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...] - self._copy_conv_state_to_page(conv_state, conv_page, mem, tp_index) - self._copy_ssm_state_to_page(ssm_state, ssm_page, mem, tp_index) - return - - def _copy_conv_state_to_page( - self, - conv_state: torch.Tensor, - conv_page: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - local_q_heads = mem.linear_config.num_linear_k_heads - local_v_heads = mem.linear_config.num_linear_v_heads - head_k_dim = mem.linear_config.head_linear_k_dim - head_v_dim = mem.linear_config.head_linear_v_dim - - local_q_state = conv_state[:, 0 : local_q_heads * head_k_dim, :] - local_k_state = conv_state[:, local_q_heads * head_k_dim : 2 * local_q_heads * head_k_dim, :] - local_v_state = conv_state[:, 2 * local_q_heads * head_k_dim :, :] - global_q_page = conv_page[:, 0 : self.global_q_dim, :] - global_k_page = conv_page[:, self.global_q_dim : self.global_q_dim + self.global_k_dim, :] - global_v_page = conv_page[:, self.global_q_dim + self.global_k_dim :, :] - - qk_head_slice = self._get_head_slice( - tp_index, local_q_heads, self.global_linear_k_heads, mem.linear_config.tp_world_size, is_write=True - ) - if qk_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = qk_head_slice - local_dim_start = local_head_start * head_k_dim - local_dim_end = local_head_end * head_k_dim - global_dim_start = global_head_start * head_k_dim - global_dim_end = global_head_end * head_k_dim - global_q_page[:, global_dim_start:global_dim_end, :].copy_( - local_q_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - global_k_page[:, global_dim_start:global_dim_end, :].copy_( - local_k_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - - v_head_slice = self._get_head_slice( - tp_index, local_v_heads, self.global_linear_v_heads, mem.linear_config.tp_world_size, is_write=True - ) - if v_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = v_head_slice - local_dim_start = local_head_start * head_v_dim - local_dim_end = local_head_end * head_v_dim - global_dim_start = global_head_start * head_v_dim - global_dim_end = global_head_end * head_v_dim - global_v_page[:, global_dim_start:global_dim_end, :].copy_( - local_v_state[:, local_dim_start:local_dim_end, :], non_blocking=True - ) - return - - def _copy_ssm_state_to_page( - self, - ssm_state: torch.Tensor, - ssm_page: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - head_slice = self._get_head_slice( - tp_index, - mem.linear_config.num_linear_v_heads, - self.global_linear_v_heads, - mem.linear_config.tp_world_size, - is_write=True, - ) - if head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = head_slice - ssm_page[:, global_head_start:global_head_end, :, :].copy_( - ssm_state[:, local_head_start:local_head_end, :, :], - non_blocking=True, - ) - return - - def _get_head_slice( - self, - tp_index: int, - local_heads: int, - global_heads: int, - tp_world_size: int, - is_write: bool, - ): - if local_heads == 0 or global_heads == 0: - return None - total_local_heads = local_heads * tp_world_size - repeat_count = max(1, total_local_heads // global_heads) - if is_write and repeat_count > 1 and tp_index % repeat_count != 0: - return None - unique_tp_index = tp_index // repeat_count - global_head_start = unique_tp_index * local_heads - global_head_end = min(global_head_start + local_heads, global_heads) - local_head_start = 0 - local_head_end = global_head_end - global_head_start - if local_head_end <= local_head_start: - return None - return local_head_start, local_head_end, global_head_start, global_head_end - - def _copy_page_to_conv_state( - self, - conv_page: torch.Tensor, - conv_state: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - local_q_heads = mem.linear_config.num_linear_k_heads - local_v_heads = mem.linear_config.num_linear_v_heads - head_k_dim = mem.linear_config.head_linear_k_dim - head_v_dim = mem.linear_config.head_linear_v_dim - - local_q_state = conv_state[:, 0 : local_q_heads * head_k_dim, :] - local_k_state = conv_state[:, local_q_heads * head_k_dim : 2 * local_q_heads * head_k_dim, :] - local_v_state = conv_state[:, 2 * local_q_heads * head_k_dim :, :] - global_q_page = conv_page[:, 0 : self.global_q_dim, :] - global_k_page = conv_page[:, self.global_q_dim : self.global_q_dim + self.global_k_dim, :] - global_v_page = conv_page[:, self.global_q_dim + self.global_k_dim :, :] - - qk_head_slice = self._get_head_slice( - tp_index, local_q_heads, self.global_linear_k_heads, mem.linear_config.tp_world_size, is_write=False - ) - if qk_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = qk_head_slice - local_dim_start = local_head_start * head_k_dim - local_dim_end = local_head_end * head_k_dim - global_dim_start = global_head_start * head_k_dim - global_dim_end = global_head_end * head_k_dim - local_q_state[:, local_dim_start:local_dim_end, :].copy_( - global_q_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - local_k_state[:, local_dim_start:local_dim_end, :].copy_( - global_k_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - - v_head_slice = self._get_head_slice( - tp_index, local_v_heads, self.global_linear_v_heads, mem.linear_config.tp_world_size, is_write=False - ) - if v_head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = v_head_slice - local_dim_start = local_head_start * head_v_dim - local_dim_end = local_head_end * head_v_dim - global_dim_start = global_head_start * head_v_dim - global_dim_end = global_head_end * head_v_dim - local_v_state[:, local_dim_start:local_dim_end, :].copy_( - global_v_page[:, global_dim_start:global_dim_end, :], non_blocking=True - ) - return - - def _copy_page_to_ssm_state( - self, - ssm_page: torch.Tensor, - ssm_state: torch.Tensor, - mem: "Qwen3NextMemManager", - tp_index: int, - ): - head_slice = self._get_head_slice( - tp_index, - mem.linear_config.num_linear_v_heads, - self.global_linear_v_heads, - mem.linear_config.tp_world_size, - is_write=False, - ) - if head_slice is not None: - local_head_start, local_head_end, global_head_start, global_head_end = head_slice - ssm_state[:, local_head_start:local_head_end, :, :].copy_( - ssm_page[:, global_head_start:global_head_end, :, :], - non_blocking=True, - ) - return - - def _read_one_rank( - self, - mem: "Qwen3NextMemManager", - tp_index: int, - req_idx: int, - conv_page: torch.Tensor, - ssm_page: torch.Tensor, - ): - conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx) - conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]] - ssm_state = mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...] - self._copy_page_to_conv_state(conv_page, conv_state, mem, tp_index) - self._copy_page_to_ssm_state(ssm_page, ssm_state, mem, tp_index) - return diff --git a/lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py b/lightllm/common/kv_cache_mem_manager/static_fp8_mem_manager.py old mode 100755 new mode 100644 similarity index 61% rename from lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py rename to lightllm/common/kv_cache_mem_manager/static_fp8_mem_manager.py index 2d5e8a1e5e..93525a88d4 --- a/lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/static_fp8_mem_manager.py @@ -1,14 +1,12 @@ import os import json import torch -import torch.distributed as dist -from typing import Tuple, Any -from lightllm.utils.config_utils import get_model_architectures from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args from lightllm.utils.dist_utils import get_dp_world_size, get_current_rank_in_dp from .mem_manager import MemoryManager -from .operator import FP8StaticPerHeadQuantMemOperator +from .operator import FP8StaticPerHeadQuantMemOperator, FP8StaticPerTensorQuantMemOperator +from lightllm.utils.config_utils import get_model_architectures logger = init_logger(__name__) @@ -92,7 +90,55 @@ def _load_and_check_config(self): assert (get_dp_world_size() * self.head_num) % q_cfg["num_head"] == 0 return cfg, scales - def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: - k = self.kv_buffer[layer_index][:, : self.head_num, :] - v = self.kv_buffer[layer_index][:, self.head_num :, :] - return k, v + +class FP8StaticPerTensorQuantMemManager(MemoryManager): + + operator_class = FP8StaticPerTensorQuantMemOperator + + def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): + # 这里用uint8存储量化后的kv，方便兼容各种torch算子。fp8量化目前采用离线方案，kv_buffer不存储scale + super().__init__(size, torch.uint8, head_num, head_dim, layer_num, always_copy, mem_fraction) + + self.qmax = torch.finfo(torch.float8_e4m3fn).max + self.qmin = torch.finfo(torch.float8_e4m3fn).min + self.scales = None + + if get_env_start_args().kv_quant_calibration_config_path is not None: + logger.info( + f"kv_quant_calibration_config_path {get_env_start_args().kv_quant_calibration_config_path} is set, " + "will load kv quant calibration config" + ) + cfg = self._load_and_check_config() + self.scales = torch.tensor(cfg["scales"], dtype=torch.float32, device="cuda").view(cfg["scales_shape"]) + else: + self.scales = torch.ones((self.kv_buffer.shape[0], 2), dtype=torch.float32, device="cuda") + + self.cpu_scales = self.scales.detach().cpu().numpy() + return + + def _load_and_check_config(self): + if os.path.exists(get_env_start_args().kv_quant_calibration_config_path): + with open(get_env_start_args().kv_quant_calibration_config_path, "r") as f: + cfg = json.load(f) + + if cfg["qmin"] != self.qmin: + raise ValueError(f"qmin {cfg['qmin']} in config not match torch.float8_e4m3fn.min {self.qmin}") + if cfg["qmax"] != self.qmax: + raise ValueError(f"qmax {cfg['qmax']} in config not match torch.float8_e4m3fn.max {self.qmax}") + model_arch = get_model_architectures(get_env_start_args().model_dir) + if cfg["architectures"] != model_arch: + raise ValueError( + f"architectures {cfg['architectures']} in config " f"not match current model_arch {model_arch}" + ) + if cfg["num_layers"] != self.layer_num: + raise ValueError( + f"num_layers {cfg['num_layers']} in config " f"not match current layer_num {self.layer_num}" + ) + assert ( + cfg["quant_type"] == "per_tensor" + ), f"quant type {cfg['quant_type']} in config not match per-tensor backend" + return cfg + else: + raise FileNotFoundError( + f"kv_quant_calibration_config {get_env_start_args().kv_quant_calibration_config_path} not found" + ) diff --git a/lightllm/common/kv_page_transfer.py b/lightllm/common/kv_page_transfer.py new file mode 100644 index 0000000000..ae294dbac6 --- /dev/null +++ b/lightllm/common/kv_page_transfer.py @@ -0,0 +1,46 @@ +import torch + +from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io, mla_page_io + + +class KVPageTransfer: + def __init__(self, view): + if view.page_shape is None: + raise ValueError("PD transfer requires a KV page layout") + self.view = view + + def alloc_paged_kv_move_buffer(self, page_num, page_size): + self.kv_move_buffer = torch.empty( + (page_num, page_size, *self.view.page_shape), dtype=self.view.page_dtype, device="cuda" + ) + self.index_buffers = [torch.empty(page_size, dtype=torch.int64, pin_memory=True) for _ in range(page_num)] + return self.kv_move_buffer + + def write_mem_to_page_kv_move_buffer(self, mem_indexes, page_index, dp_index, kv_views, dp_world_size): + self._copy(mem_indexes, page_index, dp_index, kv_views, dp_world_size, "write") + + def read_page_kv_move_buffer_to_mem(self, mem_indexes, page_index, dp_index, kv_views, dp_world_size): + self._copy(mem_indexes, page_index, dp_index, kv_views, dp_world_size, "read") + + def _copy(self, indexes, page_index, dp_index, views, dp_world_size, mode): + pinned = self.index_buffers[page_index][: len(indexes)] + pinned.numpy()[:] = indexes + indexes_gpu = pinned.cuda(non_blocking=True) + page = self.kv_move_buffer[page_index] + ranks = views[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] + assert len(ranks) == dp_world_size + if self.view.replicated: + for view in ranks[:1] if mode == "write" else ranks: + mla_page_io(mem_indexes=indexes_gpu, page_tensor=page, kv_buffer=view.kv_buffer, mode=mode) + else: + repeat_count = dp_world_size * self.view.kv_buffer.shape[2] // self.kv_move_buffer.shape[3] + for rank, view in enumerate(ranks): + if mode == "read" or rank % repeat_count == 0: + page_io( + mem_indexes=indexes_gpu, + page_tensor=page, + kv_buffer=view.kv_buffer, + tp_index=rank, + tp_world_size=dp_world_size, + mode=mode, + ) diff --git a/lightllm/common/kv_transfer_resources.py b/lightllm/common/kv_transfer_resources.py new file mode 100644 index 0000000000..a5d6fe4d0c --- /dev/null +++ b/lightllm/common/kv_transfer_resources.py @@ -0,0 +1,83 @@ +from dataclasses import dataclass +from multiprocessing.reduction import ForkingPickler +from typing import TYPE_CHECKING, Optional + +import torch +import torch.multiprocessing as mp +from filelock import FileLock + +from lightllm.utils.device_utils import kv_trans_use_p2p +from lightllm.utils.dist_utils import get_current_rank_in_node, get_node_world_size +from lightllm.utils.log_utils import init_logger +from lightllm.utils.shm_utils import create_or_link_shm, get_service_shm_name + +if TYPE_CHECKING: + from lightllm.common.linear_state_transfer import LinearAttentionStateTransfer + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class KVTransferView: + kv_buffer: torch.Tensor + page_shape: Optional[tuple] + page_dtype: torch.dtype + replicated: bool + + @classmethod + def from_mem_manager(cls, mem_manager, *, include_page_layout=False): + return cls( + mem_manager.kv_buffer, + mem_manager.get_paged_kv_move_buffer_shape(1, 1)[2:] if include_page_layout else None, + mem_manager.dtype, + mem_manager.kv_is_replicated, + ) + + +@dataclass +class KVTransferResources: + # Export GPU handles only; checkpoint pools must stay pinned in the inference process. + kv: KVTransferView + req_to_token_indexes: torch.Tensor + state_transfer: Optional["LinearAttentionStateTransfer"] = None + + def write_to_shm(self): + if kv_trans_use_p2p(): + from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor + + mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__ + + lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") + with lock: + node_world_size = get_node_world_size() + shm_name = f"mem_manager_{get_current_rank_in_node()}" + obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] + # IPC ref-counter offsets can change the serialized size between consumers. + obj_size = max(map(len, obj_bytes_array)) + shm = create_or_link_shm( + name=shm_name, expected_size=obj_size * (node_world_size * 2) + 4 + 4, force_mode="create" + ) + logger.info(f"create shm {shm.name} size {shm.size} for mem manger shared buffer") + shm.buf[0:4] = (node_world_size * 2).to_bytes(4, "little") + shm.buf[4:8] = obj_size.to_bytes(4, "little") + start_index = 8 + for obj_bytes in obj_bytes_array: + shm.buf[start_index : start_index + obj_size] = obj_bytes.ljust(obj_size, b"\0") + start_index += obj_size + + @staticmethod + def loads_from_shm(rank_in_node: int): + shm_name = f"mem_manager_{rank_in_node}" + lock = FileLock(f"/tmp/{get_service_shm_name('mem_manager_lock')}") + logger.info(f"get memmanager from shm {shm_name}") + with lock: + shm = create_or_link_shm(name=shm_name, expected_size=-1, force_mode="link") + left_num = int.from_bytes(shm.buf[0:4], "little") + obj_size = int.from_bytes(shm.buf[4:8], "little") + assert left_num > 0 + end_index = 8 + left_num * obj_size + start_index = 8 + (left_num - 1) * obj_size + obj_bytes = shm.buf[start_index:end_index].tobytes() + shm.buf[0:4] = (left_num - 1).to_bytes(4, byteorder="little") + shm.close() + return ForkingPickler.loads(obj_bytes) diff --git a/lightllm/common/linear_state.py b/lightllm/common/linear_state.py new file mode 100644 index 0000000000..7db9520895 --- /dev/null +++ b/lightllm/common/linear_state.py @@ -0,0 +1,108 @@ +import torch + +from lightllm.common.state_cache_manager import LayerCache + + +class LinearStateManager: + """Manage request-local convolution and SSM state with native MTP slots.""" + + def __init__(self, max_request_num, linear_config, args): + self.linear_config = linear_config + self.mtp_step = args.mtp_step + self.mtp_state_index = ( + torch.zeros((max_request_num + 1,), dtype=torch.int32, device="cuda") if self.mtp_step > 0 else None + ) + self.conv_state = LayerCache( + size=max_request_num + 1, + dtype=linear_config.conv_state_dtype, + shape=linear_config.get_mtp_conv_state_shape(self.mtp_step), + layer_num=linear_config.linear_layer_num, + device="cuda", + ) + self.ssm_state = LayerCache( + size=(max_request_num + 1) * (self.mtp_step + 1), + dtype=linear_config.ssm_state_dtype, + shape=linear_config.get_ssm_state_shape(), + layer_num=linear_config.linear_layer_num, + device="cuda", + ) + + def reset(self, req_idx): + ssm_start = req_idx * (self.mtp_step + 1) + self.conv_state.buffer[:, req_idx].zero_() + self.ssm_state.buffer[:, ssm_start : ssm_start + self.mtp_step + 1].zero_() + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx] = 0 + + def get_layer_buffers(self, layer_idx_in_all): + assert 0 <= layer_idx_in_all < self.linear_config.all_layer_num + layer_idx = layer_idx_in_all - layer_idx_in_all // self.linear_config.full_attention_interval + return self.conv_state.buffer[layer_idx], self.ssm_state.buffer[layer_idx] + + def copy_state_to(self, req_idx, conv_dst, ssm_dst): + conv_width = self.linear_config.get_conv_state_shape()[-1] + conv_dst.copy_(self.conv_state.buffer[:, req_idx, ..., :conv_width], non_blocking=True) + ssm_dst.copy_(self.ssm_state.buffer[:, req_idx * (self.mtp_step + 1)], non_blocking=True) + + def save_checkpoints(self, b_req_idx, req_indexes, buffer_indexes, state_cache_manager, checkpoint_lens=None): + from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer + + buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) + copy_linear_att_state_to_kv_buffer( + b_req_idx=b_req_idx, + big_page_buffer_ids=buffer_indexes, + gpu_conv_state=self.conv_state.buffer, + gpu_ssm_state=self.ssm_state.buffer, + cpu_kv_conv_state=state_cache_manager.conv_state_cache.buffer, + cpu_kv_ssm_state=state_cache_manager.ssm_state_cache.buffer, + mtp_step=self.mtp_step, + ) + + def restore(self, req_idx, conv_state, ssm_state): + self.conv_state.buffer[:, req_idx, ..., : conv_state.shape[-1]] = conv_state + self.ssm_state.buffer[:, req_idx * (self.mtp_step + 1)] = ssm_state + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx] = 0 + + def commit_accepted_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width): + from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update + + linear_att_mtp_state_index_update( + req_to_mtp_state_index=self.mtp_state_index, + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + accepted_index=accepted_index, + verify_width=verify_width, + ) + + capture_intermediate_checkpoints = False + + def create_checkpoint_pool(self, size, keep_num=0): + from lightllm.common.state_cache_manager import LinearAttCacheManager + + return LinearAttCacheManager(size, self.linear_config, keep_num=keep_num) + + def create_cpu_transfer(self, mem_manager, big_page_buffers): + from lightllm.common.hybrid_cache_transfer import HybridCpuCacheTransfer + + return HybridCpuCacheTransfer(mem_manager.kv_buffer, self.linear_config, big_page_buffers) + + def create_state_transfer(self, indexer_state_manager=None): + from lightllm.common.linear_state_transfer import LinearAttentionStateTransfer + + return LinearAttentionStateTransfer(self, indexer_state_manager) + + def save_state(self, req_idx, buffer_idx, state_cache_manager, checkpoint_len): + self.copy_state_to(req_idx, *state_cache_manager.get_state_cache(buffer_idx)) + + def restore_checkpoint(self, req_idx, restored_len, state_cache_manager, buffer_idx): + self.restore(req_idx, *state_cache_manager.get_state_cache(buffer_idx)) + + def restore_peer_state(self, req_idx, restored_len, peer_state, peer_req_idx): + conv_width = self.linear_config.get_conv_state_shape()[-1] + self.restore( + req_idx, + peer_state.req_to_conv_state.buffer[:, peer_req_idx, ..., :conv_width], + peer_state.req_to_ssm_state.buffer[:, peer_req_idx * (peer_state.mtp_step + 1)], + ) diff --git a/lightllm/common/linear_state_transfer.py b/lightllm/common/linear_state_transfer.py new file mode 100644 index 0000000000..1b8dfacb1c --- /dev/null +++ b/lightllm/common/linear_state_transfer.py @@ -0,0 +1,332 @@ +import torch +from typing import List + + +class LinearAttentionStateTransfer: + """Pack live Conv/SSM and optional indexer state into TP-independent PD pages.""" + + def __init__(self, linear_state_manager, indexer_state_manager=None): + self.linear_config = linear_state_manager.linear_config + self.req_to_conv_state = linear_state_manager.conv_state + self.req_to_ssm_state = linear_state_manager.ssm_state + self.mtp_step = linear_state_manager.mtp_step + self.req_to_indexer_tail = None if indexer_state_manager is None else indexer_state_manager.state + self.kv_move_buffer = None + self.global_linear_k_heads = self.linear_config.global_linear_k_heads + self.global_linear_v_heads = self.linear_config.global_linear_v_heads + + self.global_q_dim = self.global_linear_k_heads * self.linear_config.head_linear_k_dim + self.global_k_dim = self.global_q_dim + self.global_v_heads = self.global_linear_v_heads + self.global_v_dim = self.global_v_heads * self.linear_config.head_linear_v_dim + # conv state follows mixed_qkv layout: [q, k, v], each as a flat channel block. + self.conv_shape = ( + self.linear_config.linear_layer_num, + self.global_q_dim + self.global_k_dim + self.global_v_dim, + self.linear_config.conv_kernel_size - 1, + ) + self.ssm_shape = ( + self.linear_config.linear_layer_num, + self.global_v_heads, + self.linear_config.head_linear_k_dim, + self.linear_config.head_linear_v_dim, + ) + + self.conv_nbytes = ( + self.conv_shape[0] * self.conv_shape[1] * self.conv_shape[2] * self.req_to_conv_state.buffer.element_size() + ) + ssm_alignment = self.req_to_ssm_state.buffer.element_size() + # 做一下字节对齐，防止切出来的不对齐，导致一些操作的性能下降。 + self.ssm_offset = ((self.conv_nbytes + ssm_alignment - 1) // ssm_alignment) * ssm_alignment + self.ssm_nbytes = ( + self.ssm_shape[0] + * self.ssm_shape[1] + * self.ssm_shape[2] + * self.ssm_shape[3] + * self.req_to_ssm_state.buffer.element_size() + ) + self.state_nbytes = self.ssm_offset + self.ssm_nbytes + if self.req_to_indexer_tail is not None: + tail = self.req_to_indexer_tail.buffer + self.tail_dtype = tail.dtype + self.tail_shape = (self.linear_config.get_full_att_kv_layer_num_with_draft_model(), *tail.shape[2:]) + self.tail_offset = ((self.state_nbytes + 15) // 16) * 16 + self.tail_nbytes = self.tail_shape[0] * self.tail_shape[1] * self.tail_shape[2] * tail.element_size() + self.state_nbytes = self.tail_offset + self.tail_nbytes + + def bind_page_buffer(self, kv_move_buffer): + self.kv_move_buffer = kv_move_buffer + self.assert_page_size() + + def _view_page_to_tail(self, page_index): + page_bytes = self.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) + return ( + page_bytes[self.tail_offset : self.tail_offset + self.tail_nbytes] + .view(self.tail_dtype) + .view(self.tail_shape) + ) + + def assert_page_size(self): + kv_move_buffer = self.kv_move_buffer + page_nbytes = kv_move_buffer[0].numel() * kv_move_buffer.element_size() + assert ( + page_nbytes >= self.state_nbytes + ), f"nixl kv move page bytes {page_nbytes} is smaller than global linear att state bytes {self.state_nbytes}" + return + + def get_dp_states(self, state_transfers: List["LinearAttentionStateTransfer"], dp_index: int, dp_world_size: int): + dp_states = state_transfers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)] + assert len(dp_states) == dp_world_size + for state in dp_states: + assert hasattr(state, "req_to_conv_state") and hasattr(state, "req_to_ssm_state") + assert state.linear_config.linear_layer_num == self.linear_config.linear_layer_num + assert state.linear_config.conv_kernel_size == self.linear_config.conv_kernel_size + assert state.linear_config.head_linear_k_dim == self.linear_config.head_linear_k_dim + assert state.linear_config.head_linear_v_dim == self.linear_config.head_linear_v_dim + assert state.linear_config.num_linear_k_heads == self.linear_config.num_linear_k_heads + assert state.linear_config.num_linear_v_heads == self.linear_config.num_linear_v_heads + assert state.mtp_step == self.mtp_step + assert (state.req_to_indexer_tail is None) == (self.req_to_indexer_tail is None) + return dp_states + + def view_page_to_linear_att_state(self, page_index: int): + page_bytes = self.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) + conv_page = page_bytes[0 : self.conv_nbytes].view(self.req_to_conv_state.buffer.dtype).view(self.conv_shape) + ssm_page = ( + page_bytes[self.ssm_offset : self.ssm_offset + self.ssm_nbytes] + .view(self.req_to_ssm_state.buffer.dtype) + .view(self.ssm_shape) + ) + return conv_page, ssm_page + + def write_req_to_page( + self, + page_index: int, + req_idx: int, + dp_states: List["LinearAttentionStateTransfer"], + ): + assert req_idx is not None + conv_page, ssm_page = self.view_page_to_linear_att_state(page_index) + for tp_index, state in enumerate(dp_states): + self._write_one_rank(state, tp_index, req_idx, conv_page, ssm_page) + if self.req_to_indexer_tail is not None: + self._view_page_to_tail(page_index).copy_( + dp_states[0].req_to_indexer_tail.buffer[:, req_idx], non_blocking=True + ) + return + + def read_page_to_req( + self, + page_index: int, + req_idx: int, + dp_states: List["LinearAttentionStateTransfer"], + ): + assert req_idx is not None + conv_page, ssm_page = self.view_page_to_linear_att_state(page_index) + for tp_index, state in enumerate(dp_states): + self._read_one_rank(state, tp_index, req_idx, conv_page, ssm_page) + if self.req_to_indexer_tail is not None: + tail = self._view_page_to_tail(page_index) + for state in dp_states: + state.req_to_indexer_tail.buffer[:, req_idx].copy_(tail, non_blocking=True) + return + + def _get_req_state_indexes(self, req_idx: int): + mtp_size = self.mtp_step + 1 + # Conv is one widened slot per request; SSM keeps the historical S+1 block layout. + return req_idx, req_idx * mtp_size + + def _write_one_rank( + self, + state: "LinearAttentionStateTransfer", + tp_index: int, + req_idx: int, + conv_page: torch.Tensor, + ssm_page: torch.Tensor, + ): + conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx) + conv_state = state.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]] + ssm_state = state.req_to_ssm_state.buffer[:, ssm_req_idx, ...] + self._copy_conv_state_to_page(conv_state, conv_page, state, tp_index) + self._copy_ssm_state_to_page(ssm_state, ssm_page, state, tp_index) + return + + def _copy_conv_state_to_page( + self, + conv_state: torch.Tensor, + conv_page: torch.Tensor, + state: "LinearAttentionStateTransfer", + tp_index: int, + ): + local_q_heads = state.linear_config.num_linear_k_heads + local_v_heads = state.linear_config.num_linear_v_heads + head_k_dim = state.linear_config.head_linear_k_dim + head_v_dim = state.linear_config.head_linear_v_dim + + local_q_state = conv_state[:, 0 : local_q_heads * head_k_dim, :] + local_k_state = conv_state[:, local_q_heads * head_k_dim : 2 * local_q_heads * head_k_dim, :] + local_v_state = conv_state[:, 2 * local_q_heads * head_k_dim :, :] + global_q_page = conv_page[:, 0 : self.global_q_dim, :] + global_k_page = conv_page[:, self.global_q_dim : self.global_q_dim + self.global_k_dim, :] + global_v_page = conv_page[:, self.global_q_dim + self.global_k_dim :, :] + + qk_head_slice = self._get_head_slice( + tp_index, local_q_heads, self.global_linear_k_heads, state.linear_config.tp_world_size, is_write=True + ) + if qk_head_slice is not None: + local_head_start, local_head_end, global_head_start, global_head_end = qk_head_slice + local_dim_start = local_head_start * head_k_dim + local_dim_end = local_head_end * head_k_dim + global_dim_start = global_head_start * head_k_dim + global_dim_end = global_head_end * head_k_dim + global_q_page[:, global_dim_start:global_dim_end, :].copy_( + local_q_state[:, local_dim_start:local_dim_end, :], non_blocking=True + ) + global_k_page[:, global_dim_start:global_dim_end, :].copy_( + local_k_state[:, local_dim_start:local_dim_end, :], non_blocking=True + ) + + v_head_slice = self._get_head_slice( + tp_index, local_v_heads, self.global_linear_v_heads, state.linear_config.tp_world_size, is_write=True + ) + if v_head_slice is not None: + local_head_start, local_head_end, global_head_start, global_head_end = v_head_slice + local_dim_start = local_head_start * head_v_dim + local_dim_end = local_head_end * head_v_dim + global_dim_start = global_head_start * head_v_dim + global_dim_end = global_head_end * head_v_dim + global_v_page[:, global_dim_start:global_dim_end, :].copy_( + local_v_state[:, local_dim_start:local_dim_end, :], non_blocking=True + ) + return + + def _copy_ssm_state_to_page( + self, + ssm_state: torch.Tensor, + ssm_page: torch.Tensor, + state: "LinearAttentionStateTransfer", + tp_index: int, + ): + head_slice = self._get_head_slice( + tp_index, + state.linear_config.num_linear_v_heads, + self.global_linear_v_heads, + state.linear_config.tp_world_size, + is_write=True, + ) + if head_slice is not None: + local_head_start, local_head_end, global_head_start, global_head_end = head_slice + ssm_page[:, global_head_start:global_head_end, :, :].copy_( + ssm_state[:, local_head_start:local_head_end, :, :], + non_blocking=True, + ) + return + + def _get_head_slice( + self, + tp_index: int, + local_heads: int, + global_heads: int, + tp_world_size: int, + is_write: bool, + ): + if local_heads == 0 or global_heads == 0: + return None + total_local_heads = local_heads * tp_world_size + repeat_count = max(1, total_local_heads // global_heads) + if is_write and repeat_count > 1 and tp_index % repeat_count != 0: + return None + unique_tp_index = tp_index // repeat_count + global_head_start = unique_tp_index * local_heads + global_head_end = min(global_head_start + local_heads, global_heads) + local_head_start = 0 + local_head_end = global_head_end - global_head_start + if local_head_end <= local_head_start: + return None + return local_head_start, local_head_end, global_head_start, global_head_end + + def _copy_page_to_conv_state( + self, + conv_page: torch.Tensor, + conv_state: torch.Tensor, + state: "LinearAttentionStateTransfer", + tp_index: int, + ): + local_q_heads = state.linear_config.num_linear_k_heads + local_v_heads = state.linear_config.num_linear_v_heads + head_k_dim = state.linear_config.head_linear_k_dim + head_v_dim = state.linear_config.head_linear_v_dim + + local_q_state = conv_state[:, 0 : local_q_heads * head_k_dim, :] + local_k_state = conv_state[:, local_q_heads * head_k_dim : 2 * local_q_heads * head_k_dim, :] + local_v_state = conv_state[:, 2 * local_q_heads * head_k_dim :, :] + global_q_page = conv_page[:, 0 : self.global_q_dim, :] + global_k_page = conv_page[:, self.global_q_dim : self.global_q_dim + self.global_k_dim, :] + global_v_page = conv_page[:, self.global_q_dim + self.global_k_dim :, :] + + qk_head_slice = self._get_head_slice( + tp_index, local_q_heads, self.global_linear_k_heads, state.linear_config.tp_world_size, is_write=False + ) + if qk_head_slice is not None: + local_head_start, local_head_end, global_head_start, global_head_end = qk_head_slice + local_dim_start = local_head_start * head_k_dim + local_dim_end = local_head_end * head_k_dim + global_dim_start = global_head_start * head_k_dim + global_dim_end = global_head_end * head_k_dim + local_q_state[:, local_dim_start:local_dim_end, :].copy_( + global_q_page[:, global_dim_start:global_dim_end, :], non_blocking=True + ) + local_k_state[:, local_dim_start:local_dim_end, :].copy_( + global_k_page[:, global_dim_start:global_dim_end, :], non_blocking=True + ) + + v_head_slice = self._get_head_slice( + tp_index, local_v_heads, self.global_linear_v_heads, state.linear_config.tp_world_size, is_write=False + ) + if v_head_slice is not None: + local_head_start, local_head_end, global_head_start, global_head_end = v_head_slice + local_dim_start = local_head_start * head_v_dim + local_dim_end = local_head_end * head_v_dim + global_dim_start = global_head_start * head_v_dim + global_dim_end = global_head_end * head_v_dim + local_v_state[:, local_dim_start:local_dim_end, :].copy_( + global_v_page[:, global_dim_start:global_dim_end, :], non_blocking=True + ) + return + + def _copy_page_to_ssm_state( + self, + ssm_page: torch.Tensor, + ssm_state: torch.Tensor, + state: "LinearAttentionStateTransfer", + tp_index: int, + ): + head_slice = self._get_head_slice( + tp_index, + state.linear_config.num_linear_v_heads, + self.global_linear_v_heads, + state.linear_config.tp_world_size, + is_write=False, + ) + if head_slice is not None: + local_head_start, local_head_end, global_head_start, global_head_end = head_slice + ssm_state[:, local_head_start:local_head_end, :, :].copy_( + ssm_page[:, global_head_start:global_head_end, :, :], + non_blocking=True, + ) + return + + def _read_one_rank( + self, + state: "LinearAttentionStateTransfer", + tp_index: int, + req_idx: int, + conv_page: torch.Tensor, + ssm_page: torch.Tensor, + ): + conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx) + conv_state = state.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]] + ssm_state = state.req_to_ssm_state.buffer[:, ssm_req_idx, ...] + self._copy_page_to_conv_state(conv_page, conv_state, state, tp_index) + self._copy_page_to_ssm_state(ssm_page, ssm_state, state, tp_index) + return diff --git a/lightllm/common/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index ca1f64d5b8..711915bd04 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -1,13 +1,7 @@ from .base import ReqManager -from .linear_att import ReqManagerForMamba -from .glm5_next import Glm5NextReqManager -from .hybrid_base import HybridAttentionReqManager from .req_sampling_params import ReqSamplingParamsManager __all__ = [ "ReqManager", - "HybridAttentionReqManager", - "ReqManagerForMamba", - "Glm5NextReqManager", "ReqSamplingParamsManager", ] diff --git a/lightllm/common/req_manager/glm5_next.py b/lightllm/common/req_manager/glm5_next.py deleted file mode 100644 index b4814afbf7..0000000000 --- a/lightllm/common/req_manager/glm5_next.py +++ /dev/null @@ -1,34 +0,0 @@ -from lightllm.common.state_cache_manager import LayerCache -from lightllm.utils.envs_utils import get_env_start_args - -from .linear_att import ReqManagerForMamba - - -class Glm5NextReqManager(ReqManagerForMamba): - """KDA runtime state and the NSA indexer's incomplete four-token pool.""" - - def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config): - # Both checkpoint sizes are multiples of this hash page. A restored - # prefix therefore has no incomplete K-pool to serialize or replay. - assert ( - get_env_start_args().linear_att_hash_page_size % linear_config.index_kpool == 0 - ), "GLM K-pool requires cache pages aligned to index_kpool" - super().__init__(max_request_num, max_sequence_length, mem_manager, linear_config) - self.req_to_indexer_tail = LayerCache( - size=max_request_num + 1, - dtype=linear_config.full_att_dtype, - shape=(linear_config.index_kpool + self.mtp_step, 2 * linear_config.index_head_dim), - layer_num=linear_config.get_full_att_kv_layer_num_with_draft_model(), - device="cuda", - ) - - def get_indexer_tail_buffer(self, layer_index): - return self.req_to_indexer_tail.buffer[self.linear_config.get_full_att_kv_layer_index(layer_index)] - - def init_hybrid_attention_state(self, req): - super().init_hybrid_attention_state(req) - self.req_to_indexer_tail.buffer[:, req.req_idx].zero_() - - def restore_state(self, req, state_cache_manager, buffer_idx): - super().restore_state(req, state_cache_manager, buffer_idx) - self.req_to_indexer_tail.buffer[:, req.req_idx].zero_() diff --git a/lightllm/common/req_manager/hybrid_base.py b/lightllm/common/req_manager/hybrid_base.py deleted file mode 100644 index 948313116d..0000000000 --- a/lightllm/common/req_manager/hybrid_base.py +++ /dev/null @@ -1,70 +0,0 @@ -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, List - -import torch - -from .base import ReqManager - - -if TYPE_CHECKING: - from lightllm.server.router.model_infer.infer_batch import InferReq - - -class HybridAttentionReqManager(ReqManager, ABC): - """混合 attention 的请求运行态与大小页 checkpoint 管理接口。 - - 大小页沿同一虚拟 token 索引空间匹配前缀，full attention KV 保持 token 粒度存储。 - linear/sliding-window 状态在大页边界及请求可缓存尾部的小页边界保存 checkpoint， - 缓存命中后，再将相应 checkpoint 恢复到请求运行态。 - - 请求的 GPU 计算状态由具体实现管理；small_page_buffers 持有 CPU 小页快照； - big_page_buffers 引用 mem_manager 的 CPU 大页快照，与 full KV 容量一起创建、调整。 - 大小页不包含额外的 GPU 运行态，命中后恢复到请求状态，不重新分配 checkpoint 池。 - 公共缓存流程负责 checkpoint 槽位分配、边界、匹配与淘汰；本接口负责运行态与 checkpoint 保存/恢复。 - """ - - def __init__(self, max_request_num, max_sequence_length, mem_manager): - super().__init__(max_request_num, max_sequence_length, mem_manager) - self.small_page_buffers = None - - @property - def big_page_buffers(self): - return self.mem_manager.big_page_buffers - - @abstractmethod - def create_small_page_cache_manager(self, size: int): - """创建并持有 CPU 小页池，返回同一池供 radix 使用；size 是槽位数，不是 token 数。""" - - @abstractmethod - def init_hybrid_attention_state(self, req: "InferReq"): - """无前缀缓存命中时，初始化已分配请求槽位的 GPU 运行态。""" - - def restore_big_page_state(self, big_page_buffer_idx: int, req: "InferReq"): - """将指定大页槽位的 CPU checkpoint 恢复到请求 GPU 运行态。""" - self.restore_state(req, self.big_page_buffers, big_page_buffer_idx) - - def restore_small_page_state(self, req: "InferReq"): - """将 req.shared_kv_node 对应的小页 checkpoint 恢复到请求 GPU 运行态。""" - self.restore_state(req, self.small_page_buffers, req.shared_kv_node.small_page_buffer_idx) - - @abstractmethod - def restore_state(self, req: "InferReq", state_cache_manager, buffer_idx: int): - """CPU checkpoint → 请求 GPU 运行态；大小页共用，不负责前缀匹配或 full KV 索引恢复。""" - - def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - """批量保存请求 GPU 运行态到已分配的大页槽位，buffer_indexes 中的 -1 表示跳过。 - - b_req_idx 与 req_indexes 分别为同一批请求的 GPU 索引张量和 CPU 索引列表。 - 默认逐请求保存，模型可覆盖为批量拷贝算子。 - """ - for req_idx, buffer_idx in zip(req_indexes, buffer_indexes): - if buffer_idx != -1: - self.save_state(req_idx, buffer_idx, self.big_page_buffers) - - @abstractmethod - def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager): - """请求 GPU 运行态 → 指定 CPU checkpoint 槽位；大小页共用，调用方负责分配槽位。""" - - def update_mtp_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width): - """接受推测 token 后更新运行态位置；需要调整状态索引的模型覆写。""" - return diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py deleted file mode 100644 index ccab8a06f6..0000000000 --- a/lightllm/common/req_manager/linear_att.py +++ /dev/null @@ -1,121 +0,0 @@ -from typing import TYPE_CHECKING, List - -import torch - -from lightllm.common.state_cache_manager import LayerCache, LinearAttCacheConfig, LinearAttCacheManager -from lightllm.utils.envs_utils import get_env_start_args - -from .hybrid_base import HybridAttentionReqManager - - -if TYPE_CHECKING: - from lightllm.server.router.model_infer.infer_batch import InferReq - - -class ReqManagerForMamba(HybridAttentionReqManager): - def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_config: LinearAttCacheConfig): - super().__init__(max_request_num, max_sequence_length, mem_manager) - self.mtp_step = get_env_start_args().mtp_step - # 因为在mtp的推理中，需要标记每个请求对应的mtp index状态(conv state 和 ssm state)，在mtp对应序列中 - # 的真实位置，所以需要需要一个标记来记录，不然算子无法找到真实的处理起点。 - self.req_to_mtp_state_index = ( - torch.zeros((max_request_num + 1,), dtype=torch.int32, device="cuda") if self.mtp_step > 0 else None - ) - # 突然想到， 在linear att 开启mtp的模式中，现在的prefill linear att 算子默认是从0的位置读取信息进行操作 - # 所以不能支持 prefill decode mixed 操作了，因为一个decode过的请求，重新用prefill 算子跑，会出现读错linear - # 状态位置的问题。导致bug, 在这里加个断言，以后可以支持上 TODO - if self.mtp_step > 0: - assert get_env_start_args().enable_prefill_decode_mixed is False - - self.big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) - self.linear_config = linear_config - - self.req_to_conv_state = LayerCache( - size=(max_request_num + 1), - dtype=self.linear_config.conv_state_dtype, - shape=self.linear_config.get_mtp_conv_state_shape(mtp_step=self.mtp_step), - layer_num=self.linear_config.linear_layer_num, - device="cuda", - ) - self.req_to_ssm_state = LayerCache( - size=(max_request_num + 1) * (self.mtp_step + 1), - dtype=self.linear_config.ssm_state_dtype, - shape=self.linear_config.get_ssm_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cuda", - ) - return - - def init_hybrid_attention_state(self, req: "InferReq"): - conv_index = req.req_idx - ssm_start = req.req_idx * (self.mtp_step + 1) - self.req_to_conv_state.buffer[:, conv_index, ...].fill_(0) - # #17: zero the FULL (mtp_step + 1)-row SSM block, not just canonical row +0, so a future - # first-step verify reading offset>0 after fresh init never hits a never-written row (NaN). - self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + (self.mtp_step + 1), ...].fill_(0) - if self.req_to_mtp_state_index is not None: - self.req_to_mtp_state_index[req.req_idx] = 0 - return - - def create_small_page_cache_manager(self, size: int): - self.small_page_buffers = LinearAttCacheManager(size=size, linear_config=self.linear_config) - return self.small_page_buffers - - def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): - from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer - - buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) - state_cache_manager = self.big_page_buffers - copy_linear_att_state_to_kv_buffer( - b_req_idx=b_req_idx, - big_page_buffer_ids=buffer_indexes, - gpu_conv_state=self.req_to_conv_state.buffer, - gpu_ssm_state=self.req_to_ssm_state.buffer, - cpu_kv_conv_state=state_cache_manager.conv_state_cache.buffer, - cpu_kv_ssm_state=state_cache_manager.ssm_state_cache.buffer, - mtp_step=self.mtp_step, - ) - return - - def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearAttCacheManager): - # checkpoint 只保存标准 conv 窗口和请求的基准 SSM 状态，不包含 MTP 扩展运行态。 - conv_cache_width = self.linear_config.get_conv_state_shape()[-1] - gpu_conv_state = self.req_to_conv_state.buffer[:, req_idx, ..., :conv_cache_width] - gpu_ssm_state = self.req_to_ssm_state.buffer[:, req_idx * (self.mtp_step + 1), ...] - dst_conv_state, dst_ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) - dst_conv_state.copy_(gpu_conv_state, non_blocking=True) - dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) - - def get_mamba_cache(self, layer_idx_in_all: int): - assert ( - 0 <= layer_idx_in_all < self.linear_config.all_layer_num - ), f"invalid transformer layer index {layer_idx_in_all}" - layer_idx_in_linear = layer_idx_in_all - (layer_idx_in_all // self.linear_config.full_attention_interval) - conv_states = self.req_to_conv_state.buffer[layer_idx_in_linear] - ssm_states = self.req_to_ssm_state.buffer[layer_idx_in_linear] - return conv_states, ssm_states - - def update_mtp_state(self, b_req_mtp_start_loc, b_req_idx, b_mtp_index, accepted_index, verify_width): - from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update - - linear_att_mtp_state_index_update( - req_to_mtp_state_index=self.req_to_mtp_state_index, - b_req_mtp_start_loc=b_req_mtp_start_loc, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - accepted_index=accepted_index, - verify_width=verify_width, - ) - - def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheManager, buffer_idx: int): - conv_state, ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) - conv_dest = req.req_idx - ssm_dest = req.req_idx * (self.mtp_step + 1) - conv_cache_width = conv_state.shape[-1] - self.req_to_conv_state.buffer[:, conv_dest, ..., :conv_cache_width] = conv_state - self.req_to_ssm_state.buffer[:, ssm_dest, ...] = ssm_state - if self.req_to_mtp_state_index is not None: - self.req_to_mtp_state_index[req.req_idx] = 0 - return diff --git a/lightllm/common/state_cache_manager/base.py b/lightllm/common/state_cache_manager/base.py index aa131fe784..6839dd5681 100644 --- a/lightllm/common/state_cache_manager/base.py +++ b/lightllm/common/state_cache_manager/base.py @@ -20,6 +20,7 @@ def __init__(self, size: int, keep_num: int = 0): self.keep_num = keep_num assert 0 <= keep_num <= size, f"invalid keep_num {keep_num} for size {size}" self.free_list = collections.deque(range(size - keep_num)) + self.checkpoint_lens = [None] * size @abstractmethod def get_state_cache(self, buffer_idx: int): @@ -40,6 +41,7 @@ def free_state_cache(self, free_indexes: List[int]): assert 0 <= idx < alloc_upper_bound, ( f"free index {idx} out of alloc range [0, {alloc_upper_bound}), " f"reserved tail num {self.keep_num}" ) + self.checkpoint_lens[idx] = None self.free_list.extend(free_indexes) assert ( len(self.free_list) <= alloc_upper_bound @@ -48,6 +50,9 @@ def free_state_cache(self, free_indexes: List[int]): def get_free_cache_num(self): return len(self.free_list) + def has_checkpoint(self, buffer_idx, length): + return buffer_idx is not None and length is not None and self.checkpoint_lens[buffer_idx] == length + def get_used_cache_num(self): # Preserve the existing accounting: reserved slots count as used. return self.size - len(self.free_list) @@ -55,3 +60,4 @@ def get_used_cache_num(self): def clear_to_init_state(self): """重置空闲槽位；子类同时清零自身的 checkpoint buffer。""" self.free_list = collections.deque(range(self.size - self.keep_num)) + self.checkpoint_lens = [None] * self.size diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py index 35a37dcd28..64d84c449c 100644 --- a/lightllm/common/state_cache_manager/linear_att.py +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -59,6 +59,8 @@ def get_full_att_kv_layer_index(self, layer_index: int) -> int: """Map a global target/draft layer index to the packed full-attention cache.""" layer_index = int(layer_index) + if 0 <= layer_index < self.all_layer_num and (layer_index + 1) % self.full_attention_interval != 0: + raise ValueError(f"layer {layer_index} has no full-attention KV storage") if layer_index >= self.all_layer_num: kv_layer_index = layer_index - self.linear_layer_num else: @@ -69,6 +71,14 @@ def get_full_att_kv_layer_index(self, layer_index: int) -> int: ) return kv_layer_index + def get_full_att_kv_layer_mapping(self): + return tuple( + self.get_full_att_kv_layer_index(index) + if index >= self.all_layer_num or (index + 1) % self.full_attention_interval == 0 + else None + for index in range(self.all_layer_num + self.draft_full_att_kv_layer_num) + ) + def get_conv_state_shape(self): # Base committed sliding-window state, without speculative MTP tail. return (self.get_conv_dim(), self.conv_kernel_size - 1) @@ -91,7 +101,7 @@ def get_cpu_cache_big_page_bytes(self): b = self.get_cpu_cache_conv_bytes() c = self.get_cpu_cache_ssm_bytes() - return triton.cdiv(a + b + c, 16) * 16 + return triton.cdiv(a + b + c, 16) * 16 + 16 def get_cpu_cache_full_att_bytes(self): big_page_token_num = ( @@ -125,9 +135,13 @@ def load_from_args() -> "LinearAttCacheConfig": except: pass - n_layer = llm_config["num_hidden_layers"] + return LinearAttCacheConfig.from_model_config(llm_config, args) + + @staticmethod + def from_model_config(llm_config, args) -> "LinearAttCacheConfig": + n_layer = llm_config["n_layer"] if "n_layer" in llm_config else llm_config["num_hidden_layers"] - tp_world_size = get_env_start_args().tp // get_env_start_args().dp + tp_world_size = args.tp // args.dp full_att_dtype = ( torch.uint8 if args.llm_kv_type in {"fp8kv_sph", "fp8kv_spt"} else get_torch_dtype(args.data_type) ) diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 26a7ecd81f..289a0921c0 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -4,7 +4,7 @@ from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class +from lightllm.common.kv_cache_mem_manager import MlaMemoryManager from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args, get_added_mtp_kv_layer_num from lightllm.distributed.communication_op import dist_group_manager @@ -57,8 +57,12 @@ def _init_custom(self): def _verify_params(self): return super()._verify_params() + @classmethod + def get_mem_manager_class(cls): + return MlaMemoryManager + def _init_mem_manager(self): - manager_class = select_mem_manager_class() + manager_class = self.get_mem_manager_class() self.mem_manager = manager_class( self.max_total_token_num, diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index 415e8826f8..2a0ce95859 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -13,6 +13,15 @@ class Deepseek3_2TpPartModel(Deepseek2TpPartModel): # infer class transformer_layer_infer_class = Deepseek3_2TransformerLayerInfer + @classmethod + def get_mem_manager_class(cls): + from lightllm.common.kv_cache_mem_manager import IndexedMlaMemoryManager, FP8IndexedMlaMemoryManager + from lightllm.utils.envs_utils import get_env_start_args + + return ( + FP8IndexedMlaMemoryManager if get_env_start_args().llm_kv_type == "fp8kv_dsa" else IndexedMlaMemoryManager + ) + def _init_att_backend(self): self.prefill_att_backend = get_nsa_prefill_att_backend_class(index=0)(model=self) self.decode_att_backend = get_nsa_decode_att_backend_class(index=0)(model=self) diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index cf13ab83de..05316391ea 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -36,7 +36,7 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig k = layer_weight.k_norm_(layer_weight.wk_proj_.mm(hidden_states), eps=self.eps) gate = layer_weight.index_kpool_compress_gate.mm(hidden_states) raw = torch.cat((k, gate), -1) - tail = infer_state.req_manager.get_indexer_tail_buffer(self.layer_idx) + tail = infer_state.indexer_state_manager.get_layer_buffer(self.layer_idx) packed_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx) compress_pools( raw=raw, diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index cc8e86613e..f8c25ed18d 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -5,11 +5,14 @@ import triton from lightllm.common.build_utils import repair_config +from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager -from lightllm.common.req_manager import Glm5NextReqManager +from lightllm.common.cache_coordinator import HybridCacheCoordinator +from lightllm.common.linear_state import LinearStateManager +from lightllm.common.indexer_state import IndexerStateManager from lightllm.common.state_cache_manager import Glm5NextCacheConfig from lightllm.distributed.communication_op import dist_group_manager from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer @@ -58,25 +61,35 @@ def _init_some_value(self): self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - def _init_req_manager(self): + def _init_cache_coordinator(self): self.linear_config = Glm5NextCacheConfig.from_model_config(self.config, self.args) - self.req_manager = Glm5NextReqManager( - self.max_req_num, - max(self.batch_max_tokens or 0, self.max_seq_length or 0), - None, - linear_config=self.linear_config, + assert ( + self.args.linear_att_hash_page_size % self.linear_config.index_kpool == 0 + ), "GLM K-pool requires cache pages aligned to index_kpool" + if self.args.mtp_step > 0: + assert self.args.enable_prefill_decode_mixed is False + self.cache_coordinator = HybridCacheCoordinator( + linear_state_manager=LinearStateManager(self.max_req_num, self.linear_config, self.args), + indexer_state_manager=IndexerStateManager(self.max_req_num, self.linear_config, self.args.mtp_step), ) + @classmethod + def get_mem_manager_class(cls): + kv_type = get_env_start_args().llm_kv_type + if kv_type != "None": + raise ValueError(f"GLM5Next does not support llm_kv_type={kv_type}") + return Glm5NextMemManager + def _init_mem_manager(self): - self.mem_manager = Glm5NextMemManager( + self.mem_manager = self.get_mem_manager_class()( size=self.max_total_token_num, dtype=self.data_type, - num_kv_heads=1, + head_num=1, head_dim=self.linear_config.full_att_head_dim, - full_att_layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), - linear_config=self.linear_config, + layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), mem_fraction=self.mem_fraction, ) + self.mem_manager.set_kv_layer_mapping(self.linear_config.get_full_att_kv_layer_mapping()) def _init_att_backend(self): self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) diff --git a/lightllm/models/qwen3_5_dflash/model.py b/lightllm/models/qwen3_5_dflash/model.py index 422ccf61cf..4dc26a1e96 100644 --- a/lightllm/models/qwen3_5_dflash/model.py +++ b/lightllm/models/qwen3_5_dflash/model.py @@ -32,7 +32,7 @@ def _init_mem_manager(self): target_mem_manager = self.main_model.mem_manager draft_kv_shape = (self.config["num_key_value_heads"], self.config["head_dim"]) target_kv_shape = ( - target_mem_manager.linear_config.full_att_all_num_kv_heads, + self.main_model.linear_config.full_att_all_num_kv_heads, target_mem_manager.head_dim, ) assert draft_kv_shape == target_kv_shape, ( diff --git a/lightllm/models/qwen3_5_dspark/model.py b/lightllm/models/qwen3_5_dspark/model.py index 896949d1cf..27b4766925 100644 --- a/lightllm/models/qwen3_5_dspark/model.py +++ b/lightllm/models/qwen3_5_dspark/model.py @@ -29,7 +29,7 @@ def _init_mem_manager(self): target_mem_manager = self.main_model.mem_manager draft_kv_shape = (self.config["num_key_value_heads"], self.config["head_dim"]) target_kv_shape = ( - target_mem_manager.linear_config.full_att_all_num_kv_heads, + self.main_model.linear_config.full_att_all_num_kv_heads, target_mem_manager.head_dim, ) assert draft_kv_shape == target_kv_shape, ( diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 92d68c9fd2..39f9c42919 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -8,7 +8,6 @@ from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo from lightllm.utils.log_utils import init_logger from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor -from lightllm.common.kv_cache_mem_manager import Qwen3NextMemManager from lightllm.common.basemodel.attention.base_att import AttControl from typing import Tuple from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ @@ -245,7 +244,6 @@ def context_attention_forward( if not self.is_linear_attention_layer: return super().context_attention_forward(input_embdings, infer_state, layer_weight) - assert isinstance(infer_state.mem_manager, Qwen3NextMemManager) mixed_qkvzba = self._linear_in_proj(input_embdings, layer_weight) if torch.cuda.is_current_stream_capturing(): @@ -281,7 +279,6 @@ def token_attention_forward( if not self.is_linear_attention_layer: return super().token_attention_forward(input_embdings, infer_state, layer_weight) - assert isinstance(infer_state.mem_manager, Qwen3NextMemManager) mixed_qkvzba = self._linear_in_proj(input_embdings, layer_weight) core_attn_out, z = infer_state.decode_att_state1.decode_att( q=None, diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index 08233444eb..3051e073fc 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -11,14 +11,14 @@ ) from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args -from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( - FP8StaticPerHeadQuantQwen3NextMemManager, - FP8StaticPerTensorQuantQwen3NextMemManager, - Qwen3NextMemManager, +from lightllm.common.kv_cache_mem_manager import ( + FP8StaticPerHeadQuantMemManager, + FP8StaticPerTensorQuantMemManager, + MemoryManager, ) -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.common.req_manager import ReqManagerForMamba +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.common.cache_coordinator import HybridCacheCoordinator +from lightllm.common.linear_state import LinearStateManager from lightllm.common.state_cache_manager import LinearAttCacheConfig logger = init_logger(__name__) @@ -57,64 +57,37 @@ def _init_config(self): super()._init_config() self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - def _init_mem_manager(self): - assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 - start_args: StartArgs = get_env_start_args() - ssm_dtype_dict = {"bfloat16": torch.bfloat16, "float32": torch.float32} - draft_full_att_kv_layer_num = get_added_mtp_kv_layer_num() - full_att_dtype = torch.uint8 if start_args.llm_kv_type in {"fp8kv_sph", "fp8kv_spt"} else self.data_type - self.linear_config = LinearAttCacheConfig( - tp_world_size=self.tp_world_size_, - full_att_all_num_kv_heads=self.config["num_key_value_heads"], - full_att_dtype=full_att_dtype, - full_att_num_kv_heads=self.num_kv_heads, - full_att_head_dim=self.config["head_dim"], - global_linear_k_heads=self.config["linear_num_key_heads"], - global_linear_v_heads=self.config["linear_num_value_heads"], - num_linear_k_heads=max(1, self.config["linear_num_key_heads"] // self.tp_world_size_), - num_linear_v_heads=max(1, self.config["linear_num_value_heads"] // self.tp_world_size_), - head_linear_k_dim=self.config["linear_key_head_dim"], - head_linear_v_dim=self.config["linear_value_head_dim"], - conv_kernel_size=self.config["linear_conv_kernel_dim"], - linear_layer_num=self.config["n_layer"] - - (self.config["n_layer"] // self.config["full_attention_interval"]), - conv_state_dtype=self.data_type, - ssm_state_dtype=ssm_dtype_dict[start_args.linear_att_ssm_data_type], - full_attention_interval=self.config["full_attention_interval"], - all_layer_num=self.config["n_layer"], - draft_full_att_kv_layer_num=draft_full_att_kv_layer_num, - ) - + @classmethod + def get_mem_manager_class(cls): + start_args = get_env_start_args() if start_args.llm_kv_type == "None": - mem_manager_class = Qwen3NextMemManager + mem_manager_class = MemoryManager elif start_args.llm_kv_type == "fp8kv_sph": - mem_manager_class = FP8StaticPerHeadQuantQwen3NextMemManager + mem_manager_class = FP8StaticPerHeadQuantMemManager elif start_args.llm_kv_type == "fp8kv_spt": - mem_manager_class = FP8StaticPerTensorQuantQwen3NextMemManager + mem_manager_class = FP8StaticPerTensorQuantMemManager else: raise ValueError(f"Qwen3Next does not support llm_kv_type={start_args.llm_kv_type}") - self.mem_manager = mem_manager_class( + return mem_manager_class + + def _init_mem_manager(self): + assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 + self.mem_manager = self.get_mem_manager_class()( size=self.max_total_token_num, dtype=self.data_type, - num_kv_heads=self.num_kv_heads, + head_num=self.num_kv_heads, head_dim=self.config["head_dim"], - full_att_layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), - linear_config=self.linear_config, + layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), mem_fraction=self.mem_fraction, ) - - def _init_req_manager(self): - create_max_seq_len = 0 - - if self.batch_max_tokens is not None: - create_max_seq_len = max(create_max_seq_len, self.batch_max_tokens) - if self.max_seq_length is not None: - create_max_seq_len = max(create_max_seq_len, self.max_seq_length) - - self.req_manager = ReqManagerForMamba( - self.max_req_num, create_max_seq_len, None, linear_config=LinearAttCacheConfig.load_from_args() - ) - return + self.mem_manager.set_kv_layer_mapping(self.linear_config.get_full_att_kv_layer_mapping()) + + def _init_cache_coordinator(self): + if self.args.mtp_step > 0: + assert self.args.enable_prefill_decode_mixed is False + self.linear_config = LinearAttCacheConfig.from_model_config(self.config, self.args) + state = LinearStateManager(self.max_req_num, self.linear_config, self.args) + self.cache_coordinator = HybridCacheCoordinator(linear_state_manager=state) def _init_att_backend1(self): # MTP draft models only keep full-attention layers and reuse the main diff --git a/lightllm/server/multi_level_kv_cache/disk_cache_worker.py b/lightllm/server/multi_level_kv_cache/disk_cache_worker.py index 542ddbd877..43a70ab210 100644 --- a/lightllm/server/multi_level_kv_cache/disk_cache_worker.py +++ b/lightllm/server/multi_level_kv_cache/disk_cache_worker.py @@ -6,7 +6,8 @@ from typing import List, Optional import torch -from lightllm.utils.envs_utils import get_unique_server_name +from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args +from lightllm.utils.config_utils import is_hybrid_att_model from lightllm.utils.log_utils import init_logger from .cpu_cache_client import CpuKvCacheClient @@ -43,7 +44,7 @@ def __init__( self._pages_all_idle = False assert disk_cache_storage_size > 0 - storage_size = int(disk_cache_storage_size * (1024 ** 3)) + storage_size = int(disk_cache_storage_size * (1024**3)) # 为了发挥afs性能，2TiB以内使用8个shard；容量超过2TiB后，每增加1TiB增加4个shard。 num_shard = max(8, math.ceil(disk_cache_storage_size / 1024) * 4) num_worker = 24 @@ -54,7 +55,8 @@ def __init__( if not cache_dir: cache_dir = os.path.join(tempfile.gettempdir(), f"lightllm_disk_cache_{get_unique_server_name()}") os.makedirs(cache_dir, exist_ok=True) - cache_file = os.path.join(cache_dir, "cache_file") + filename = "cache_file_hybrid_state_v1" if is_hybrid_att_model(get_env_start_args().model_dir) else "cache_file" + cache_file = os.path.join(cache_dir, filename) self.max_concurrent_write_tasks = max_concurrent_write_tasks self._page_major_tensor = self._prepare_tensor(cpu_cache_client.cpu_kv_cache_tensor) diff --git a/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py b/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py index f8516bce2c..995df2153e 100644 --- a/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py @@ -119,6 +119,7 @@ def __init__( big_page_num: int, kv_cache_mem_manager=None, small_page_buffers=None, + big_page_buffers=None, ): from lightllm.common.kv_cache_mem_manager import MemoryManager @@ -132,7 +133,7 @@ def __init__( self.mem_manager: MemoryManager = kv_cache_mem_manager - self.big_page_buffers: StateCacheManager = self.mem_manager.big_page_buffers + self.big_page_buffers: StateCacheManager = big_page_buffers self._key_dtype = torch.int64 self._value_dtype = torch.int64 @@ -247,6 +248,13 @@ def _insert_helper( # 提前释放 len_to_big_page_id 对应的buffer资源 new_big_page_buffer_id = len_to_big_page_id.pop(child.node_prefix_total_len, None) if new_big_page_buffer_id is not None: + if not self.big_page_buffers.has_checkpoint( + child.big_page_buffer_idx, child.node_prefix_total_len + ) and self.big_page_buffers.has_checkpoint(new_big_page_buffer_id, child.node_prefix_total_len): + child.big_page_buffer_idx, new_big_page_buffer_id = ( + new_big_page_buffer_id, + child.big_page_buffer_idx, + ) # 因为节点已经存在，所以无法插入，但是要释放对应的buffer_id 节点 self.big_page_buffers.free_state_cache([new_big_page_buffer_id]) @@ -364,12 +372,10 @@ def match_prefix( ans_node_list=ans_node_list, update_refs=update_refs, ) - if len(ans_node_list) == 0: - return None, 0, None - # 判定真正可以用的匹配节点。 ans_node_list = self._trim_unusable_match_tail(ans_node_list, mm_input_spans=mm_input_spans) if len(ans_node_list) == 0: + self.dec_node_ref_counter(self.root_node) return None, 0, None ans_node = ans_node_list[-1] @@ -442,7 +448,12 @@ def _trim_unusable_match_tail( """ removed_list = [] for node in reversed(nodes): - has_state = node.is_big_page_node() or node.small_page_buffer_idx is not None + if node.is_big_page_node(): + has_state = self.big_page_buffers.has_checkpoint(node.big_page_buffer_idx, node.node_prefix_total_len) + else: + has_state = self.small_page_buffers.has_checkpoint( + node.small_page_buffer_idx, node.node_prefix_total_len + ) outside_mm = clamp_mm_cache_len(node.node_prefix_total_len, mm_input_spans) == node.node_prefix_total_len if has_state and outside_mm: break @@ -478,7 +489,7 @@ def flush_cache(self): return def deref_to_first_big_page_node( - self, node: HybridAttPagedTreeNode, mm_input_spans: Sequence[Tuple[int, int]] = () + self, node: HybridAttPagedTreeNode, mm_input_spans: Sequence[Tuple[int, int]] = (), require_state=True ) -> Optional[HybridAttPagedTreeNode]: """回退到最近的大页；用于恢复 state 时，传入图片区间以跳过图内边界。""" assert not node.is_big_page_node() @@ -486,7 +497,11 @@ def deref_to_first_big_page_node( while iter_node is not self.root_node: node_end = iter_node.node_prefix_total_len outside_mm = clamp_mm_cache_len(node_end, mm_input_spans) == node_end - if iter_node.is_big_page_node() and outside_mm: + if ( + iter_node.is_big_page_node() + and outside_mm + and (not require_state or self.big_page_buffers.has_checkpoint(iter_node.big_page_buffer_idx, node_end)) + ): break self._discard_node(iter_node) @@ -500,6 +515,7 @@ def deref_to_first_big_page_node( iter_node = iter_node.parent if iter_node is self.root_node: + self.dec_node_ref_counter(self.root_node) return None else: return iter_node diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 35fb8acda4..25404a6516 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -8,7 +8,8 @@ from sortedcontainers import SortedDict from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union -from lightllm.common.req_manager import ReqManager, HybridAttentionReqManager +from lightllm.common.req_manager import ReqManager +from lightllm.common.cache_coordinator import CacheCoordinator 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 @@ -39,6 +40,7 @@ class InferenceContext: backend: Optional["ModeBackend"] = None req_manager: ReqManager = None # gpu 请求管理 + cache_coordinator: Optional[CacheCoordinator] = None radix_cache: Union[HybridAttPagedRadixCache, RadixCache] = None shm_req_manager: ShmReqManager = None # 共享内存请求对象管理 requests_mapping: Dict[int, "InferReq"] = None @@ -59,11 +61,16 @@ def register( radix_cache: Union[HybridAttPagedRadixCache, RadixCache], shm_req_manager: ShmReqManager, vocab_size: int, + cache_coordinator: CacheCoordinator, cache_placement_controller: Optional[CachePlacementController] = None, ): self.args = get_env_start_args() self.backend: ModeBackend = backend self.req_manager = req_manager + self.cache_coordinator = cache_coordinator + if cache_coordinator.mem_manager is not req_manager.mem_manager: + raise ValueError("Request manager and coordinator must share the memory manager") + cache_coordinator.validate_radix_cache(radix_cache) self.req_sampling_manager = self.req_manager.req_sampling_params_manager self.radix_cache = radix_cache self.shm_req_manager = shm_req_manager @@ -74,7 +81,7 @@ def register( self.vocab_size = vocab_size - self.is_hybrid_att_model = isinstance(self.req_manager, HybridAttentionReqManager) + self.is_hybrid_att_model = self.cache_coordinator.has_state_cache self.output_constraint_state_factory = self._init_output_constraint_state_factory() return @@ -152,158 +159,17 @@ 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"): - 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]) - elif CacheTier.GPU not in req.cache_tiers: - self._free_req_mem_without_radix_insert(free_token_index=free_token_index, req=req) - else: - if not self.is_hybrid_att_model: - self._full_att_free_req(free_token_index=free_token_index, req=req) - else: - self._hybrid_att_free_req(free_token_index=free_token_index, req=req) - assert len(req.hybrid_len_to_big_page_id) == 0 - 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"): - 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]) - - if self.is_hybrid_att_model: - # 释放请求尾部尚未移交给 radix cache 的 hybrid attention 小页状态。 - if req.tail_small_page_buffer_id is not None: - self.radix_cache.small_page_buffers.free_state_cache([req.tail_small_page_buffer_id]) - req.tail_small_page_buffer_id = None - # 释放请求执行期间申请、但不再插入 radix cache 的大页状态。 - 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() - - # 解除请求对已命中 GPU radix cache 前缀节点的引用。 - if req.shared_kv_node is not None: - self.radix_cache.dec_node_ref_counter(req.shared_kv_node) - req.shared_kv_node = None - return - - def _full_att_free_req(self, free_token_index: List, req: "InferReq"): - 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) - 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 - - 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 - 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) - ) - 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: - assert req.hybrid_cache_len <= req.cur_kv_len - - if req.cur_kv_len == 0 and req.hold_kv_len == 0: - return - - 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] + free_token_index.extend( + self.cache_coordinator.release_request_cache( + req, + self.req_manager, + self.radix_cache, + self.args, + publish_cache=self.radix_cache is not None and CacheTier.GPU in req.cache_tiers, ) - 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 - - 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 - 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 False, f"error state: cur_kv_len: {req.cur_kv_len}" + ) + req.cur_kv_len = req.hold_kv_len = 0 + req.shm_req.shm_cur_kv_len = 0 def _save_promptcache_kvbuffer(self): """ @@ -403,10 +269,7 @@ def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bo if alloc_token_num > can_alloc_token_num: break - if g_infer_context.is_hybrid_att_model: - req._hybrid_match_radix_cache() - else: - req._match_radix_cache() + req._match_radix_cache() assert req.paused is True req.paused = False @@ -424,58 +287,8 @@ def get_can_alloc_token_num(self): ) return self.req_manager.mem_manager.allocator.can_use_mem_size + radix_cache_unref_token_num - def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): - """Snapshot request-level attention state at big/small-page boundaries.""" - if not self.is_hybrid_att_model: - return - - # Request-state snapshot at a big-page boundary. - big_page_token_num = self.args.linear_att_hash_page_size * self.args.linear_att_page_block_num - big_page_buffer_ids = [] - for req in reqs: - # 保存时尚未更新 req.cur_kv_len；它是本轮 chunk 起点，cur_input_len 是终点。 - cur_input_len = req.get_chuncked_input_token_len() - first_page_end = (req.cur_kv_len // big_page_token_num + 1) * big_page_token_num - state_buffer_id = -1 - last_page_end = min(cur_input_len, req.hybrid_cache_len) - # 完整图片可能跨过多个大页结束点，例如页长 256、chunk [128, 768) 覆盖 256/512/768。 - # 这些页的 KV 都有效，仍需为 radix 插入和 CPU offload 分配对应槽位。 - for page_end in range(first_page_end, last_page_end + 1, big_page_token_num): - big_page_id = self.radix_cache.big_page_buffers.alloc_one_state_cache() - assert big_page_id is not None - assert page_end not in req.hybrid_len_to_big_page_id - req.hybrid_len_to_big_page_id[page_end] = big_page_id - # 当前运行态只对应 chunk 终点；图片内部页尾不保存 state，匹配时也不能在此恢复。 - if page_end == cur_input_len: - state_buffer_id = big_page_id - big_page_buffer_ids.append(state_buffer_id) - - assert len(b_req_idx) == len(big_page_buffer_ids) - if any(buffer_id != -1 for buffer_id in big_page_buffer_ids): - self.req_manager.save_big_page_states( - b_req_idx=b_req_idx, - req_indexes=[req.req_idx for req in reqs], - buffer_indexes=big_page_buffer_ids, - ) - - assert not self.args.disable_chunked_prefill, "chunked prefill must be enabled for hybrid attention models" - - # Request-state snapshot at the final small-page boundary. - for req in reqs: - # 判断本次prefill 完以后 kv 的长度是否到达 hybrid checkpoint 的存储边界。 - if req.get_chuncked_input_token_len() == req.hybrid_cache_len: - assert req.tail_small_page_buffer_id is None - if req.hybrid_cache_len % big_page_token_num != 0: - self.radix_cache.free_one_small_page_buffer() - req.tail_small_page_buffer_id = self.radix_cache.small_page_buffers.alloc_one_state_cache() - if req.tail_small_page_buffer_id is not None: - dst_buffer_idx = req.tail_small_page_buffer_id - self.req_manager.save_state( - req_idx=req.req_idx, - buffer_idx=dst_buffer_idx, - state_cache_manager=self.radix_cache.small_page_buffers, - ) - return + def save_cache_state(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): + self.cache_coordinator.save_state_cache(b_req_idx, reqs, self.radix_cache, self.args) g_infer_context = InferenceContext() @@ -606,10 +419,7 @@ def __init__( self.generator.manual_seed(self.sampling_param.shm_param.seed) if init_prefix_cache: - if g_infer_context.is_hybrid_att_model: - self._hybrid_match_radix_cache() - else: - self._match_radix_cache() + self._match_radix_cache() return def _init_all_state(self): @@ -651,155 +461,7 @@ def _init_all_state(self): return 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..." - enable_prompt_cache = (not self.sampling_param.disable_prompt_cache) and g_infer_context.radix_cache is not None - if enable_prompt_cache and self.get_cur_total_len() > 1 and self.cur_kv_len == 0: - input_token_ids = self.shm_req.shm_prompt_ids.arr[0 : self.get_cur_total_len()] - key = torch.tensor(input_token_ids, dtype=torch.int64, device="cpu") - 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, mm_input_spans=self.mm_input_spans - ) - if share_node is not None: - self.shared_kv_node = share_node - ready_cache_len = share_node.node_prefix_total_len - # 从 cpu 到 gpu 是流内阻塞操作 - g_infer_context.req_manager.req_to_token_indexs[self.req_idx, 0:ready_cache_len] = value_tensor - self.cur_kv_len = int(ready_cache_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换 - self.hold_kv_len = self.cur_kv_len - assert self.hold_kv_len % self.args.page_size == 0 - self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 - - self.shm_req.shm_cur_kv_len = self.cur_kv_len - return - - 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..." - enable_prompt_cache = (not self.sampling_param.disable_prompt_cache) and g_infer_context.radix_cache is not None - block_hashs = self.shm_req.hybrid_token_hash_list.get_all() - hash_page_size = self.args.linear_att_hash_page_size - match_tokens = min(len(block_hashs) * hash_page_size, self.get_cur_total_len() - 1) - match_tokens = max(0, match_tokens) - match_tokens = (match_tokens // hash_page_size) * hash_page_size - 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( - key, block_hashs=block_hashs, update_refs=True, mm_input_spans=self.mm_input_spans - ) - 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, - ) - else: - # 如果 大页本质是被启用的，则需要使用小页的匹配结果, 将小页的kv 复制到的新申请的kv位置，同时释放 - # 对应的小页对应的节点，递归找到对应最近的大叶节点进行返回,然后赋值到req.shared_node 对象上 - shared_kv_len = share_node.node_prefix_total_len - cur_big_page_tokens = (shared_kv_len // big_page_token_num) * big_page_token_num - need_tokens = shared_kv_len - cur_big_page_tokens - radix_cache = g_infer_context.radix_cache - if g_infer_context.get_can_alloc_token_num() > need_tokens: - # 有充足的token 容量时 - assert need_tokens % self.args.page_size == 0 - radix_cache.free_radix_cache_to_get_enough_token(need_token_num=need_tokens) - tail_mems = radix_cache.mem_manager.alloc(need_size=need_tokens) - g_infer_context.req_manager.req_to_token_indexs[ - self.req_idx, 0:cur_big_page_tokens - ] = value_tensor[0:cur_big_page_tokens] - g_infer_context.req_manager.req_to_token_indexs[ - self.req_idx, cur_big_page_tokens:shared_kv_len - ] = tail_mems - - # 将 对应的 value_tensors 中的 kv 数据 拷贝到 tail_mems 中对应的数据去 - radix_cache.mem_manager.operator.copy_mem_to_mem( - value_tensor[cur_big_page_tokens:shared_kv_len], tail_mems - ) - # 尾部 KV 换到新 mem 后，同步拷贝已捕获的 top-k prompt logprobs。 - self.prompt_selected_logprobs.copy_capture_slots_if_needed( - source_indexes=value_tensor[cur_big_page_tokens:shared_kv_len], - destination_indexes=tail_mems, - ) - - self.shared_kv_node = share_node # 只是为了保证 restore_small_page_state 正确调用 - g_infer_context.req_manager.restore_small_page_state( - req=self, - ) - self.shared_kv_node = None - - big_page_shared_node = radix_cache.deref_to_first_big_page_node(node=share_node) - self.shared_kv_node = big_page_shared_node - self.cur_kv_len = int(shared_kv_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换 - self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 - else: - # 没有充足的token 容量时， 直接找到最接近的大页，进行大页恢复 - share_node = radix_cache.deref_to_first_big_page_node( - node=share_node, mm_input_spans=self.mm_input_spans - ) - if share_node is not None: - assert share_node.is_big_page_node() - # 大页匹配 - self.shared_kv_node = share_node - ready_cache_len = share_node.node_prefix_total_len - # 从 cpu 到 gpu 是流内阻塞操作 - g_infer_context.req_manager.req_to_token_indexs[ - self.req_idx, 0:ready_cache_len - ] = value_tensor[0:ready_cache_len] - self.cur_kv_len = int(ready_cache_len) # 序列化问题, 该对象可能为numpy.int64，用 int(*)转换 - self.shm_req.prompt_cache_len = self.cur_kv_len # 记录 prompt cache 的命中长度 - assert self.tail_small_page_buffer_id is None - # 恢复 hybrid checkpoint - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=share_node.big_page_buffer_idx, req=self - ) - - self.shm_req.shm_cur_kv_len = self.cur_kv_len - - # hybrid radix cache 使用自己的 hash page；启动检查保证它是模型 - # page_size 的整数倍，因此命中前缀和尾部重建后的持有长度均可直接 - # 作为请求分页预留边界。 - self.hold_kv_len = self.cur_kv_len - 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 + g_infer_context.cache_coordinator.restore_prefix(self, g_infer_context.req_manager, g_infer_context.radix_cache) def is_master_req(self): """ diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 6c621b6adb..7313eafa6d 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -7,6 +7,7 @@ import torch.distributed as dist from typing import List, Tuple, Callable, Optional, Union from transformers.configuration_utils import PretrainedConfig +from lightllm.common.kv_transfer_resources import KVTransferResources from lightllm.utils.infer_utils import set_random_seed from lightllm.utils.log_utils import init_logger from lightllm.models import get_draft_model_class, get_model @@ -15,7 +16,6 @@ from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.logprobs_manager import PromptLogprobsCaptureManager from lightllm.common.basemodel.moe_route_info_manager import MoeRouteInfoManager -from lightllm.common.req_manager import HybridAttentionReqManager from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput @@ -147,11 +147,15 @@ def init_model(self, kvargs): } self.model, self.is_multimodal = get_model(model_cfg, model_kvargs) self.model: TpPartBaseModel = self.model # for easy typing + if self.args.enable_dp_prompt_cache_fetch and not self.model.mem_manager.operator.supports_dp_cache_fetch: + raise ValueError( + f"--enable_dp_prompt_cache_fetch is not supported by {type(self.model.mem_manager).__name__}" + ) set_random_seed(2147483647) - self.is_hybrid_att_model = isinstance(self.model.req_manager, HybridAttentionReqManager) + self.is_hybrid_att_model = self.model.cache_coordinator.has_state_cache if self.is_hybrid_att_model: - self.small_page_buffers = self.model.req_manager.create_small_page_cache_manager( + self.small_page_buffers = self.model.cache_coordinator.create_small_page_cache_manager( size=self.args.linear_att_cache_size, ) else: @@ -168,6 +172,7 @@ def init_model(self, kvargs): big_page_num=self.args.linear_att_page_block_num, kv_cache_mem_manager=self.model.mem_manager, small_page_buffers=self.small_page_buffers, + big_page_buffers=self.model.cache_coordinator.big_page_buffers, ) else: self.radix_cache = RadixCache( @@ -195,6 +200,7 @@ def init_model(self, kvargs): shm_req_manager=self.shm_req_manager, vocab_size=self.model.vocab_size, cache_placement_controller=cache_placement_controller, + cache_coordinator=self.model.cache_coordinator, ) # 初始化 dp 模式使用的通信 tensor, 对于非dp模式，不会使用到 if self.dp_size > 1: @@ -219,9 +225,12 @@ def init_model(self, kvargs): ) if self.args.run_mode in ["prefill", "decode"] or self.args.enable_dp_prompt_cache_fetch: - # 如果存在需要跨进程使用mem manger的特性，则将mem manager写入到 shm中，方便 - # 读取 - self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager) + self.kv_transfer_resources = KVTransferResources( + self.model.mem_manager.get_transfer_view(include_page_layout=self.is_pd_mode), + self.model.req_manager.req_to_token_indexs, + self.model.cache_coordinator.create_state_transfer(), + ) + self.kv_transfer_resources.write_to_shm() dist.barrier(group=self.node_nccl_group) # 同一 DP 组内只需主 rank 初始化真实的 capture buffer 并执行后续相关操作； @@ -274,7 +283,6 @@ def init_custom(self): def init_dp_kv_shared(self): from lightllm.server.router.model_infer.mode_backend.dp_backend.dp_shared_kv_trans import DPKVSharedMoudle - from lightllm.common.kv_cache_mem_manager import MemoryManager torch.cuda.set_device(get_current_device_id()) @@ -284,13 +292,13 @@ def init_dp_kv_shared(self): backend=self, ) - # Collect mem_managers from all ranks - self.mem_managers = [] - for rank_idx in range(self.node_world_size): - if rank_idx != self.rank_in_node: - self.mem_managers.append(MemoryManager.loads_from_shm(rank_idx)) - else: - self.mem_managers.append(self.model.mem_manager) + resources = [ + self.kv_transfer_resources if rank == self.rank_in_node else KVTransferResources.loads_from_shm(rank) + for rank in range(self.node_world_size) + ] + self.kv_views = [resource.kv for resource in resources] + self.peer_states = [resource.state_transfer for resource in resources] + self.req_to_token_indexes_by_rank = [resource.req_to_token_indexes for resource in resources] return def get_max_total_token_num(self): @@ -722,7 +730,6 @@ 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 diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 9bc8c16cf7..09380963ef 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -123,7 +123,7 @@ def prefill_normal( is_prefill=True, b_prefill_has_output_cpu=model_input.b_prefill_has_output_cpu, ) - g_infer_context.save_hybrid_state_to_cache( + g_infer_context.save_cache_state( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) @@ -222,7 +222,7 @@ def prefill_mtp( target_model_output=model_output, target_next_token_ids=next_token_ids, ) - g_infer_context.save_hybrid_state_to_cache( + g_infer_context.save_cache_state( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py index 413a86b94e..ba7ec5fbe3 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 @@ -83,10 +83,8 @@ def build_shared_kv_trans_tasks( 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_req_idx, req.cur_kv_len : req.cur_kv_len + trans_size - ] + token_indexes = self.backend.req_to_token_indexes_by_rank[max_kv_len_mem_manager_index] + max_kv_len_mem_indexes = token_indexes[max_kv_len_req_idx, req.cur_kv_len : req.cur_kv_len + trans_size] trans_tasks.append( TransTask( req=req, @@ -94,6 +92,7 @@ def build_shared_kv_trans_tasks( max_kv_len_dp_rank=int(max_kv_len_dp_rank), max_kv_len_mem_manager_index=int(max_kv_len_mem_manager_index), max_kv_len_mem_indexes=max_kv_len_mem_indexes, + source_req_idx=max_kv_len_req_idx, ) ) @@ -115,7 +114,7 @@ def kv_trans(self, trans_tasks: List["TransTask"]): max_kv_len_dp_ranks_tensor = torch.tensor(max_kv_len_dp_ranks, dtype=torch.int32, device="cuda") mem_indexes_tensor = torch.cat(mem_indexes).to(dtype=torch.int64, device="cuda") self.backend.model.mem_manager.operator.copy_kv_from_other_dp_ranks( - mem_managers=self.backend.mem_managers, + kv_views=self.backend.kv_views, move_token_indexes=max_kv_len_mem_indexes_tensor, token_dp_indexes=max_kv_len_dp_ranks_tensor, mem_indexes=mem_indexes_tensor, @@ -126,7 +125,14 @@ def kv_trans(self, trans_tasks: List["TransTask"]): self.backend.logger.debug(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) + restored_len = trans_task.req.cur_kv_len + len(trans_task.mem_indexes) + self.backend.model.cache_coordinator.restore_peer_state( + trans_task.req.req_idx, + restored_len, + self.backend.peer_states[trans_task.max_kv_len_mem_manager_index], + trans_task.source_req_idx, + ) + trans_task.req.cur_kv_len = restored_len assert trans_task.req.cur_kv_len <= trans_task.req.hold_kv_len if self.backend.is_master_in_dp: trans_task.req.shm_req.shm_cur_kv_len = trans_task.req.cur_kv_len @@ -139,3 +145,4 @@ class TransTask: max_kv_len_dp_rank: int max_kv_len_mem_manager_index: int max_kv_len_mem_indexes: torch.Tensor + source_req_idx: int diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 0d2330d8a4..5d0675d9a1 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -201,7 +201,7 @@ def prefill_normal( is_prefill=True, b_prefill_has_output_cpu=model_input.b_prefill_has_output_cpu, ) - g_infer_context.save_hybrid_state_to_cache( + g_infer_context.save_cache_state( b_req_idx=model_input.b_req_idx, reqs=run_reqs, ) @@ -324,7 +324,7 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer ) if g_infer_context.is_hybrid_att_model: - g_infer_context.save_hybrid_state_to_cache(b_req_idx=b_req_idx, reqs=run_reqs) + g_infer_context.save_cache_state(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() @@ -448,7 +448,7 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] target_next_token_ids=next_token_ids, ) if req_num > 0: - g_infer_context.save_hybrid_state_to_cache(b_req_idx=b_req_idx, reqs=run_reqs) + g_infer_context.save_cache_state(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() @@ -691,7 +691,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I ) if req_num > 0 and g_infer_context.is_hybrid_att_model: - g_infer_context.save_hybrid_state_to_cache(b_req_idx=b_req_idx, reqs=run_reqs) + g_infer_context.save_cache_state(b_req_idx=b_req_idx, reqs=run_reqs) sync_event = torch.cuda.Event() sync_event.record() diff --git a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py index 736e52f439..fe1fb1d6ca 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 @@ -87,8 +87,14 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): match_tokens = page_len_list[matched_page_num - 1] cache_len = clamp_mm_cache_len(match_tokens, req.mm_input_spans) if cache_len == match_tokens: - break - matched_page_num = bisect.bisect_right(page_len_list, cache_len, hi=matched_page_num) + valid_page_num = self.backend.model.cache_coordinator.get_cpu_cache_match_page_num( + page_list[:matched_page_num], page_len_list[:matched_page_num], self.cpu_cache_client + ) + if valid_page_num == matched_page_num: + break + matched_page_num = valid_page_num + else: + matched_page_num = bisect.bisect_right(page_len_list, cache_len, hi=matched_page_num) match_tokens = page_len_list[matched_page_num - 1] if matched_page_num else 0 # 更新命中的 cpu kv cache 长度, 减去radix cache和disk cache的部分. @@ -121,7 +127,6 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream()) # g_infer_context.get_overlap_stream().synchronize() - mem_manager = self.backend.model.mem_manager req_manager = self.backend.model.req_manager mem_indexes_cuda = mem_indexes.cuda(non_blocking=True) @@ -144,11 +149,12 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): idle_token_num -= need_token_num req.cur_kv_len = req.cur_kv_len + need_token_num - mem_manager.operator.load_cpu_cache_to_gpu( + self.backend.model.cache_coordinator.load_cpu_cache_to_gpu( mem_indexes=mem_indexes_cuda, page_indexes=page_indexes_cuda, cpu_cache_client=self.cpu_cache_client, req=req, + page_indexes_cpu=need_pages, ) torch.cuda.current_stream().synchronize() @@ -228,13 +234,13 @@ def _start_kv_cache_offload_task( assert CacheTier.CPU in req.cache_tiers disk_offload_enable = CacheTier.DISK in req.cache_tiers with torch.cuda.stream(cpu_kv_cache_stream): + cpu_kv_cache_stream.wait_stream(g_infer_context.get_overlap_stream()) # 综合考虑后只对prompt做缓存管理，不包含decode内容，这里与radix cache不一致 token_hash_list = req.shm_req.token_hash_list.get_all() page_len_list = req.shm_req.token_hash_page_len_list.get_all() 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 @@ -243,6 +249,8 @@ def _start_kv_cache_offload_task( req=req, move_block_size=move_block_size, page_len_list=page_len_list ) + cacheable_len = self.backend.model.cache_coordinator.get_cacheable_prefix_len(req, self.args) + move_block_size = min(move_block_size, bisect.bisect_right(page_len_list, cacheable_len)) if move_block_size == 0: dist.broadcast_object_list([0], group=self.gloo_group, group_src=0) req.cpu_cache_task_status = InferReq._CpuCacheTaskStatus.FINISHED @@ -288,14 +296,15 @@ def _start_kv_cache_offload_task( assert req.cur_kv_len >= move_token_num 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 - - mem_manager.operator.offload_gpu_kv_to_cpu_cache( + self.backend.model.cache_coordinator.offload_gpu_kv_to_cpu_cache( mem_indexes=token_indexes, page_indexes=cuda_page_indexes, page_readies=cuda_page_readies, + radix_cache=self.backend.radix_cache, cpu_cache_client=self.cpu_cache_client, req=req, + page_indexes_cpu=page_list, + page_readies_cpu=ready_list, ) # 这个操作只是为了在offload 对应的cuda stream中，同步标记下对应的kv cache offload 操作已经完成， diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index 783a58523a..961219407a 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -1,3 +1,4 @@ +from lightllm.common.kv_page_transfer import KVPageTransfer import torch import time import inspect @@ -7,8 +8,9 @@ import queue import pickle from typing import List, Dict, Union, Optional +from lightllm.common.kv_transfer_resources import KVTransferResources from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.kv_transfer_resources import KVTransferView from lightllm.server.pd_io_struct import ( PDChunckedTransTask, PDChunckedTransTaskGroup, @@ -85,9 +87,10 @@ def _init_env( # 从共享内存读取所有rank的mem_manager node_world_size = args.tp // args.nnodes - mem_managers: List[MemoryManager] = [ - MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size) - ] + resources = [KVTransferResources.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size)] + + kv_views = [resource.kv for resource in resources] + state_transfers = [resource.state_transfer for resource in resources] task_out_queue.put("get_mem_managers_ok") @@ -96,7 +99,8 @@ def _init_env( device_id=device_id, task_in_queue=task_in_queue, task_out_queue=task_out_queue, - mem_managers=mem_managers, + kv_views=kv_views, + state_transfers=state_transfers, up_status_in_queue=up_status_in_queue, ) assert manager is not None @@ -119,7 +123,8 @@ def __init__( device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, - mem_managers: List[MemoryManager], + kv_views: List[KVTransferView], + state_transfers, up_status_in_queue: Optional[mp.SimpleQueue], ): self.args = args @@ -127,12 +132,15 @@ def __init__( self.device_id = device_id self.task_in_queue = task_in_queue self.task_out_queue = task_out_queue - self.mem_managers = mem_managers + self.kv_views = kv_views + self.state_transfers = state_transfers self.up_status_in_queue = up_status_in_queue - cur_mem_manager: MemoryManager = self.mem_managers[device_id] - kv_move_buffer = cur_mem_manager.alloc_paged_kv_move_buffer( + self.kv_page_transfer = KVPageTransfer(self.kv_views[device_id]) + kv_move_buffer = self.kv_page_transfer.alloc_paged_kv_move_buffer( page_num=self.args.pd_kv_page_num, page_size=self.args.pd_kv_page_size ) + if self.state_transfers[device_id] is not None: + self.state_transfers[device_id].bind_page_buffer(kv_move_buffer) self.copy_cuda_stream = torch.cuda.Stream(priority=-1) self.transporter = create_kv_transporter( args=self.args, @@ -170,12 +178,12 @@ def __init__( def _warmup(self): for dp_index in range(self.args.dp // self.args.nnodes): with torch.cuda.stream(stream=self.copy_cuda_stream): - cur_mem = self.mem_managers[self.device_id] + cur_mem = self.kv_page_transfer cur_mem.read_page_kv_move_buffer_to_mem( mem_indexes=[0], page_index=0, dp_index=dp_index, - mem_managers=self.mem_managers, + kv_views=self.kv_views, dp_world_size=self.dp_world_size, ) torch.cuda.current_stream().synchronize() @@ -396,16 +404,21 @@ def read_page_to_mems_loop(self): copy_end_event = torch.cuda.Event(enable_timing=True) with torch.cuda.stream(stream=self.copy_cuda_stream): copy_start_event.record(self.copy_cuda_stream) - cur_mem = self.mem_managers[self.device_id] - cur_mem.read_page_kv_move_buffer_to_mem( - trans_task.mem_indexes, - page_index=trans_task.dst_page_index, - dp_index=trans_task.decode_dp_index, - mem_managers=self.mem_managers, - dp_world_size=self.dp_world_size, - page_kind=trans_task.page_kind, - req_idx=trans_task.req_idx, - ) + cur_mem = self.kv_page_transfer + if trans_task.page_kind == "att_state": + helper = self.state_transfers[self.device_id] + assert helper is not None + states = helper.get_dp_states(self.state_transfers, trans_task.decode_dp_index, self.dp_world_size) + helper.read_page_to_req(trans_task.dst_page_index, trans_task.req_idx, states) + else: + assert trans_task.page_kind == "kv" + cur_mem.read_page_kv_move_buffer_to_mem( + trans_task.mem_indexes, + page_index=trans_task.dst_page_index, + dp_index=trans_task.decode_dp_index, + kv_views=self.kv_views, + dp_world_size=self.dp_world_size, + ) copy_end_event.record(self.copy_cuda_stream) self.success_queue.put((copy_end_event, copy_start_event, trans_task)) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py index 6e2084b04d..2409ac6074 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py @@ -1,3 +1,4 @@ +from lightllm.common.kv_page_transfer import KVPageTransfer import torch import time import inspect @@ -7,8 +8,9 @@ import queue import pickle from typing import List, Dict, Optional, Union +from lightllm.common.kv_transfer_resources import KVTransferResources from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.kv_transfer_resources import KVTransferView from lightllm.server.pd_io_struct import PDAbortReq, PDChunckedTransTask from lightllm.utils.graceful_utils import graceful_registry from lightllm.server.core.objs import StartArgs @@ -60,9 +62,10 @@ def _init_env( # 从共享内存读取所有rank的mem_manager node_world_size = args.tp // args.nnodes - mem_managers: List[MemoryManager] = [ - MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size) - ] + resources = [KVTransferResources.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size)] + + kv_views = [resource.kv for resource in resources] + state_transfers = [resource.state_transfer for resource in resources] task_out_queue.put("get_mem_managers_ok") @@ -71,7 +74,8 @@ def _init_env( device_id=device_id, task_in_queue=task_in_queue, task_out_queue=task_out_queue, - mem_managers=mem_managers, + kv_views=kv_views, + state_transfers=state_transfers, ) assert manager is not None task_out_queue.put("module_ready") @@ -93,19 +97,23 @@ def __init__( device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, - mem_managers: List[MemoryManager], + kv_views: List[KVTransferView], + state_transfers, ) -> None: self.args = args self.dp_world_size = self.args.tp // self.args.dp self.device_id = device_id self.task_in_queue = task_in_queue self.task_out_queue = task_out_queue - self.mem_managers = mem_managers + self.kv_views = kv_views + self.state_transfers = state_transfers - cur_mem_manager: MemoryManager = self.mem_managers[device_id] - kv_move_buffer = cur_mem_manager.alloc_paged_kv_move_buffer( + self.kv_page_transfer = KVPageTransfer(self.kv_views[device_id]) + kv_move_buffer = self.kv_page_transfer.alloc_paged_kv_move_buffer( page_num=self.args.pd_kv_page_num, page_size=self.args.pd_kv_page_size ) + if self.state_transfers[device_id] is not None: + self.state_transfers[device_id].bind_page_buffer(kv_move_buffer) self.copy_cuda_stream = torch.cuda.Stream(priority=-1) self.transporter = create_kv_transporter( args=self.args, @@ -145,12 +153,12 @@ def __init__( def _warmup(self): for dp_index in range(self.args.dp // self.args.nnodes): with torch.cuda.stream(stream=self.copy_cuda_stream): - cur_mem = self.mem_managers[self.device_id] + cur_mem = self.kv_page_transfer cur_mem.write_mem_to_page_kv_move_buffer( mem_indexes=[0], page_index=0, dp_index=dp_index, - mem_managers=self.mem_managers, + kv_views=self.kv_views, dp_world_size=self.dp_world_size, ) torch.cuda.current_stream().synchronize() @@ -203,16 +211,21 @@ def local_copy_kv_loop(self): # 将kv 数据拷贝到 page 上，然后传输给 decode node，让其进行读取。 with torch.cuda.stream(stream=self.copy_cuda_stream): - cur_mem = self.mem_managers[self.device_id] - cur_mem.write_mem_to_page_kv_move_buffer( - trans_task.mem_indexes, - page_index=trans_task.src_page_index, - dp_index=trans_task.prefill_dp_index, - mem_managers=self.mem_managers, - dp_world_size=self.dp_world_size, - page_kind=trans_task.page_kind, - req_idx=trans_task.req_idx, - ) + cur_mem = self.kv_page_transfer + if trans_task.page_kind == "att_state": + helper = self.state_transfers[self.device_id] + assert helper is not None + states = helper.get_dp_states(self.state_transfers, trans_task.prefill_dp_index, self.dp_world_size) + helper.write_req_to_page(trans_task.src_page_index, trans_task.req_idx, states) + else: + assert trans_task.page_kind == "kv" + cur_mem.write_mem_to_page_kv_move_buffer( + trans_task.mem_indexes, + page_index=trans_task.src_page_index, + dp_index=trans_task.prefill_dp_index, + kv_views=self.kv_views, + dp_world_size=self.dp_world_size, + ) sync_event = torch.cuda.Event() sync_event.record() diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index edde24602d..632da90624 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -31,8 +31,9 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) - if backend.is_hybrid_att_model: - backend.model.req_manager.update_mtp_state( + state_manager = backend.model.cache_coordinator.state_manager + if state_manager is not None: + state_manager.commit_accepted_state( b_req_mtp_start_loc=b_req_mtp_start_loc, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index b3047b488e..63df930c00 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -17,12 +17,12 @@ ) from lightllm.utils.log_utils import init_logger from lightllm.utils.config_utils import get_num_key_value_heads, get_head_dim, get_layer_num, is_hybrid_att_model -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class +from lightllm.utils.llm_utils import get_llm_model_class from lightllm.common.kv_cache_mem_manager import ( MemoryManager, PPLINT8KVMemoryManager, PPLINT4KVMemoryManager, - Deepseek2MemoryManager, + MlaMemoryManager, ) from typing import List, Tuple, Optional @@ -63,7 +63,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": assert args.enable_cpu_cache is_hybrid_model = is_hybrid_att_model(args.model_dir) - mem_manager_class = None if is_hybrid_model else select_mem_manager_class() + mem_manager_class = None if is_hybrid_model else get_llm_model_class().get_mem_manager_class() if is_hybrid_model: hybrid_config = get_hybrid_cache_config() cpu_cache_meta = CpuKVCacheMeta( @@ -76,7 +76,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": scale_head_dim=0, scale_data_type=get_llm_data_type(), ) - elif mem_manager_class is Deepseek2MemoryManager: + elif mem_manager_class is MlaMemoryManager: cpu_cache_meta = CpuKVCacheMeta( page_num=0, token_page_size=args.cpu_cache_token_page_size, diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index b21c897d10..dfbd2abdae 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -73,8 +73,8 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st infer_state.b_mtp_index = torch.zeros_like(b_req_idx) infer_state.req_manager = SimpleNamespace( HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, - req_to_mtp_state_index=req_to_mtp_state_index, ) + infer_state.linear_state_manager = SimpleNamespace(mtp_state_index=req_to_mtp_state_index) state = gdn.LinearAttDecodeAttState( backend=SimpleNamespace( model=model, diff --git a/unit_tests/common/test_cache_coordinator.py b/unit_tests/common/test_cache_coordinator.py new file mode 100644 index 0000000000..4130ea1544 --- /dev/null +++ b/unit_tests/common/test_cache_coordinator.py @@ -0,0 +1,447 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.cache_coordinator import HybridCacheCoordinator + + +@pytest.mark.parametrize( + "hit", ["big", "small", "small_only", "small_no_ancestor", "fallback", "fallback_no_ancestor", "cold", "existing"] +) +def test_restore_prefix(hit): + events = [] + big = SimpleNamespace(node_prefix_total_len=8, big_page_buffer_idx=1, is_big_page_node=lambda: True) + small = SimpleNamespace(node_prefix_total_len=12, small_page_buffer_idx=2, is_big_page_node=lambda: False) + req = SimpleNamespace( + req_idx=1, + args=SimpleNamespace( + linear_att_hash_page_size=4, + linear_att_page_block_num=4 if hit == "small_only" else 2, + max_req_total_len=12 if hit == "small_only" else 64, + page_size=4, + ), + sampling_param=SimpleNamespace(disable_prompt_cache=False), + shm_req=SimpleNamespace( + hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [1, 2, 3], size=3), + shm_prompt_ids=SimpleNamespace(arr=list(range(13))), + prompt_cache_len=0, + ), + cur_kv_len=0, + get_cur_total_len=lambda: 13, + shared_kv_node=None, + tail_small_page_buffer_id=None, + mm_input_spans=[], + prompt_selected_logprobs=SimpleNamespace( + copy_capture_slots_if_needed=lambda **kwargs: events.append("logprobs") + ), + ) + if hit == "existing": + req.cur_kv_len = 8 + req.shared_kv_node = big + req.shm_req.prompt_cache_len = 4 + if hit.endswith("no_ancestor"): + small.node_prefix_total_len = 4 + source = None if hit == "cold" else big if hit == "big" else small + + def match(key, **kwargs): + assert kwargs["update_refs"] + assert key.tolist() == list(range(12)) + events.append("match") + count = source.node_prefix_total_len if source is not None else 0 + return source, count, torch.arange(count, dtype=torch.int32) + + def restore(req_idx, conv, ssm): + assert req_idx == req.req_idx + expected_slot = 1 if hit in {"big", "fallback"} else 2 + assert conv == ssm == expected_slot + if hit in {"small", "small_no_ancestor"}: + assert req.shared_kv_node is None + events.append("restore") + + manager = HybridCacheCoordinator( + SimpleNamespace( + restore_checkpoint=lambda req_idx, length, pool, slot: restore(req_idx, *pool.get_state_cache(slot)), + reset=lambda req_idx: events.append("reset"), + conv_state=None, + ssm_state=None, + ), + SimpleNamespace(reset=lambda req_idx: events.append("tail"), state=None), + ) + manager.small_page_buffers = SimpleNamespace( + get_state_cache=lambda slot: (slot, slot), has_checkpoint=lambda slot, length: slot is not None + ) + mem = SimpleNamespace( + allocator=SimpleNamespace(can_use_mem_size=4 if hit.startswith("fallback") else 8), + alloc=lambda need_size: torch.arange(20, 20 + need_size, dtype=torch.int32), + operator=SimpleNamespace(copy_mem_to_mem=lambda src, dst: events.append("copy")), + ) + manager.mem_manager = mem + manager.big_page_buffers = SimpleNamespace( + get_state_cache=lambda slot: (slot, slot), has_checkpoint=lambda slot, length: slot is not None + ) + requests = SimpleNamespace(mem_manager=mem, req_to_token_indexs=torch.full((2, 16), -1, dtype=torch.int32)) + radix = SimpleNamespace( + match_prefix=match, + mem_manager=mem, + get_tree_total_tokens_num=lambda: 12, + get_refed_tokens_num=lambda: 12, + free_radix_cache_to_get_enough_token=lambda **kwargs: events.append("reserve"), + deref_to_first_big_page_node=lambda **kwargs: None if hit.endswith("no_ancestor") else big, + ) + manager.restore_prefix(req, requests, radix) + + if hit == "existing": + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == 8 + assert req.shared_kv_node is big + assert req.shm_req.prompt_cache_len == 4 + assert not events + return + length = ( + 0 + if hit in {"cold", "fallback_no_ancestor"} + else 8 + if hit in {"big", "fallback"} + else 4 + if hit == "small_no_ancestor" + else 12 + ) + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == length + assert req.shm_req.prompt_cache_len == length + assert req.shared_kv_node is ( + None if hit in {"cold", "small_no_ancestor", "fallback_no_ancestor"} else small if hit == "small_only" else big + ) + if hit == "small_no_ancestor": + assert requests.req_to_token_indexs[1, :4].tolist() == list(range(20, 24)) + assert events == ["match", "reserve", "copy", "logprobs", "restore", "tail"] + elif hit == "small": + assert requests.req_to_token_indexs[1, :12].tolist() == list(range(8)) + list(range(20, 24)) + assert events == ["match", "reserve", "copy", "logprobs", "restore", "tail"] + elif hit in {"cold", "fallback_no_ancestor"}: + assert events == ["match", "reset", "tail"] + assert (requests.req_to_token_indexs == -1).all() + else: + assert requests.req_to_token_indexs[1, :length].tolist() == list(range(length)) + assert events == ["match", "restore", "tail"] + + +def test_memory_binding_rejects_another_pool(): + manager = HybridCacheCoordinator() + mem = SimpleNamespace() + manager.bind_mem_manager(mem) + manager.bind_mem_manager(mem) + with pytest.raises(ValueError, match="different memory manager"): + manager.bind_mem_manager(SimpleNamespace()) + assert manager.mem_manager is mem + + +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_sph", "fp8kv_spt"]) +def test_qwen_runtime_and_cache_share_config(monkeypatch, kv_type): + from lightllm.models.qwen3next import model as qwen + from lightllm.common.state_cache_manager import linear_att + from lightllm.server.core.objs.start_args_type import StartArgs + + args = StartArgs(tp=4, dp=1, data_type="bfloat16", llm_kv_type=kv_type) + monkeypatch.setattr(qwen, "get_env_start_args", lambda: args) + model = object.__new__(qwen.Qwen3NextTpPartModel) + model.args = args + model.config = dict( + n_layer=8, + num_attention_heads=8, + num_key_value_heads=2, + head_dim=16, + linear_num_key_heads=2, + linear_num_value_heads=8, + linear_key_head_dim=8, + linear_value_head_dim=16, + linear_conv_kernel_dim=4, + full_attention_interval=4, + ) + model.max_req_num = 2 + model.tp_world_size_ = 4 + model.num_kv_heads = 1 + model.data_type = torch.bfloat16 + model.max_total_token_num = 32 + model.mem_fraction = 0.9 + events = [] + + def runtime(size, config, args): + events.append("runtime") + return SimpleNamespace(linear_config=config, conv_state=object(), ssm_state=object()) + + def memory(**kwargs): + events.append("memory") + assert events == ["runtime", "memory"] + return SimpleNamespace( + **kwargs, kv_buffer=object(), set_kv_layer_mapping=lambda mapping: events.append(mapping) + ) + + monkeypatch.setattr(linear_att, "get_added_mtp_kv_layer_num", lambda: 1) + monkeypatch.setattr(qwen, "LinearStateManager", runtime) + for name in ( + "MemoryManager", + "FP8StaticPerHeadQuantMemManager", + "FP8StaticPerTensorQuantMemManager", + ): + monkeypatch.setattr(qwen, name, memory) + model._init_cache_coordinator() + model._init_mem_manager() + monkeypatch.setattr("lightllm.common.cache_coordinator.get_env_start_args", lambda: args) + pools = [] + + def checkpoint_pool(size, config, keep_num): + pools.append((size, config, keep_num)) + return SimpleNamespace(size=size) + + monkeypatch.setattr( + model.cache_coordinator.state_manager, + "create_checkpoint_pool", + lambda size, keep_num=0: checkpoint_pool(size, model.linear_config, keep_num), + raising=False, + ) + monkeypatch.setattr(model.cache_coordinator.state_manager, "create_cpu_transfer", lambda *args: None, raising=False) + model.cache_coordinator.bind_mem_manager(model.mem_manager) + pool = model.cache_coordinator.big_page_buffers + model.cache_coordinator.bind_mem_manager(model.mem_manager) + assert model.cache_coordinator.big_page_buffers is pool + assert len(pools) == 1 + assert pools[0][1] is model.linear_config + assert pools[0][2] == 2 + assert not hasattr(model.mem_manager, "big_page_buffers") + assert not hasattr(model.mem_manager, "req_to_conv_state") + config = model.linear_config + assert config is model.cache_coordinator.linear_state_manager.linear_config + assert events[-1] == config.get_full_att_kv_layer_mapping() + assert config.num_linear_k_heads == 1 + assert config.num_linear_v_heads == 2 + assert config.full_att_dtype == (torch.bfloat16 if kv_type == "None" else torch.uint8) + assert config.get_full_att_kv_layer_index(8) == 2 + + +@pytest.mark.parametrize("mismatch", [None, "requests", "radix"]) +def test_context_requires_shared_memory(monkeypatch, mismatch): + from lightllm.server.router.model_infer import infer_batch + + monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: SimpleNamespace(output_constraint_mode="none")) + mem = SimpleNamespace() + coordinator = HybridCacheCoordinator() + coordinator.bind_mem_manager(mem) + args = dict( + backend=SimpleNamespace(), + req_manager=SimpleNamespace( + mem_manager=object() if mismatch == "requests" else mem, req_sampling_params_manager=None + ), + radix_cache=SimpleNamespace(mem_manager=object() if mismatch == "radix" else mem), + shm_req_manager=None, + vocab_size=1, + cache_coordinator=coordinator, + ) + context = infer_batch.InferenceContext() + if mismatch: + with pytest.raises(ValueError, match="share the memory manager"): + context.register(**args) + else: + context.register(**args) + assert context.cache_coordinator is coordinator + args.pop("cache_coordinator") + with pytest.raises(TypeError, match="cache_coordinator"): + context.register(**args) + + +@pytest.mark.parametrize("pool", ["big_page_buffers", "small_page_buffers"]) +def test_radix_cannot_use_a_different_checkpoint_pool(pool): + coordinator = HybridCacheCoordinator(SimpleNamespace()) + coordinator.mem_manager = object() + coordinator.big_page_buffers = object() + coordinator.small_page_buffers = object() + radix = SimpleNamespace( + mem_manager=coordinator.mem_manager, + big_page_buffers=coordinator.big_page_buffers, + small_page_buffers=coordinator.small_page_buffers, + ) + coordinator.validate_radix_cache(radix) + setattr(radix, pool, object()) + with pytest.raises(ValueError, match="checkpoint pools"): + coordinator.validate_radix_cache(radix) + + +def test_checkpoint_pool_cannot_be_replaced(): + coordinator = HybridCacheCoordinator() + pool = coordinator.small_page_buffers = SimpleNamespace(size=8) + assert coordinator.create_small_page_cache_manager(8) is pool + with pytest.raises(ValueError, match="replace"): + coordinator.create_small_page_cache_manager(16) + + +def test_full_attention_uses_the_coordinator_for_restore_and_release(): + from lightllm.common.cache_coordinator import CacheCoordinator + + coordinator = CacheCoordinator() + assert coordinator.linear_state_manager is coordinator.indexer_state_manager is None + node = SimpleNamespace(node_prefix_total_len=4) + references = [] + radix = SimpleNamespace( + match_prefix=lambda *args, **kwargs: (node, 4, torch.tensor([10, 11, 12, 13])), + insert=lambda key, value: (6, node), + dec_node_ref_counter=references.append, + ) + reqs = SimpleNamespace(req_to_token_indexs=torch.arange(20, 28).reshape(1, 8)) + req = SimpleNamespace( + req_idx=0, + args=SimpleNamespace(page_size=2), + cur_kv_len=0, + hold_kv_len=0, + shared_kv_node=None, + mm_input_spans=[], + sampling_param=SimpleNamespace(disable_prompt_cache=False), + shm_req=SimpleNamespace(shm_prompt_ids=SimpleNamespace(arr=list(range(9)))), + get_cur_total_len=lambda: 9, + get_input_token_ids=lambda: list(range(9)), + ) + coordinator.restore_prefix(req, reqs, radix) + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.prompt_cache_len == 4 + req.cur_kv_len, req.hold_kv_len = 7, 8 + released = coordinator.release_request_cache(req, reqs, radix, req.args, True) + assert [part.tolist() for part in released] == [[26, 27], [24, 25]] + assert references == [node] + assert req.shared_kv_node is None + + +@pytest.mark.parametrize("local_prefix", [0, 8]) +def test_peer_runtime_without_intermediate_checkpoints_is_not_published(local_prefix): + from unittest.mock import Mock + + coordinator = HybridCacheCoordinator() + coordinator.big_page_buffers = SimpleNamespace( + free_state_cache=Mock(), has_checkpoint=lambda slot, length: slot is not None + ) + coordinator.small_page_buffers = SimpleNamespace( + free_state_cache=Mock(), has_checkpoint=lambda slot, length: slot is not None + ) + shared = SimpleNamespace(node_prefix_total_len=local_prefix, big_page_buffer_idx=0) if local_prefix else None + req = SimpleNamespace( + req_idx=0, + cur_kv_len=28, + hold_kv_len=28, + shared_kv_node=shared, + hybrid_cache_len=28, + hybrid_len_to_big_page_id={24: 3}, + tail_small_page_buffer_id=7, + ) + args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) + requests = SimpleNamespace(req_to_token_indexs=torch.arange(28).reshape(1, 28)) + radix = SimpleNamespace(insert=Mock(), dec_node_ref_counter=Mock()) + assert coordinator.get_cacheable_prefix_len(req, args) == local_prefix + freed = coordinator.release_request_cache(req, requests, radix, args, True) + assert torch.cat(freed).tolist() == list(range(local_prefix, 28)) + radix.insert.assert_not_called() + coordinator.big_page_buffers.free_state_cache.assert_called_once_with([3]) + coordinator.small_page_buffers.free_state_cache.assert_called_once_with([7]) + assert not req.hybrid_len_to_big_page_id and req.tail_small_page_buffer_id is None + + +def test_release_keeps_only_the_checkpoint_complete_prefix(): + from unittest.mock import Mock + + coordinator = HybridCacheCoordinator() + coordinator.big_page_buffers = SimpleNamespace( + free_state_cache=Mock(), has_checkpoint=lambda slot, length: slot is not None + ) + coordinator.small_page_buffers = SimpleNamespace( + free_state_cache=Mock(), has_checkpoint=lambda slot, length: slot is not None + ) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=28, + hold_kv_len=28, + shared_kv_node=None, + hybrid_cache_len=28, + hybrid_len_to_big_page_id={8: 2, 24: 3}, + tail_small_page_buffer_id=7, + get_input_token_ids=lambda: list(range(28)), + shm_req=SimpleNamespace(hybrid_token_hash_list=SimpleNamespace(get_all=lambda: list(range(7)))), + ) + args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) + requests = SimpleNamespace(req_to_token_indexs=torch.arange(28).reshape(1, 28)) + + def insert(key, value, **kwargs): + assert len(key) == len(value) == 8 + assert kwargs["len_to_big_page_id"] == {8: 2} + kwargs["len_to_big_page_id"].clear() + return 0, None + + assert coordinator.get_cacheable_prefix_len(req, args) == 8 + freed = coordinator.release_request_cache(req, requests, SimpleNamespace(insert=insert), args, True) + assert torch.cat(freed).tolist() == list(range(8, 28)) + coordinator.big_page_buffers.free_state_cache.assert_called_once_with([3]) + coordinator.small_page_buffers.free_state_cache.assert_called_once_with([7]) + + +def test_dp_restore_copies_linear_state_and_indexer_tail(): + from lightllm.common.linear_state import LinearStateManager + + state = LinearStateManager.__new__(LinearStateManager) + state.linear_config = SimpleNamespace(get_conv_state_shape=lambda: (4, 3)) + state.mtp_step = 2 + state.mtp_state_index = torch.full((2,), 2, dtype=torch.int32) + state.conv_state = SimpleNamespace(buffer=torch.zeros(1, 2, 4, 5)) + state.ssm_state = SimpleNamespace(buffer=torch.zeros(1, 6, 4, 4)) + indexer = SimpleNamespace(state=SimpleNamespace(buffer=torch.zeros(1, 2, 6, 4))) + coordinator = HybridCacheCoordinator(state, indexer) + peer = SimpleNamespace( + mtp_step=2, + req_to_conv_state=SimpleNamespace(buffer=torch.randn(1, 2, 4, 5)), + req_to_ssm_state=SimpleNamespace(buffer=torch.randn(1, 6, 4, 4)), + req_to_indexer_tail=SimpleNamespace(buffer=torch.randn(1, 2, 6, 4)), + ) + coordinator.restore_peer_state(0, 11, peer, 1) + torch.testing.assert_close(state.conv_state.buffer[:, 0, ..., :3], peer.req_to_conv_state.buffer[:, 1, ..., :3]) + torch.testing.assert_close(state.ssm_state.buffer[:, 0], peer.req_to_ssm_state.buffer[:, 3]) + torch.testing.assert_close(indexer.state.buffer[:, 0], peer.req_to_indexer_tail.buffer[:, 1]) + assert not state.conv_state.buffer[:, 0, ..., 3:].any() + assert state.mtp_state_index.tolist() == [0, 2] + + +def test_checkpoint_lengths(monkeypatch): + from unittest.mock import Mock + from lightllm.common.state_cache_manager import StateCacheManager + + class Pool(StateCacheManager): + def get_state_cache(self, buffer_idx): + return buffer_idx + + state = SimpleNamespace(save_state=Mock(), restore_checkpoint=Mock()) + coordinator = HybridCacheCoordinator(state_manager=state) + pool = coordinator.big_page_buffers = Pool(4, keep_num=1) + coordinator.small_page_buffers = Pool(1) + slot = pool.alloc_one_state_cache() + with pytest.raises(AssertionError): + coordinator.restore_state(0, pool, slot, 8) + coordinator.save_state(0, slot, pool, 8) + with pytest.raises(AssertionError): + coordinator.restore_state(0, pool, slot, 16) + coordinator.restore_state(0, pool, slot, 8) + state.restore_checkpoint.assert_called_once_with(0, 8, pool, slot) + pool.free_state_cache([slot]) + assert not pool.has_checkpoint(slot, 8) + + args = SimpleNamespace(cpu_cache_token_page_size=8, linear_att_hash_page_size=4, linear_att_page_block_num=2) + monkeypatch.setattr("lightllm.common.cache_coordinator.get_env_start_args", lambda: args) + coordinator.cpu_cache_transfer = SimpleNamespace( + load=Mock(), get_checkpoint_lens=lambda pages, client: [None, 16][: len(pages)] + ) + assert coordinator.get_cpu_cache_match_page_num([0], [8], None) == 0 + assert coordinator.get_cpu_cache_match_page_num([0, 1], [8, 16], None) == 2 + req = SimpleNamespace(req_idx=0, cur_kv_len=16, mm_input_spans=[], hybrid_len_to_big_page_id={}) + coordinator.load_cpu_cache_to_gpu(torch.arange(16), torch.arange(2), None, req) + first, last = [req.hybrid_len_to_big_page_id[length] for length in (8, 16)] + assert not pool.has_checkpoint(first, 8) + assert pool.has_checkpoint(last, 16) + state.restore_checkpoint.assert_called_with(0, 16, pool, last) + req.hybrid_cache_len = req.cur_kv_len + req.tail_small_page_buffer_id = req.shared_kv_node = None + assert coordinator.get_cacheable_prefix_len(req, args) == 16 + req.cur_kv_len = 8 + assert coordinator.get_cacheable_prefix_len(req, args) == 0 + pool.clear_to_init_state() + assert pool.checkpoint_lens == [None] * pool.size diff --git a/unit_tests/common/test_hybrid_cache_transfer.py b/unit_tests/common/test_hybrid_cache_transfer.py new file mode 100644 index 0000000000..ba4b20e2b7 --- /dev/null +++ b/unit_tests/common/test_hybrid_cache_transfer.py @@ -0,0 +1,65 @@ +from types import SimpleNamespace + +import torch + +from lightllm.common import hybrid_cache_transfer as transfer +from lightllm.common.basemodel.triton_kernel import linear_att_cpu_cache_copy as kernels + + +def test_cpu_checkpoint_metadata(monkeypatch): + class CudaIndexes(list): + is_cuda = True + + args = SimpleNamespace(cpu_cache_token_page_size=8, linear_att_hash_page_size=4) + monkeypatch.setattr(transfer, "get_env_start_args", lambda: args) + monkeypatch.setattr(transfer, "get_current_rank_in_dp", lambda: 0) + monkeypatch.setattr(transfer, "get_dp_world_size", lambda: 1) + copies = [] + monkeypatch.setattr(kernels, "copy_kv_buffer_to_cpu_cache", lambda **kwargs: copies.append(kwargs)) + tensor = torch.tensor + + def cpu_tensor(data, **kwargs): + kwargs.pop("pin_memory", None) + return tensor(data, **kwargs) + + monkeypatch.setattr(torch, "tensor", cpu_tensor) + pages = torch.full((4, 1, 1, 1, 64), 37, dtype=torch.uint8) + pages[1].reshape(-1)[-16:].view(torch.int64).copy_(torch.tensor([24, 1])) + before = pages.clone() + pool = SimpleNamespace( + checkpoint_lens=[None, 16, None, 32], + conv_state_cache=SimpleNamespace(buffer=None), + ssm_state_cache=SimpleNamespace(buffer=None), + ) + writer = transfer.HybridCpuCacheTransfer(None, None, pool) + writer.big_page_ids_buffer_store = torch.empty(4, dtype=torch.int64) + client = SimpleNamespace(cpu_kv_cache_tensor=pages) + page_indexes = [2, 0, 1, -1] + page_readies = [False, False, True, False] + + writer.offload( + CudaIndexes(range(32)), + CudaIndexes(page_indexes), + CudaIndexes(page_readies), + client, + [0, 1, 2, 3], + None, + page_indexes_cpu=page_indexes, + page_readies_cpu=page_readies, + ) + + assert len(copies) == 1 + assert copies[0]["big_page_buffer_ids"].tolist() == [0, 1, 2, 3] + assert pages[2].reshape(-1)[-16:].view(torch.int64).tolist() == [-1, 1] + assert pages[0].reshape(-1)[-16:].view(torch.int64).tolist() == [16, 1] + assert torch.equal(pages[1], before[1]) + assert torch.equal(pages[3], before[3]) + assert torch.equal(pages.reshape(4, -1)[:, :-16], before.reshape(4, -1)[:, :-16]) + + serialized = pages.numpy().tobytes() + restored = torch.frombuffer(bytearray(serialized), dtype=torch.uint8).reshape_as(pages) + reader = transfer.HybridCpuCacheTransfer(None, None, None) + restored_client = SimpleNamespace(cpu_kv_cache_tensor=restored) + assert reader.get_checkpoint_lens([2, 0, 1], restored_client) == [None, 16, 24] + restored[0].reshape(-1)[-8:].view(torch.int64).fill_(0) + assert reader.get_checkpoint_lens([0], restored_client) == [None] diff --git a/unit_tests/common/test_kv_transfer_resources.py b/unit_tests/common/test_kv_transfer_resources.py new file mode 100644 index 0000000000..600c9b2a04 --- /dev/null +++ b/unit_tests/common/test_kv_transfer_resources.py @@ -0,0 +1,93 @@ +from contextlib import nullcontext +import pickle +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from lightllm.common import kv_transfer_resources as shared + + +def test_dp_export_without_pd_layout(): + mem = SimpleNamespace( + kv_buffer=object(), + dtype="bfloat16", + kv_is_replicated=False, + get_paged_kv_move_buffer_shape=Mock(side_effect=RuntimeError("PD layout unavailable")), + ) + view = shared.KVTransferView.from_mem_manager(mem) + assert view.kv_buffer is mem.kv_buffer + assert view.page_shape is None + mem.get_paged_kv_move_buffer_shape.assert_not_called() + with pytest.raises(RuntimeError, match="PD layout unavailable"): + shared.KVTransferView.from_mem_manager(mem, include_page_layout=True) + + +def test_pd_export_with_page_layout(): + mem = SimpleNamespace( + kv_buffer=object(), + dtype="bfloat16", + kv_is_replicated=True, + get_paged_kv_move_buffer_shape=Mock(return_value=(1, 1, 2, 1, 576)), + ) + view = shared.KVTransferView.from_mem_manager(mem, include_page_layout=True) + assert view.page_shape == (2, 1, 576) + assert view.page_dtype == mem.dtype + assert view.replicated + mem.get_paged_kv_move_buffer_shape.assert_called_once_with(1, 1) + + +def test_shared_resource_exports_have_independent_consumer_slots(monkeypatch): + segment = SimpleNamespace(name="test", buf=None, close=lambda: None) + + def shm(**kwargs): + if kwargs["force_mode"] == "create": + segment.size = kwargs["expected_size"] + segment.buf = memoryview(bytearray(segment.size)) + return segment + + monkeypatch.setattr(shared, "create_or_link_shm", shm) + monkeypatch.setattr(shared, "FileLock", lambda _: nullcontext()) + monkeypatch.setattr(shared, "get_service_shm_name", lambda name: name) + monkeypatch.setattr(shared, "get_node_world_size", lambda: 2) + monkeypatch.setattr(shared, "get_current_rank_in_node", lambda: 0) + monkeypatch.setattr(shared, "kv_trans_use_p2p", lambda: False) + monkeypatch.setattr( + shared, "ForkingPickler", SimpleNamespace(dumps=lambda obj: memoryview(pickle.dumps(obj)), loads=pickle.loads) + ) + resources = shared.KVTransferResources(SimpleNamespace(size=32), [[1, 2]], SimpleNamespace(mtp_step=2)) + resources.write_to_shm() + for remaining in (3, 2, 1, 0): + restored = shared.KVTransferResources.loads_from_shm(0) + assert restored.kv.size == 32 + assert restored.req_to_token_indexes == [[1, 2]] + assert restored.state_transfer.mtp_step == 2 + assert int.from_bytes(segment.buf[:4], "little") == remaining + with pytest.raises(AssertionError): + shared.KVTransferResources.loads_from_shm(0) + + +def test_shared_memory_exports_preserve_variable_pickle_lengths(monkeypatch, tmp_path): + from lightllm.common import kv_transfer_resources as module + + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", f"{tmp_path.parent.name}_{tmp_path.name}") + monkeypatch.setattr(module, "get_node_world_size", lambda: 4) + monkeypatch.setattr(module, "get_current_rank_in_node", lambda: 0) + monkeypatch.setattr(module, "kv_trans_use_p2p", lambda: False) + # CUDA IPC ref-counter offsets can cross pickle's one-byte integer boundary + # between exports. Each consumer still needs its own reference-counter token. + offsets = [254, 255, 256, 257, 258, 259, 260, 261] + payloads = [pickle.dumps({"ref_counter_offset": offset}) for offset in offsets] + serialized = iter(payloads) + monkeypatch.setattr(module.ForkingPickler, "dumps", lambda obj: memoryview(next(serialized))) + manager = module.KVTransferResources(None, None) + try: + manager.write_to_shm() + for offset in reversed(offsets): + assert manager.loads_from_shm(0) == {"ref_counter_offset": offset} + with pytest.raises(AssertionError): + manager.loads_from_shm(0) + finally: + shm = module.create_or_link_shm("mem_manager_0", -1, force_mode="link") + shm.close() + shm.unlink() diff --git a/unit_tests/common/test_mem_manager_layouts.py b/unit_tests/common/test_mem_manager_layouts.py new file mode 100644 index 0000000000..dd8b16543f --- /dev/null +++ b/unit_tests/common/test_mem_manager_layouts.py @@ -0,0 +1,284 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.kv_cache_mem_manager import ( + MemoryManager, + FP8StaticPerHeadQuantMemManager, + FP8StaticPerTensorQuantMemManager, + PPLINT4KVMemoryManager, + PPLINT8KVMemoryManager, +) + + +def test_layer_mapping_validates_storage_bounds(): + mem = MemoryManager.__new__(MemoryManager) + mem.layer_num = 3 + assert mem.get_kv_layer_index(2) == 2 + for mapping in ([], [0, -1], [0, 3]): + with pytest.raises(ValueError, match="allocated layers"): + mem.set_kv_layer_mapping(mapping) + mem.set_kv_layer_mapping([0, 0, 1, 1, 2]) + assert [mem.get_kv_layer_index(i) for i in (1, 3, 4)] == [0, 1, 2] + + +@pytest.mark.parametrize("manager_class", [PPLINT4KVMemoryManager, PPLINT8KVMemoryManager]) +def test_integer_cache_capacity_matches_allocated_buffers(monkeypatch, manager_class): + allocate = torch.empty + monkeypatch.setattr(torch, "empty", lambda shape, **kw: allocate(shape, dtype=kw["dtype"])) + mem = manager_class.__new__(manager_class) + mem.page_size, mem.head_num, mem.head_dim, mem.layer_num = 4, 2, 16, 3 + mem.group_quant_size, mem.kv_dtype, mem.dtype = 8, torch.int8, torch.bfloat16 + mem._init_buffers(8, mem.dtype, mem.head_num, mem.head_dim, mem.layer_num) + storage_bytes = sum(t.numel() * t.element_size() for t in (mem.kv_buffer, mem.scale_buffer)) + assert storage_bytes == mem.get_cell_size() * 12 + mem.kv_buffer.random_(-128, 128) + mem.scale_buffer.normal_() + saved = mem.get_index_kv_buffer(slice(1, 3)) + mem.load_index_kv_buffer(slice(5, 7), saved) + for name, expected in saved.items(): + assert torch.equal(getattr(mem, name)[:, 5:7], expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "manager_class", [MemoryManager, FP8StaticPerHeadQuantMemManager, FP8StaticPerTensorQuantMemManager] +) +def test_mapped_target_and_draft_kv_writes(manager_class): + mem = manager_class.__new__(manager_class) + mem.layer_num, mem.head_num = 3, 1 + quantized = manager_class is not MemoryManager + mem.kv_buffer = torch.zeros((3, 8, 2, 16), dtype=torch.uint8 if quantized else torch.bfloat16, device="cuda") + mem.scales = torch.tensor([[1, 1], [2, 4], [4, 2]], dtype=torch.float32, device="cuda") + mem.set_kv_layer_mapping([0, 0, 1, 1, 2]) + operator = mem.operator_class(mem) + indexes = torch.tensor([1, 3], dtype=torch.int32, device="cuda") + kv = torch.arange(64, dtype=torch.bfloat16, device="cuda").reshape(2, 2, 16) + for logical_layer, physical_layer in ((3, 1), (4, 2)): + operator.copy_kv_to_mem_manager(logical_layer, indexes, kv) + k, v = mem.get_att_input_params(logical_layer) + actual = torch.cat((k[indexes], v[indexes]), dim=1) + expected = kv + if quantized: + actual = actual.view(torch.float8_e4m3fn).float() + expected = (kv.float() / mem.scales[physical_layer].view(1, 2, 1)).to(torch.float8_e4m3fn).float() + assert torch.equal(actual, expected) + assert not mem.kv_buffer[0].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("manager_class", [PPLINT4KVMemoryManager, PPLINT8KVMemoryManager]) +def test_integer_cache_copy_includes_scales(manager_class): + mem = manager_class.__new__(manager_class) + mem.kv_buffer = torch.randint(-128, 128, (2, 8, 4, 16), dtype=torch.int8, device="cuda") + mem.scale_buffer = torch.randn((2, 8, 4, 2), dtype=torch.bfloat16, device="cuda") + expected = mem.get_index_kv_buffer([0, 2]) + mem.operator_class(mem).copy_mem_to_mem(torch.tensor([0, 2]), torch.tensor([4, 6])) + for name, values in expected.items(): + assert torch.equal(getattr(mem, name)[:, [4, 6]], values) + + +@pytest.mark.parametrize("quantized", [False, True]) +def test_models_select_mla_layout_without_storage_inspecting_models(monkeypatch, quantized): + from lightllm.models.deepseek2.model import Deepseek2TpPartModel + from lightllm.models.deepseek3_2.model import Deepseek3_2TpPartModel + from lightllm.common.kv_cache_mem_manager import ( + MlaMemoryManager, + IndexedMlaMemoryManager, + FP8IndexedMlaMemoryManager, + ) + + monkeypatch.setattr( + "lightllm.utils.envs_utils.get_env_start_args", + lambda: SimpleNamespace(llm_kv_type="fp8kv_dsa" if quantized else "None"), + ) + assert Deepseek2TpPartModel.get_mem_manager_class() is MlaMemoryManager + assert Deepseek3_2TpPartModel.get_mem_manager_class() is ( + FP8IndexedMlaMemoryManager if quantized else IndexedMlaMemoryManager + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("global_heads,prefill_tp,decode_tp", [(4, 1, 2), (4, 2, 1), (1, 1, 4), (1, 4, 1)]) +def test_kv_transfer_view_roundtrip(global_heads, prefill_tp, decode_tp): + from lightllm.common.kv_transfer_resources import KVTransferView + from lightllm.common.kv_page_transfer import KVPageTransfer + + expected = torch.randn((2, 3, 2, global_heads, 16), dtype=torch.bfloat16, device="cuda") + + def views(tp, initialize): + result = [] + for rank in range(tp): + local = max(1, global_heads // tp) + first = rank // max(1, tp // global_heads) * local + buffer = torch.zeros((2, 8, 2 * local, 16), dtype=torch.bfloat16, device="cuda") + if initialize: + buffer[:, [1, 3, 6]] = expected[:, :, :, first : first + local].reshape(2, 3, 2 * local, 16) + result.append(KVTransferView(buffer, (2, 2 * global_heads, 16), torch.bfloat16, False)) + return result + + sources, destinations = views(prefill_tp, True), views(decode_tp, False) + writer, reader = KVPageTransfer(sources[0]), KVPageTransfer(destinations[0]) + writer.alloc_paged_kv_move_buffer(1, 8) + reader.alloc_paged_kv_move_buffer(1, 8) + writer.write_mem_to_page_kv_move_buffer([1, 3, 6], 0, 0, sources, prefill_tp) + reader.kv_move_buffer.copy_(writer.kv_move_buffer) + reader.read_page_kv_move_buffer_to_mem([0, 2, 5], 0, 0, destinations, decode_tp) + for rank, view in enumerate(destinations): + local = max(1, global_heads // decode_tp) + first = rank // max(1, decode_tp // global_heads) * local + assert torch.equal( + view.kv_buffer[:, [0, 2, 5]], expected[:, :, :, first : first + local].reshape(2, 3, 2 * local, 16) + ) + assert not view.kv_buffer[:, [1, 3, 4, 6, 7]].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("indexed", [False, True]) +def test_mla_writer_respects_the_storage_view(indexed): + from lightllm.common.kv_cache_mem_manager import MlaMemoryManager, IndexedMlaMemoryManager + + cls = IndexedMlaMemoryManager if indexed else MlaMemoryManager + mem = cls.__new__(cls) + mem.head_dim = 576 + (72 if indexed else 0) + mem.layer_num = 2 + mem.kv_buffer = torch.full((2, 8, 1, mem.head_dim), 13, dtype=torch.bfloat16, device="cuda") + mem.set_kv_layer_mapping([0, 0, 1]) + kv = torch.randn((3, 1, 576), dtype=torch.bfloat16, device="cuda") + indexes = torch.tensor([1, 3, 6], dtype=torch.int32, device="cuda") + mem.operator_class(mem).copy_kv_to_mem_manager(2, indexes, kv) + assert torch.equal(mem.get_att_input_params(2)[indexes], kv) + assert (mem.kv_buffer[0] == 13).all() + if indexed: + assert (mem.kv_buffer[..., 576:] == 13).all() + + +@pytest.mark.parametrize( + "kv_type,expected", + [ + ("None", MemoryManager), + ("fp8kv_sph", FP8StaticPerHeadQuantMemManager), + ("fp8kv_spt", FP8StaticPerTensorQuantMemManager), + ], +) +def test_qwen_storage_selection(monkeypatch, kv_type, expected): + from lightllm.models.qwen3next import model + + monkeypatch.setattr(model, "get_env_start_args", lambda: SimpleNamespace(llm_kv_type=kv_type)) + assert model.Qwen3NextTpPartModel.get_mem_manager_class() is expected + for unsupported in ("int4kv", "int8kv"): + monkeypatch.setattr(model, "get_env_start_args", lambda: SimpleNamespace(llm_kv_type=unsupported)) + with pytest.raises(ValueError, match="does not support"): + model.Qwen3NextTpPartModel.get_mem_manager_class() + + +def test_glm_storage_selection(monkeypatch): + from lightllm.models.glm5_next import model + from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager + + monkeypatch.setattr(model, "get_env_start_args", lambda: SimpleNamespace(llm_kv_type="None")) + assert model.Glm5NextTpPartModel.get_mem_manager_class() is Glm5NextMemManager + monkeypatch.setattr(model, "get_env_start_args", lambda: SimpleNamespace(llm_kv_type="fp8kv_sph")) + with pytest.raises(ValueError, match="does not support"): + model.Glm5NextTpPartModel.get_mem_manager_class() + + +@pytest.mark.parametrize("manager_class", [FP8StaticPerTensorQuantMemManager, PPLINT8KVMemoryManager]) +def test_dp_fetch_startup_rejects_unsupported_storage(monkeypatch, manager_class): + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.server.router.model_infer.mode_backend import base_backend + + mem = manager_class.__new__(manager_class) + mem.operator = mem.operator_class(mem) + model = SimpleNamespace(mem_manager=mem) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + monkeypatch.setattr(backend, "init_rank_infos", lambda: None) + monkeypatch.setattr(base_backend, "init_distributed_env", lambda _: None) + monkeypatch.setattr(base_backend.dist_group_manager, "create_groups", lambda group_size: None) + monkeypatch.setattr(base_backend, "TokenLoad", lambda *args: None) + monkeypatch.setattr(base_backend.PretrainedConfig, "get_config_dict", lambda _: ({}, {})) + monkeypatch.setattr(base_backend, "get_model", lambda *args: (model, False)) + with pytest.raises(ValueError, match=f"--enable_dp_prompt_cache_fetch.*{manager_class.__name__}"): + backend.init_model( + { + "args": StartArgs(tp=2, dp=2, enable_dp_prompt_cache_fetch=True), + "world_size": 2, + "load_way": "HF", + "weight_dir": "/model", + "max_total_token_num": 16, + } + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_glm_dp_fetch_copies_packed_indexer_bytes(): + from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager + from lightllm.common.kv_transfer_resources import KVTransferView + + sources = [ + torch.randint(0, 256, (2, 8, 1, 1168), dtype=torch.uint8, device="cuda").view(torch.bfloat16) for _ in range(4) + ] + mem = Glm5NextMemManager.__new__(Glm5NextMemManager) + mem.kv_buffer = torch.zeros_like(sources[0]) + operator = mem.operator_class(mem) + views = [KVTransferView(source, None, source.dtype, True) for source in sources] + operator.copy_kv_from_other_dp_ranks( + kv_views=views, + move_token_indexes=torch.tensor([1, 3, 6], dtype=torch.int64, device="cuda"), + token_dp_indexes=torch.tensor([0, 1, 1], dtype=torch.int32, device="cuda"), + mem_indexes=torch.tensor([2, 4, 7], dtype=torch.int64, device="cuda"), + dp_size_in_node=2, + rank_in_dp=1, + ) + expected = torch.stack([sources[1][:, 1], sources[3][:, 3], sources[3][:, 6]], dim=1) + assert torch.equal(mem.kv_buffer.view(torch.uint8)[:, [2, 4, 7]], expected.view(torch.uint8)) + assert not mem.kv_buffer.view(torch.uint8)[:, [0, 1, 3, 5, 6]].any() + + +@pytest.mark.parametrize("layer_count", [5, 8]) +@pytest.mark.parametrize("draft_layers", [0, 2]) +def test_hybrid_mapping_excludes_linear_layers(layer_count, draft_layers): + from lightllm.common.state_cache_manager import LinearAttCacheConfig + + config = LinearAttCacheConfig( + tp_world_size=1, + full_att_all_num_kv_heads=1, + full_att_dtype=torch.bfloat16, + full_att_num_kv_heads=1, + full_att_head_dim=8, + global_linear_k_heads=1, + global_linear_v_heads=1, + num_linear_k_heads=1, + num_linear_v_heads=1, + head_linear_k_dim=8, + head_linear_v_dim=8, + conv_kernel_size=4, + linear_layer_num=layer_count - layer_count // 4, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4, + all_layer_num=layer_count, + draft_full_att_kv_layer_num=draft_layers, + ) + expected = (None, None, None, 0, None) + if layer_count == 8: + expected += (None, None, 1) + expected += tuple(range(layer_count // 4, layer_count // 4 + draft_layers)) + assert config.get_full_att_kv_layer_mapping() == expected + mem = MemoryManager.__new__(MemoryManager) + mem.layer_num, mem.head_num = config.get_full_att_kv_layer_num_with_draft_model(), 1 + mem.kv_buffer = torch.randn(mem.layer_num, 2, 2, 8) + mem.set_kv_layer_mapping(config.get_full_att_kv_layer_mapping()) + for layer, slot in enumerate(expected): + if slot is None: + with pytest.raises(ValueError, match="no full-attention KV storage"): + config.get_full_att_kv_layer_index(layer) + with pytest.raises(ValueError, match="no KV storage"): + mem.get_att_input_params(layer) + with pytest.raises(ValueError, match="no KV storage"): + mem.operator_class(mem).copy_kv_to_mem_manager(layer, torch.tensor([0]), torch.empty(1, 2, 8)) + else: + assert config.get_full_att_kv_layer_index(layer) == slot + k, v = mem.get_att_input_params(layer) + torch.testing.assert_close(torch.cat((k, v), dim=1), mem.kv_buffer[slot]) diff --git a/unit_tests/common/test_paged_mem_manager.py b/unit_tests/common/test_paged_mem_manager.py index 68bc514991..b8646b1c7f 100644 --- a/unit_tests/common/test_paged_mem_manager.py +++ b/unit_tests/common/test_paged_mem_manager.py @@ -4,9 +4,9 @@ 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 +from lightllm.common.kv_cache_mem_manager.mla_mem_manager import MlaMemoryManager +from lightllm.common.kv_cache_mem_manager.group_quantized_mem_manager import PPLINT4KVMemoryManager +from lightllm.common.kv_cache_mem_manager.group_quantized_mem_manager import PPLINT8KVMemoryManager def test_kv_cache_allocation_requires_complete_pages(): @@ -22,7 +22,7 @@ def test_kv_cache_allocation_requires_complete_pages(): @pytest.mark.parametrize( "manager_class", - [MemoryManager, Deepseek2MemoryManager, PPLINT4KVMemoryManager, PPLINT8KVMemoryManager], + [MemoryManager, MlaMemoryManager, PPLINT4KVMemoryManager, PPLINT8KVMemoryManager], ) def test_kv_cache_buffer_size_requires_complete_pages(manager_class): manager = manager_class.__new__(manager_class) diff --git a/unit_tests/common/test_quantized_paged_mem_manager.py b/unit_tests/common/test_quantized_paged_mem_manager.py index 84fe266613..84d734a83e 100644 --- a/unit_tests/common/test_quantized_paged_mem_manager.py +++ b/unit_tests/common/test_quantized_paged_mem_manager.py @@ -2,8 +2,8 @@ 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 +from lightllm.common.kv_cache_mem_manager.group_quantized_mem_manager import PPLINT4KVMemoryManager +from lightllm.common.kv_cache_mem_manager.group_quantized_mem_manager import PPLINT8KVMemoryManager def test_quantized_kv_buffers_include_the_full_hold_page(monkeypatch): diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index 224c308f28..79398728a2 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -1,3 +1,4 @@ +from lightllm.common.cache_coordinator import HybridCacheCoordinator, CacheCoordinator from types import MethodType, SimpleNamespace import pytest @@ -394,6 +395,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): def test_page_size_one_frees_all_preallocated_indexes(): infer_context = InferenceContext.__new__(InferenceContext) + infer_context.cache_coordinator = CacheCoordinator() 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, :]) @@ -401,6 +403,7 @@ def test_page_size_one_frees_all_preallocated_indexes(): req_idx=0, cur_kv_len=3, hold_kv_len=12, + shared_kv_node=None, shm_req=SimpleNamespace(shm_cur_kv_len=3), ) free_token_indexes = [] @@ -413,6 +416,7 @@ def test_page_size_one_frees_all_preallocated_indexes(): def test_linear_attention_frees_reserved_page_tail(monkeypatch): infer_context = InferenceContext.__new__(InferenceContext) + infer_context.cache_coordinator = CacheCoordinator() 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, :]) @@ -429,13 +433,18 @@ def test_linear_attention_frees_reserved_page_tail(monkeypatch): 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) + free_token_indexes.extend( + HybridCacheCoordinator().release_request_cache( + req, infer_context.req_manager, infer_context.radix_cache, infer_context.args, publish_cache=True + ) + ) 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.cache_coordinator = CacheCoordinator() 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, :]) @@ -452,7 +461,11 @@ def test_linear_attention_frees_preallocated_pages_before_first_forward(monkeypa 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) + free_token_indexes.extend( + HybridCacheCoordinator().release_request_cache( + req, infer_context.req_manager, infer_context.radix_cache, infer_context.args, publish_cache=True + ) + ) assert free_token_indexes[0].tolist() == list(range(8)) @@ -488,6 +501,7 @@ def match_radix_cache(): def test_recover_paused_reqs_uses_page_allocation_need(monkeypatch): infer_context = InferenceContext.__new__(InferenceContext) + infer_context.cache_coordinator = CacheCoordinator() infer_context.args = SimpleNamespace(page_size=4) infer_context.backend = SimpleNamespace(disable_chunked_prefill=False) infer_context.radix_cache = None @@ -515,6 +529,7 @@ def test_recover_paused_reqs_uses_page_allocation_need(monkeypatch): def test_recover_paused_reqs_requires_capacity_for_the_full_sequence(monkeypatch): infer_context = InferenceContext.__new__(InferenceContext) + infer_context.cache_coordinator = CacheCoordinator() infer_context.args = SimpleNamespace(page_size=4) infer_context.radix_cache = None infer_context.req_manager = SimpleNamespace( @@ -535,6 +550,7 @@ def test_recover_paused_reqs_requires_capacity_for_the_full_sequence(monkeypatch def test_recover_paused_reqs_accounts_for_rematched_prefix(monkeypatch): infer_context = InferenceContext.__new__(InferenceContext) + infer_context.cache_coordinator = CacheCoordinator() infer_context.args = SimpleNamespace(page_size=4) infer_context.backend = SimpleNamespace(disable_chunked_prefill=False) infer_context.get_can_alloc_token_num = lambda: 12 diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index ae40702d75..58f35cac3e 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -5,7 +5,10 @@ import torch from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager -from lightllm.common.req_manager import Glm5NextReqManager +from lightllm.common.req_manager import ReqManager +from lightllm.common.cache_coordinator import HybridCacheCoordinator +from lightllm.common.linear_state import LinearStateManager +from lightllm.common.indexer_state import IndexerStateManager from lightllm.common.state_cache_manager import Glm5NextCacheConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args, set_unique_server_name @@ -75,8 +78,10 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t draft_full_att_kv_layer_num=int(mtp_step > 0), ) full_layers = config.get_full_att_kv_layer_num_with_draft_model() - mem = Glm5NextMemManager(16, torch.bfloat16, 1, 584, full_layers, config) - req = Glm5NextReqManager(3, 16, mem, config) + mem = Glm5NextMemManager(16, torch.bfloat16, 1, 584, full_layers) + mem.set_kv_layer_mapping(config.get_full_att_kv_layer_mapping()) + req = HybridCacheCoordinator(LinearStateManager(3, config, args), IndexerStateManager(3, config, mtp_step)) + req.bind_mem_manager(mem) att_kv = mem.get_att_input_params(3) assert att_kv.shape == (17, 1, 512) assert att_kv.stride(0) == 584 @@ -88,45 +93,56 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t mem.operator.copy_kv_to_mem_manager(3, destinations, new_kv) assert torch.equal(att_kv[destinations], new_kv) assert torch.equal(index_bytes, expected_index_bytes) - assert req.get_indexer_tail_buffer(3).shape == (4, 4 + mtp_step, 256) + assert req.indexer_state_manager.get_layer_buffer(3).shape == (4, 4 + mtp_step, 256) if mtp_step: - assert req.get_indexer_tail_buffer(4).shape == (4, 4 + mtp_step, 256) + assert req.indexer_state_manager.get_layer_buffer(4).shape == (4, 4 + mtp_step, 256) assert mem.get_att_input_params(4).shape == att_kv.shape - cache = req.create_small_page_cache_manager(2) if small_page else mem.big_page_buffers + cache = req.create_small_page_cache_manager(2) if small_page else req.big_page_buffers slot = cache.alloc_one_state_cache() source_req = SimpleNamespace(req_idx=0) - req.init_hybrid_attention_state(source_req) - req.req_to_conv_state.buffer[:, 0].normal_() - req.req_to_ssm_state.buffer[:, 0].normal_() - conv = req.req_to_conv_state.buffer[:, 0, ..., :3].clone() - ssm = req.req_to_ssm_state.buffer[:, 0].clone() + req.reset(source_req.req_idx) + req.linear_state_manager.conv_state.buffer[:, 0].normal_() + req.linear_state_manager.ssm_state.buffer[:, 0].normal_() + conv = req.linear_state_manager.conv_state.buffer[:, 0, ..., :3].clone() + ssm = req.linear_state_manager.ssm_state.buffer[:, 0].clone() if small_page: - req.save_state(0, slot, cache) + req.save_state(0, slot, cache, checkpoint_len=4) else: - req.save_big_page_states(torch.tensor([0], dtype=torch.int32, device="cuda"), [0], [slot]) + req.save_big_page_states(torch.tensor([0], dtype=torch.int32, device="cuda"), [0], [slot], checkpoint_lens=[4]) torch.cuda.synchronize() - req.req_to_conv_state.buffer[:, 0].zero_() - req.req_to_ssm_state.buffer[:, 0].zero_() - req.req_to_indexer_tail.buffer.normal_() - unchanged = req.req_to_indexer_tail.buffer[:, [0, 1, 3]].clone() + req.linear_state_manager.conv_state.buffer[:, 0].zero_() + req.linear_state_manager.ssm_state.buffer[:, 0].zero_() + req.indexer_state_manager.state.buffer.normal_() + unchanged = req.indexer_state_manager.state.buffer[:, [0, 1, 3]].clone() + if mtp_step: + state = req.linear_state_manager + state.mtp_state_index.fill_(mtp_step) + tensor = lambda values: torch.tensor(values, dtype=torch.int32, device="cuda") + state.commit_accepted_state( + b_req_mtp_start_loc=tensor([0, 2]), + b_req_idx=tensor([0, 0, 1, 1, 1]), + b_mtp_index=tensor([0, 1, 0, 1, 2]), + accepted_index=tensor([1, 0, 1, 1, 0]), + verify_width=mtp_step + 1, + ) + assert state.mtp_state_index.tolist() == [0, 1, mtp_step, mtp_step] dest_req = SimpleNamespace(req_idx=2, shared_kv_node=SimpleNamespace(small_page_buffer_idx=slot)) - if small_page: - req.restore_small_page_state(dest_req) - else: - req.restore_big_page_state(slot, dest_req) + req.restore_state(dest_req.req_idx, cache, slot, restored_len=4) + if mtp_step: + assert req.linear_state_manager.mtp_state_index.tolist() == [0, 1, 0, mtp_step] torch.cuda.synchronize() - assert torch.equal(req.req_to_conv_state.buffer[:, 2, ..., :3], conv) - assert torch.equal(req.req_to_ssm_state.buffer[:, 2 * (mtp_step + 1)], ssm) + assert torch.equal(req.linear_state_manager.conv_state.buffer[:, 2, ..., :3], conv) + assert torch.equal(req.linear_state_manager.ssm_state.buffer[:, 2 * (mtp_step + 1)], ssm) # Both page sizes are aligned to complete pools. Restoring a prefix must # discard stale tail values from a previously allocated request slot. - assert not req.req_to_indexer_tail.buffer[:, 2].any() - assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) - req.req_to_indexer_tail.buffer[:, 2].normal_() - req.init_hybrid_attention_state(dest_req) - assert not req.req_to_conv_state.buffer[:, 2].any() - assert not req.req_to_ssm_state.buffer[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() - assert not req.req_to_indexer_tail.buffer[:, 2].any() - assert torch.equal(req.req_to_indexer_tail.buffer[:, [0, 1, 3]], unchanged) + assert not req.indexer_state_manager.state.buffer[:, 2].any() + assert torch.equal(req.indexer_state_manager.state.buffer[:, [0, 1, 3]], unchanged) + req.indexer_state_manager.state.buffer[:, 2].normal_() + req.reset(dest_req.req_idx) + assert not req.linear_state_manager.conv_state.buffer[:, 2].any() + assert not req.linear_state_manager.ssm_state.buffer[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() + assert not req.indexer_state_manager.state.buffer[:, 2].any() + assert torch.equal(req.indexer_state_manager.state.buffer[:, [0, 1, 3]], unchanged) # KV moves carry MLA latents and pooled FP8 bytes; raw keys/gates only # belong to live requests. Bytewise equality checks FP8 scale integrity. packed_bytes = mem.kv_buffer.view(torch.uint8) @@ -145,7 +161,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t indexes = torch.arange(8, dtype=torch.int32, device="cuda") zero = torch.zeros(1, dtype=torch.int64, device="cuda") ready = torch.zeros(1, dtype=torch.int32, pin_memory=True) - big = mem.big_page_buffers + big = req.big_page_buffers packed_bytes[:, :8].random_(0, 256) big.conv_state_cache.buffer[0].normal_() big.ssm_state_cache.buffer[0].normal_() diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 1a5e32567b..a72d15bc4c 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -596,7 +596,8 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): # A large capacity also exercises query chunking and uninitialized # logits beyond each query's actual pool range. max_kv_seq_len=max_kv_seq_len, - req_manager=SimpleNamespace(req_to_token_indexs=ragged[None], get_indexer_tail_buffer=lambda _: tail), + req_manager=SimpleNamespace(req_to_token_indexs=ragged[None]), + indexer_state_manager=SimpleNamespace(get_layer_buffer=lambda _: tail), b_req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), b_seq_len=torch.tensor([tokens], device="cuda", dtype=torch.int32), b1_cu_q_seq_len=torch.tensor([0, tokens], device="cuda", dtype=torch.int32), diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 00d4e326f0..2509d855b4 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -192,13 +192,14 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): split_qkv=lambda x: x.split(hidden, -1), ) manager = SimpleNamespace( - req_to_mtp_state_index=accepted_offsets, HOLD_REQUEST_ID=3, get_mamba_cache=lambda layer: (conv, ssm) + mtp_state_index=accepted_offsets, HOLD_REQUEST_ID=3, get_layer_buffers=lambda layer: (conv, ssm) ) infer = InferStateInfo() infer.batch_size = tokens infer.b_req_idx = torch.tensor(request_rows, device="cuda", dtype=torch.int32) infer.b_mtp_index = torch.tensor(offsets, device="cuda", dtype=torch.int32) infer.req_manager = manager + infer.linear_state_manager = manager weight = SimpleNamespace( get_merged_kda_conv_weight=lambda: conv_weight, linear_A_log=SimpleNamespace(weight=a), diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 57665cbbe3..b481e85763 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -1,10 +1,14 @@ +from lightllm.common.kv_page_transfer import KVPageTransfer +from lightllm.common.kv_transfer_resources import KVTransferView import dataclasses import pytest import torch from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, MemoryManager -from lightllm.common.req_manager import Glm5NextReqManager +from lightllm.common.cache_coordinator import HybridCacheCoordinator +from lightllm.common.linear_state import LinearStateManager +from lightllm.common.indexer_state import IndexerStateManager from lightllm.common.state_cache_manager import Glm5NextCacheConfig from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools from lightllm.server.core.objs.start_args_type import StartArgs @@ -31,14 +35,6 @@ def make_mems(monkeypatch, request): get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - def check_shm_refs(mem, req): - assert mem.req_to_indexer_tail is req.req_to_indexer_tail - assert mem.req_to_conv_state is req.req_to_conv_state - assert mem.req_to_ssm_state is req.req_to_ssm_state - assert mem.big_page_buffers is None - - monkeypatch.setattr(MemoryManager, "write_to_shm", check_shm_refs) - def create(tp): config = Glm5NextCacheConfig( tp_world_size=tp, @@ -61,14 +57,24 @@ def create(tp): draft_full_att_kv_layer_num=int(mtp_step > 0), ) mems = [] + states = [] for _ in range(tp): - mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0), config) - req = Glm5NextReqManager(3, 32, mem, config) - mem.write_to_shm(req) - assert mem.big_page_buffers is not None + mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0)) + mem.set_kv_layer_mapping(config.get_full_att_kv_layer_mapping()) + coordinator = HybridCacheCoordinator( + LinearStateManager(3, config, args), IndexerStateManager(3, config, mtp_step) + ) + coordinator.bind_mem_manager(mem) + assert not hasattr(mem, "req_to_conv_state") + assert not hasattr(mem, "req_to_ssm_state") + assert not hasattr(mem, "req_to_indexer_tail") + assert not hasattr(mem, "big_page_buffers") + states.append(coordinator.create_state_transfer()) mems.append(mem) - mems[0].alloc_paged_kv_move_buffer(1, 2048) - return mems + transfer = KVPageTransfer(KVTransferView.from_mem_manager(mems[0], include_page_layout=True)) + transfer.alloc_paged_kv_move_buffer(1, 2048) + states[0].bind_page_buffer(transfer.kv_move_buffer) + return mems, states, transfer return create @@ -76,7 +82,9 @@ def create(tp): @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) @pytest.mark.parametrize("remainder", range(4)) def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder): - source, dest = make_mems(prefill_tp), make_mems(decode_tp) + (source, source_states, source_transfer), (dest, dest_states, dest_transfer) = make_mems(prefill_tp), make_mems( + decode_tp + ) mtp_size = get_env_start_args().mtp_step + 1 layers = source[0].layer_num src_req, dst_req = 0, 2 @@ -89,47 +97,47 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem # Non-live slots deliberately contain stale data. The sequence length, not # the bytes in these slots, controls subsequent K-pool completion. tail = torch.randn(layers, 3 + mtp_size, 256, dtype=torch.bfloat16, device="cuda") - for rank, mem in enumerate(source): + for rank, (mem, state) in enumerate(zip(source, source_states)): heads = slice(rank * 8 // prefill_tp, (rank + 1) * 8 // prefill_tp) mem.kv_buffer.view(torch.uint8)[:, src_indexes] = packed - mem.req_to_conv_state.buffer[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) - mem.req_to_ssm_state.buffer[:, src_req * mtp_size].copy_(ssm[:, heads]) - mem.req_to_indexer_tail.buffer[:, src_req].copy_(tail) + state.req_to_conv_state.buffer[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) + state.req_to_ssm_state.buffer[:, src_req * mtp_size].copy_(ssm[:, heads]) + state.req_to_indexer_tail.buffer[:, src_req].copy_(tail) untouched = [] - for mem in dest: + for mem, state in zip(dest, dest_states): mem.kv_buffer.zero_() - mem.req_to_conv_state.buffer.normal_() - mem.req_to_ssm_state.buffer.normal_() - mem.req_to_indexer_tail.buffer.normal_() + state.req_to_conv_state.buffer.normal_() + state.req_to_ssm_state.buffer.normal_() + state.req_to_indexer_tail.buffer.normal_() untouched.append( [ x.buffer[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( - (mem.req_to_conv_state, 1), - (mem.req_to_ssm_state, mtp_size), - (mem.req_to_indexer_tail, 1), + (state.req_to_conv_state, 1), + (state.req_to_ssm_state, mtp_size), + (state.req_to_indexer_tail, 1), ) ] ) - source[0].write_mem_to_page_kv_move_buffer(src_indexes, 0, 0, source, prefill_tp) + source_transfer.write_mem_to_page_kv_move_buffer(src_indexes, 0, 0, source, prefill_tp) # Model a byte-preserving transport between separate P/D page buffers. - dest[0].kv_move_buffer.view(torch.uint8).copy_(source[0].kv_move_buffer.view(torch.uint8)) - dest[0].read_page_kv_move_buffer_to_mem(dst_indexes, 0, 0, dest, decode_tp) - source[0].write_mem_to_page_kv_move_buffer([], 0, 0, source, prefill_tp, "att_state", src_req) - dest[0].kv_move_buffer.view(torch.uint8).copy_(source[0].kv_move_buffer.view(torch.uint8)) - dest[0].read_page_kv_move_buffer_to_mem([], 0, 0, dest, decode_tp, "att_state", dst_req) + dest_transfer.kv_move_buffer.view(torch.uint8).copy_(source_transfer.kv_move_buffer.view(torch.uint8)) + dest_transfer.read_page_kv_move_buffer_to_mem(dst_indexes, 0, 0, dest, decode_tp) + source_states[0].write_req_to_page(0, src_req, source_states[0].get_dp_states(source_states, 0, prefill_tp)) + dest_transfer.kv_move_buffer.view(torch.uint8).copy_(source_transfer.kv_move_buffer.view(torch.uint8)) + dest_states[0].read_page_to_req(0, dst_req, dest_states[0].get_dp_states(dest_states, 0, decode_tp)) torch.cuda.synchronize() - for rank, mem in enumerate(dest): + for rank, (mem, state) in enumerate(zip(dest, dest_states)): heads = slice(rank * 8 // decode_tp, (rank + 1) * 8 // decode_tp) assert torch.equal(mem.kv_buffer.view(torch.uint8)[:, dst_indexes], packed) - assert torch.equal(mem.req_to_conv_state.buffer[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) - assert torch.equal(mem.req_to_ssm_state.buffer[:, dst_req * mtp_size], ssm[:, heads]) - assert torch.equal(mem.req_to_indexer_tail.buffer[:, dst_req], tail) + assert torch.equal(state.req_to_conv_state.buffer[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) + assert torch.equal(state.req_to_ssm_state.buffer[:, dst_req * mtp_size], ssm[:, heads]) + assert torch.equal(state.req_to_indexer_tail.buffer[:, dst_req], tail) for (state, stride), expected in zip( - ((mem.req_to_conv_state, 1), (mem.req_to_ssm_state, mtp_size), (mem.req_to_indexer_tail, 1)), + ((state.req_to_conv_state, 1), (state.req_to_ssm_state, mtp_size), (state.req_to_indexer_tail, 1)), untouched[rank], ): indexes = torch.tensor([0, 1, 3], device="cuda") * stride @@ -142,12 +150,15 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem ape = torch.randn(4, 128, device="cuda") tensor = lambda values: torch.tensor(values, dtype=torch.int32, device="cuda") outputs = [] - for mem, req_idx, indexes in ((source[0], src_req, src_indexes), (dest[0], dst_req, dst_indexes)): + for mem, state, req_idx, indexes in ( + (source[0], source_states[0], src_req, src_indexes), + (dest[0], dest_states[0], dst_req, dst_indexes), + ): new_indexes = [i for i in range(32) if i not in indexes][:count] all_indexes = indexes + new_indexes compress_pools( raw=raw, - tail=mem.req_to_indexer_tail.buffer[0], + tail=state.req_to_indexer_tail.buffer[0], packed_buffer=mem.get_indexer_k_buffer(3), ape=ape, lengths=tensor(range(length + 1, length + count + 1)), @@ -163,10 +174,10 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem def test_pd_page_capacity_includes_tail_without_resizing(make_mems): - mem = make_mems(1)[0] - helper = mem.att_state_page_helper + mems, states, transfer = make_mems(1) + mem, helper = mems[0], states[0] # A page that fits Conv+SSM but not the appended tail must be rejected. - mem.kv_move_buffer = torch.empty((1, helper.tail_offset), dtype=torch.uint8, device="cuda") + helper.kv_move_buffer = torch.empty((1, helper.tail_offset), dtype=torch.uint8, device="cuda") with pytest.raises(AssertionError, match="smaller than global linear att state"): helper.assert_page_size() shape = mem.get_paged_kv_move_buffer_shape(2, 2048) diff --git a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py index c68d7e3439..403de188ff 100644 --- a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py +++ b/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py @@ -4,6 +4,7 @@ import pytest import torch +from lightllm.common.cache_coordinator import CacheCoordinator from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager from lightllm.utils.kv_cache_utils import CpuKVCacheMeta @@ -85,6 +86,8 @@ def allocate(req, count): ), ), ) + module.backend.model.cache_coordinator = CacheCoordinator() + module.backend.model.cache_coordinator.bind_mem_manager(module.backend.model.mem_manager) monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 1024) monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) diff --git a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py index f04341a983..9d2e09cecf 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -1,3 +1,4 @@ +from lightllm.common.cache_coordinator import HybridCacheCoordinator from types import SimpleNamespace from sortedcontainers import SortedDict @@ -95,7 +96,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")) @@ -250,7 +250,13 @@ def test_case10(): 测试场景：测试 flush_cache 函数 """ print("\nTest Case 10: Testing flush_cache function\n") - tree = RadixCache(100, 0) + freed = [] + mem_manager = SimpleNamespace( + page_size=1, + allocator=SimpleNamespace(can_use_mem_size=95), + free=lambda indexes: freed.extend(indexes.tolist()), + ) + tree = RadixCache(100, 0, mem_manager=mem_manager) tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) tree.insert(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)) tree_node, size, values = tree.match_prefix( @@ -258,7 +264,9 @@ def test_case10(): ) assert tree_node is not None assert size == 3 + tree.dec_node_ref_counter(tree_node) tree.flush_cache() + assert sorted(freed) == [1, 2, 3, 4, 5] tree_node, size, values = tree.match_prefix( torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu"), update_refs=True ) @@ -409,14 +417,18 @@ def test_hybrid_match_refs_only_valid_prefix_after_tail_state_eviction(): small_page_buffers = SimpleNamespace( get_free_cache_num=lambda: 0, free_state_cache=lambda free_indexes: freed_state_ids.extend(free_indexes), + has_checkpoint=lambda slot, length: slot is not None, + ) + mem_manager = SimpleNamespace( + big_page_buffers=SimpleNamespace(has_checkpoint=lambda slot, length: slot is not None) ) - mem_manager = SimpleNamespace(big_page_buffers=SimpleNamespace()) tree = HybridAttPagedRadixCache( total_token_num=100, rank_in_node=101, hash_page_size=4, big_page_num=2, kv_cache_mem_manager=mem_manager, + big_page_buffers=mem_manager.big_page_buffers, small_page_buffers=small_page_buffers, ) key = torch.arange(12, dtype=torch.int64) @@ -457,10 +469,13 @@ def hybrid_tree(): rank_in_node=102, hash_page_size=4, big_page_num=2, - kv_cache_mem_manager=SimpleNamespace( - big_page_buffers=SimpleNamespace(free_state_cache=freed_big_states.extend) + kv_cache_mem_manager=SimpleNamespace(), + big_page_buffers=SimpleNamespace( + free_state_cache=freed_big_states.extend, has_checkpoint=lambda slot, length: slot is not None + ), + small_page_buffers=SimpleNamespace( + free_state_cache=lambda **kwargs: None, has_checkpoint=lambda slot, length: slot is not None ), - small_page_buffers=SimpleNamespace(free_state_cache=lambda **kwargs: None), ) return tree, freed_big_states @@ -494,7 +509,6 @@ def test_hybrid_match_traverses_pages_inside_image(hybrid_tree): tree.dec_node_ref_counter(node) assert tree.get_refed_tokens_num() == 0 - # Duplicate requests release their slots; no per-page validity metadata is needed. tree.insert(key[:8], key[:8] + 100, [1, 2], [None, None], SortedDict({8: 8})) tree.insert(key[:16], key[:16] + 100, [1, 2, 3, 4], [None] * 4, SortedDict({8: 9, 16: 10})) assert first_page.big_page_buffer_idx == 5 @@ -505,6 +519,55 @@ def test_hybrid_match_traverses_pages_inside_image(hybrid_tree): assert not tree.root_node.children +def test_hybrid_checkpoint_validity(hybrid_tree): + from lightllm.common.state_cache_manager.base import StateCacheManager + + class Pool(StateCacheManager): + def get_state_cache(self, buffer_idx): + raise NotImplementedError + + tree, freed = hybrid_tree + pool = Pool(3) + tree.big_page_buffers = pool + invalid, endpoint, replacement = pool.alloc_state_cache(3) + pool.checkpoint_lens[endpoint] = 16 + pool.checkpoint_lens[replacement] = 8 + free_slots = pool.free_state_cache + + def free_state_cache(slots): + freed.extend(slots) + free_slots(slots) + + pool.free_state_cache = free_state_cache + key = torch.arange(16) + tree.insert(key, key + 100, [1, 2, 3, 4], [None] * 4, SortedDict({8: invalid, 16: endpoint})) + first = tree.root_node.children[2] + second = first.children[4] + + node, length, values = tree.match_prefix(key[:8], block_hashs=[1, 2], update_refs=True, mm_input_spans=[]) + assert node is None and length == 0 and values is None + assert tree.get_refed_tokens_num() == 0 + node, length, values = tree.match_prefix(key, block_hashs=[1, 2, 3, 4], update_refs=True, mm_input_spans=[]) + assert node is second and length == 16 + assert values.tolist() == list(range(100, 116)) + assert tree.get_refed_tokens_num() == 16 + tree.dec_node_ref_counter(node) + + prefix, _ = tree.insert(key[:8], key[:8] + 200, [1, 2], [None] * 2, SortedDict({8: replacement})) + assert prefix == 8 and first.big_page_buffer_idx == replacement + assert freed == [invalid] + assert pool.checkpoint_lens[invalid] is None + node, length, values = tree.match_prefix(key[:8], block_hashs=[1, 2], update_refs=True) + assert node is first and length == 8 + assert values.tolist() == list(range(100, 108)) + tree.dec_node_ref_counter(node) + tree._evict(16, lambda *args: None) + assert sorted(freed) == [invalid, endpoint, replacement] + assert pool.get_free_cache_num() == 3 + assert pool.checkpoint_lens == [None] * 3 + assert tree.get_refed_tokens_num() == 0 and tree.root_node.ref_counter == 1 + + @pytest.mark.parametrize( "small_state_valid,mm_spans,match_len,expected", [ @@ -526,10 +589,15 @@ def test_hybrid_match_falls_back_along_matched_branch(monkeypatch, small_state_v rank_in_node=104, hash_page_size=64, big_page_num=4, - kv_cache_mem_manager=SimpleNamespace( - big_page_buffers=SimpleNamespace(free_state_cache=freed_big_states.extend) + kv_cache_mem_manager=SimpleNamespace(), + big_page_buffers=SimpleNamespace( + free_state_cache=freed_big_states.extend, has_checkpoint=lambda slot, length: slot is not None + ), + small_page_buffers=SimpleNamespace( + get_free_cache_num=lambda: 0, + free_state_cache=lambda **kwargs: None, + has_checkpoint=lambda slot, length: slot is not None, ), - small_page_buffers=SimpleNamespace(get_free_cache_num=lambda: 0, free_state_cache=lambda **kwargs: None), ) # The older request leaves a small-page state at 384, before the image. old_values = torch.arange(1000, 1384) @@ -600,6 +668,7 @@ def record_visit(node): @pytest.mark.parametrize("require_state", [False, True]) def test_hybrid_fallback_distinguishes_kv_ownership_from_state_restore(hybrid_tree, first_outside_image, require_state): tree, freed = hybrid_tree + tree.big_page_buffers.has_checkpoint = lambda slot, length: first_outside_image and length == 8 tree.insert( torch.arange(20), block_hashs=[1, 2, 3, 4, 5], @@ -610,7 +679,7 @@ def test_hybrid_fallback_distinguishes_kv_ownership_from_state_restore(hybrid_tr node, _, _ = tree.match_prefix( torch.arange(20), block_hashs=[1, 2, 3, 4, 5], update_refs=True, mm_input_spans=spans ) - node = tree.deref_to_first_big_page_node(node, mm_input_spans=spans if require_state else ()) + node = tree.deref_to_first_big_page_node(node, require_state=require_state) expected = 16 if not require_state else (8 if first_outside_image else 0) assert (0 if node is None else node.node_prefix_total_len) == expected assert tree.get_refed_tokens_num() == expected @@ -652,7 +721,11 @@ def test_hybrid_request_release_frees_or_transfers_image_page_slots( hybrid_token_hash_list=SimpleNamespace(get_all=lambda: block_hashs), ) released = [] - context._hybrid_att_free_req(released, req) + coordinator = HybridCacheCoordinator() + coordinator.bind_mem_manager(tree.mem_manager) + coordinator.small_page_buffers = tree.small_page_buffers + coordinator.big_page_buffers = tree.big_page_buffers + released.extend(coordinator.release_request_cache(req, context.req_manager, tree, context.args, publish_cache=True)) assert torch.cat(released).tolist() == list(range(100 + retained, 100 + max(24, computed))) assert not req.hybrid_len_to_big_page_id assert req.tail_small_page_buffer_id is None @@ -670,3 +743,42 @@ def test_hybrid_request_release_frees_or_transfers_image_page_slots( if __name__ == "__main__": pytest.main() + + +@pytest.mark.parametrize("publish_cache", [False, True]) +@pytest.mark.parametrize("duplicate", [False, True]) +def test_coordinator_releases_each_checkpoint_once(hybrid_tree, publish_cache, duplicate): + tree, freed_big = hybrid_tree + freed_small = [] + tree.small_page_buffers.free_state_cache = lambda free_indexes: freed_small.extend(free_indexes) + tree.insert(torch.arange(8), torch.arange(100, 108), [1, 2], [None, None], SortedDict({8: 5})) + if duplicate: + tree.insert(torch.arange(12), torch.arange(100, 112), [1, 2, 3], [None, None, 7], SortedDict()) + node, _, _ = tree.match_prefix(torch.arange(8), block_hashs=[1, 2], update_refs=True) + assert tree.get_refed_tokens_num() == 8 + coordinator = HybridCacheCoordinator() + coordinator.bind_mem_manager(tree.mem_manager) + coordinator.small_page_buffers = tree.small_page_buffers + coordinator.big_page_buffers = tree.big_page_buffers + req = SimpleNamespace( + req_idx=0, + cur_kv_len=12, + hold_kv_len=16, + hybrid_cache_len=12, + shared_kv_node=node, + tail_small_page_buffer_id=8, + hybrid_len_to_big_page_id=SortedDict({8: 6}), + get_input_token_ids=lambda: list(range(12)), + shm_req=SimpleNamespace(hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [1, 2, 3])), + ) + requests = SimpleNamespace(req_to_token_indexs=torch.arange(100, 116).view(1, -1)) + args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) + released = coordinator.release_request_cache(req, requests, tree, args, publish_cache) + assert sorted(torch.cat(released).tolist()) == list(range(108 if duplicate or not publish_cache else 112, 116)) + assert freed_big == [6] + assert freed_small == ([8] if duplicate or not publish_cache else []) + assert req.shared_kv_node is None + assert req.tail_small_page_buffer_id is None + assert not req.hybrid_len_to_big_page_id + assert tree.get_refed_tokens_num() == 0 + assert tree.get_tree_total_tokens_num() == (12 if publish_cache or duplicate else 8) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_shared_kv_trans.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_shared_kv_trans.py new file mode 100644 index 0000000000..ae9ef2a7b9 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_shared_kv_trans.py @@ -0,0 +1,211 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import torch + +from lightllm.server.router.model_infer.mode_backend import generic_pre_process +from lightllm.common.cache_coordinator import HybridCacheCoordinator +from lightllm.common.linear_state import LinearStateManager +from lightllm.common.state_cache_manager.base import StateCacheManager +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 + + +class _CheckpointPool(StateCacheManager): + def get_state_cache(self, buffer_idx): + raise AssertionError("This test only allocates checkpoint slots") + + +def _advance(tokens, conv, ssm, indexer): + for token in tokens: + value = torch.tensor([token, token * 2], dtype=torch.float64).reshape(2, 1, 1) + ssm.mul_(0.5).add_(value).add_(conv.sum(-1, keepdim=True) * 0.25) + conv.copy_(torch.cat((conv[..., 1:], value), dim=-1)) + indexer.copy_(torch.cat((indexer[:, 1:], value.square()), dim=1)) + + +def test_dp_fetch_cpu_transport_restores_and_releases_complete_request_state(monkeypatch): + args = SimpleNamespace(page_size=4, linear_att_hash_page_size=2, linear_att_page_block_num=2) + state = LinearStateManager.__new__(LinearStateManager) + state.linear_config = SimpleNamespace(get_conv_state_shape=lambda: (1, 2)) + state.mtp_step = 2 + state.mtp_state_index = torch.full((4,), 2, dtype=torch.int32) + state.conv_state = SimpleNamespace(buffer=torch.full((2, 4, 1, 4), -10.0, dtype=torch.float64)) + state.ssm_state = SimpleNamespace(buffer=torch.full((2, 12, 1, 1), -20.0, dtype=torch.float64)) + indexer = SimpleNamespace(state=SimpleNamespace(buffer=torch.full((2, 4, 2, 1), -30.0, dtype=torch.float64))) + coordinator = HybridCacheCoordinator(state, indexer) + coordinator.big_page_buffers = _CheckpointPool(4) + coordinator.small_page_buffers = _CheckpointPool(2) + tables = [torch.arange(96, dtype=torch.int32).reshape(4, 24) for _ in range(6)] + tables[1] = torch.full((4, 24), -1, dtype=torch.int32) + views = [SimpleNamespace(kv_buffer=torch.full((2, 128), -100.0)) for _ in range(6)] + peers = [ + SimpleNamespace( + mtp_step=2, + req_to_conv_state=SimpleNamespace(buffer=torch.full_like(state.conv_state.buffer, -40)), + req_to_ssm_state=SimpleNamespace(buffer=torch.full_like(state.ssm_state.buffer, -50)), + req_to_indexer_tail=SimpleNamespace(buffer=torch.full_like(indexer.state.buffer, -60)), + ) + for _ in range(6) + ] + next_slot = 0 + allocated = {} + + def allocate(req, count): + nonlocal next_slot + indexes = torch.arange(next_slot, next_slot + count, dtype=torch.int32) + tables[1][req.req_idx, req.hold_kv_len : req.hold_kv_len + count] = indexes + req.hold_kv_len += count + next_slot += count + allocated.setdefault(req.req_idx, []).extend(indexes.tolist()) + return indexes + + requests = [] + prefixes = [4, 8] + restored_lens = [12, 16] + source_ranks = [3, 5] + source_rows = [2, 3] + for req_idx, local_len, restored_len, rank, row in zip([0, 2], prefixes, restored_lens, source_ranks, source_rows): + tokens = list(range(1 + req_idx, restored_len + 8 + req_idx)) + shared_slot = coordinator.big_page_buffers.alloc_one_state_cache() + coordinator.big_page_buffers.checkpoint_lens[shared_slot] = local_len + req = SimpleNamespace( + req_idx=req_idx, + args=args, + cur_kv_len=local_len, + hold_kv_len=0, + shared_kv_node=SimpleNamespace( + node_prefix_total_len=local_len, big_page_buffer_idx=shared_slot, parent=None + ), + hybrid_cache_len=restored_len + 6, + hybrid_len_to_big_page_id={}, + tail_small_page_buffer_id=None, + shm_req=SimpleNamespace(shm_cur_kv_len=local_len), + get_cur_total_len=lambda tokens=tokens: len(tokens), + tokens=tokens, + mm_input_spans=[], + ) + req._kv_cache_alloc_need = lambda length, req=req: InferReq._kv_cache_alloc_need(req, length) + local_indexes = allocate(req, local_len) + kv = torch.tensor([tokens, [token * 2 for token in tokens]], dtype=torch.float32) + views[1].kv_buffer[:, local_indexes] = kv[:, :local_len] + source_indexes = tables[rank][row, :restored_len] + views[rank].kv_buffer[:, source_indexes] = kv[:, :restored_len] + peer = peers[rank] + conv = peer.req_to_conv_state.buffer[:, row, ..., :2].zero_() + ssm = peer.req_to_ssm_state.buffer[:, row * 3].zero_() + tail = peer.req_to_indexer_tail.buffer[:, row].zero_() + _advance(tokens[:restored_len], conv, ssm, tail) + requests.append(req) + + pool_baselines = [ + (pool.get_used_cache_num(), pool.checkpoint_lens.copy()) + for pool in (coordinator.big_page_buffers, coordinator.small_page_buffers) + ] + copies = [] + + def copy_kv(**kwargs): + assert [req.cur_kv_len for req in requests] == prefixes + assert kwargs["dp_size_in_node"] == 3 and kwargs["rank_in_dp"] == 1 + copies.append(kwargs) + for source, dp_rank, target in zip( + kwargs["move_token_indexes"], kwargs["token_dp_indexes"], kwargs["mem_indexes"] + ): + views[1].kv_buffer[:, target] = kwargs["kv_views"][int(dp_rank) * 2 + 1].kv_buffer[:, source] + + backend = SimpleNamespace( + model=SimpleNamespace( + mem_manager=SimpleNamespace(operator=SimpleNamespace(copy_kv_from_other_dp_ranks=copy_kv)), + cache_coordinator=coordinator, + ), + dp_world_size=2, + dp_size_in_node=3, + rank_in_dp=1, + node_nccl_group=None, + is_master_in_dp=True, + req_to_token_indexes_by_rank=tables, + kv_views=views, + peer_states=peers, + logger=Mock(), + _alloc_req_kv_mem=allocate, + ) + 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], [12, 2], [8, 1]], [[8, 2], [4, 1], [16, 3]]])) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=lambda: 128)) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) + tasks = module.build_shared_kv_trans_tasks(requests, req_dp_ranks=[0, 0]) + assert [task.max_kv_len_mem_manager_index for task in tasks] == source_ranks + assert [task.source_req_idx for task in tasks] == source_rows + tensor_to = torch.Tensor.to + tensor = torch.tensor + + def cpu_to(value, *args, **kwargs): + if kwargs.get("device") == "cuda": + kwargs["device"] = "cpu" + return tensor_to(value, *args, **kwargs) + + with monkeypatch.context() as transport: + transport.setattr(torch.Tensor, "to", cpu_to) + transport.setattr( + dp_shared_kv_trans, + "torch", + SimpleNamespace( + cat=torch.cat, + tensor=lambda data, **kwargs: tensor(data, dtype=kwargs["dtype"]), + int32=torch.int32, + int64=torch.int64, + ), + ) + module.kv_trans(tasks) + + assert len(copies) == 1 + assert state.mtp_state_index.tolist() == [0, 2, 0, 2] + radix = SimpleNamespace(insert=Mock(), dec_node_ref_counter=Mock()) + held_nodes = [req.shared_kv_node for req in requests] + freed = [] + for req, local_len, restored_len in zip(requests, prefixes, restored_lens): + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == restored_len + indexes = tables[1][req.req_idx, :restored_len] + tokens = req.tokens + torch.testing.assert_close( + views[1].kv_buffer[:, indexes], + tensor([tokens[:restored_len], [x * 2 for x in tokens[:restored_len]]], dtype=torch.float32), + ) + conv = state.conv_state.buffer[:, req.req_idx, ..., :2] + ssm = state.ssm_state.buffer[:, req.req_idx * 3] + tail = indexer.state.buffer[:, req.req_idx] + cold = (torch.zeros_like(conv), torch.zeros_like(ssm), torch.zeros_like(tail)) + _advance(tokens[:restored_len], *cold) + for actual, expected in zip((conv, ssm, tail), cold): + torch.testing.assert_close(actual, expected) + _advance(tokens[restored_len:], *cold) + _advance(tokens[restored_len:], conv, ssm, tail) + for actual, expected in zip((conv, ssm, tail), cold): + torch.testing.assert_close(actual, expected) + allocate(req, req._kv_cache_alloc_need(len(tokens))) + req.cur_kv_len = len(tokens) + big_slot = coordinator.big_page_buffers.alloc_one_state_cache() + small_slot = coordinator.small_page_buffers.alloc_one_state_cache() + req.hybrid_len_to_big_page_id[restored_len + 4] = big_slot + req.tail_small_page_buffer_id = small_slot + coordinator.big_page_buffers.checkpoint_lens[big_slot] = restored_len + 4 + coordinator.small_page_buffers.checkpoint_lens[small_slot] = req.hybrid_cache_len + assert coordinator.get_cacheable_prefix_len(req, args) == local_len + released = coordinator.release_request_cache( + req, SimpleNamespace(req_to_token_indexs=tables[1]), radix, args, True + ) + freed.extend(torch.cat(released).tolist()) + assert torch.cat(released).tolist() == allocated[req.req_idx][local_len:] + assert not req.hybrid_len_to_big_page_id + assert req.tail_small_page_buffer_id is req.shared_kv_node is None + assert len(freed) == len(set(freed)) + radix.insert.assert_not_called() + assert [call.args[0] for call in radix.dec_node_ref_counter.call_args_list] == held_nodes + for pool, (used, checkpoint_lens) in zip( + (coordinator.big_page_buffers, coordinator.small_page_buffers), pool_baselines + ): + assert pool.get_used_cache_num() == used + assert pool.checkpoint_lens == checkpoint_lens diff --git a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py index 7fcbb94827..d10b36cbd3 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 @@ -1,3 +1,4 @@ +from lightllm.common.cache_coordinator import HybridCacheCoordinator, CacheCoordinator from types import SimpleNamespace from collections import deque @@ -48,7 +49,9 @@ def test_non_gpu_cache_tiers_release_owned_tokens_without_radix_insert(): req = SimpleNamespace(req_idx=0, cur_kv_len=5, hold_kv_len=7, shared_kv_node=shared_node) free_token_indexes = [] - context._free_req_mem_without_radix_insert(free_token_indexes, req) + free_token_indexes.extend( + CacheCoordinator().release_request_cache(req, context.req_manager, context.radix_cache, None, False) + ) assert free_token_indexes[0].tolist() == [12, 13, 14, 15, 16] assert released_refs == [shared_node] @@ -60,7 +63,11 @@ 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.req_manager = None + context.args = None + context.cache_coordinator = SimpleNamespace( + release_request_cache=lambda req, *args, **kwargs: inserted_reqs.append(req) or [] + ) req = SimpleNamespace( cache_tiers=(CacheTier.GPU, CacheTier.CPU, CacheTier.DISK), cur_kv_len=3, @@ -131,7 +138,8 @@ def start_offload(req, cpu_kv_cache_stream): assert len(module.cpu_cache_handle_queue) == 2 -def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): +@pytest.mark.parametrize("with_radix", [False, True]) +def test_non_gpu_linear_cache_tiers_release_pending_state_pages(with_radix): freed_small_pages = [] freed_big_pages = [] context = InferenceContext() @@ -148,19 +156,31 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): shared_kv_node=None, tail_small_page_buffer_id=7, hybrid_len_to_big_page_id={128: 8, 256: 9}, + cache_tiers=(CacheTier.CPU,), + shm_req=SimpleNamespace(shm_cur_kv_len=3), ) free_token_indexes = [] - context._free_req_mem_without_radix_insert(free_token_indexes, req) + coordinator = HybridCacheCoordinator() + coordinator.bind_mem_manager(context.radix_cache) + coordinator.small_page_buffers = context.radix_cache.small_page_buffers + coordinator.big_page_buffers = context.radix_cache.big_page_buffers + context.cache_coordinator = coordinator + context.args = None + if not with_radix: + context.radix_cache = None + context.free_a_req_mem(free_token_indexes, req) 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 == {} + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == 0 -def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch): +@pytest.mark.parametrize("checkpoint_len", ["full", None, 128, 132]) +def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch, checkpoint_len): table = torch.full((1, 136), -1, dtype=torch.int32) table[0, :4] = torch.arange(4, dtype=torch.int32) next_index = 4 @@ -182,13 +202,17 @@ def alloc_req_kv_mem(req, alloc_token_num): return mem_indexes operator = SimpleNamespace( - load_cpu_cache_to_gpu=lambda mem_indexes, **kwargs: loaded_indexes.extend(mem_indexes.tolist()) + load_cpu_cache_to_gpu=lambda mem_indexes, **kwargs: loaded_indexes.extend(mem_indexes.tolist()), ) + coordinator = CacheCoordinator() if checkpoint_len == "full" else HybridCacheCoordinator() + coordinator.cpu_cache_transfer = SimpleNamespace(get_checkpoint_lens=lambda *args: [checkpoint_len]) + coordinator.load_cpu_cache_to_gpu = operator.load_cpu_cache_to_gpu backend = SimpleNamespace( is_master_in_dp=False, radix_cache=None, model=SimpleNamespace( mem_manager=SimpleNamespace(operator=operator), + cache_coordinator=coordinator, req_manager=SimpleNamespace(req_to_token_indexs=table), ), _alloc_req_kv_mem=alloc_req_kv_mem, @@ -219,9 +243,14 @@ def alloc_req_kv_mem(req, alloc_token_num): 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) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None)) module.load_cpu_cache_to_reqs([req]) + if checkpoint_len not in ("full", 132): + assert not alloc_sizes and not loaded_indexes + assert req.cur_kv_len == req.hold_kv_len == 4 + return assert alloc_sizes == [128] assert req.cur_kv_len == 132 assert req.hold_kv_len == 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 index 89e15fc93d..03dcdb01fc 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 @@ -121,7 +121,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=source_table)] * 2, + req_to_token_indexes_by_rank=[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) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py index 2e488ebb30..8ab560e2b0 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py @@ -200,6 +200,42 @@ def test_mode_proposals_own_their_schedule_metadata(): assert "schedule_scores_cpu" in DSparkSpecProposal.__dataclass_fields__ +@pytest.mark.parametrize("has_state", [False, True]) +def test_verify_mtp_state_provider(monkeypatch, has_state): + from lightllm.common.cache_coordinator import CacheCoordinator, HybridCacheCoordinator + + commit_args = {} + state = SimpleNamespace(commit_accepted_state=lambda **kwargs: commit_args.update(kwargs)) + coordinator = HybridCacheCoordinator(state_manager=state) if has_state else CacheCoordinator() + backend = SimpleNamespace( + max_draft_step=1, + model=SimpleNamespace( + cache_coordinator=coordinator, + req_manager=SimpleNamespace(req_sampling_params_manager=SimpleNamespace(req_to_next_token_ids=None)), + ), + ) + accept_lengths = torch.tensor([2, 1], dtype=torch.int32) + accepted_index = torch.tensor([True, True, True, False]) + monkeypatch.setattr(mtp_utils, "mtp_verify", lambda **kwargs: (accept_lengths, accepted_index)) + req_idx = torch.tensor([1, 1, 2, 2], dtype=torch.int32) + start_loc = torch.tensor([0, 2], dtype=torch.int32) + mtp_index = torch.tensor([0, 1, 0, 1], dtype=torch.int32) + + lengths, accepted = mtp_utils.verify_mtp_tokens( + backend, torch.tensor([10, 11, 12, 13]), req_idx, start_loc, mtp_index + ) + + assert lengths is accept_lengths and accepted is accepted_index + if has_state: + assert commit_args["b_req_idx"] is req_idx + assert commit_args["b_req_mtp_start_loc"] is start_loc + assert commit_args["b_mtp_index"] is mtp_index + assert commit_args["accepted_index"] is accepted_index + assert commit_args["verify_width"] == 2 + else: + assert not commit_args + + def test_scatter_mtp_next_tokens_consumes_mode_proposal(monkeypatch): scatter_args = {} monkeypatch.setattr( diff --git a/unit_tests/server/router/model_infer/structured_output/test_overlap.py b/unit_tests/server/router/model_infer/structured_output/test_overlap.py index 930cc5c1e8..947373c5ad 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_overlap.py +++ b/unit_tests/server/router/model_infer/structured_output/test_overlap.py @@ -67,7 +67,7 @@ def test_overlap_handoff_completes_postprocessing_before_sampling( impl.torch.cuda, "Event", lambda: SimpleNamespace(record=lambda: None, synchronize=lambda: None) ) monkeypatch.setattr(impl.g_infer_context, "get_overlap_stream", lambda: None) - monkeypatch.setattr(impl.g_infer_context, "save_hybrid_state_to_cache", lambda **kwargs: None) + monkeypatch.setattr(impl.g_infer_context, "save_cache_state", lambda **kwargs: None) monkeypatch.setattr(impl.g_infer_context, "is_hybrid_att_model", False) second_forward = Event() order, errors = [], [] diff --git a/unit_tests/server/router/model_infer/structured_output/test_state.py b/unit_tests/server/router/model_infer/structured_output/test_state.py index 4843d4dac3..46c19d8b03 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_state.py +++ b/unit_tests/server/router/model_infer/structured_output/test_state.py @@ -49,7 +49,8 @@ def test_context_initializes_state_factory_with_resolved_eos_at_startup(monkeypa ): context.register( backend=SimpleNamespace(dp_rank_in_node=0), - req_manager=SimpleNamespace(req_sampling_params_manager=None), + req_manager=SimpleNamespace(req_sampling_params_manager=None, mem_manager=None), + cache_coordinator=infer_batch.HybridCacheCoordinator(), radix_cache=None, shm_req_manager=None, vocab_size=257, diff --git a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py index 7b60c15c4a..0cba1e0d0b 100644 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -5,8 +5,10 @@ import pytest import torch -from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager -from lightllm.common.req_manager.linear_att import ReqManagerForMamba +from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.req_manager import ReqManager +from lightllm.common.cache_coordinator import HybridCacheCoordinator +from lightllm.common.linear_state import LinearStateManager from lightllm.common.state_cache_manager import LinearAttCacheConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.server.router.model_infer.infer_batch import g_infer_context @@ -56,9 +58,14 @@ def linear_cache(monkeypatch, tmp_path, request): full_attention_interval=2, all_layer_num=2, ) - mem = Qwen3NextMemManager(384, torch.bfloat16, 1, 8, 1, config) - requests = ReqManagerForMamba(2, 512, mem, config) - requests.create_small_page_cache_manager(2) + monkeypatch.setattr("lightllm.common.kv_cache_mem_manager.mem_manager.get_num_key_value_heads", lambda _: 1) + mem = MemoryManager(384, torch.bfloat16, 1, 8, 1) + mem.set_kv_layer_mapping(config.get_full_att_kv_layer_mapping()) + requests = ReqManager(2, 512, mem) + coordinator = HybridCacheCoordinator(LinearStateManager(2, config, args)) + coordinator.bind_mem_manager(mem) + coordinator.create_small_page_cache_manager(2) + monkeypatch.setattr(g_infer_context, "cache_coordinator", None) monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(g_infer_context, "req_manager", requests) monkeypatch.setattr( @@ -67,10 +74,10 @@ def linear_cache(monkeypatch, tmp_path, request): SimpleNamespace( get_big_page_ids_by_node=lambda node: [], free_one_small_page_buffer=lambda: None, - small_page_buffers=requests.small_page_buffers, + small_page_buffers=coordinator.small_page_buffers, ), ) - yield mem, requests, config + yield mem, requests, coordinator, config get_env_start_args.cache_clear() get_unique_server_name.cache_clear() @@ -83,22 +90,25 @@ def test_cpu_roundtrip_keeps_kv_continuous_and_restores_only_saved_states( from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule - mem, requests, config = linear_cache + mem, requests, coordinator, config = linear_cache disable_chunked_mm_input = get_env_start_args().disable_chunked_mm_input first_page_has_state = not disable_chunked_mm_input - big = mem.big_page_buffers + big = coordinator.big_page_buffers big_states = {} for end in range(128, token_num + 1, 128): slot = big.alloc_one_state_cache() big_states[end] = slot + if end != 128 or first_page_has_state: + big.checkpoint_lens[slot] = end for tensor in big.get_state_cache(slot): tensor.normal_() tail_slot = None if token_num % 128: - tail_slot = requests.small_page_buffers.alloc_one_state_cache() - for tensor in requests.small_page_buffers.get_state_cache(tail_slot): + tail_slot = coordinator.small_page_buffers.alloc_one_state_cache() + coordinator.small_page_buffers.checkpoint_lens[tail_slot] = token_num + for tensor in coordinator.small_page_buffers.get_state_cache(tail_slot): tensor.normal_() - final_states = requests.small_page_buffers.get_state_cache(tail_slot) + final_states = coordinator.small_page_buffers.get_state_cache(tail_slot) else: final_states = big.get_state_cache(big_states[token_num]) expected_states = [tensor.clone() for tensor in final_states] @@ -129,14 +139,16 @@ def test_cpu_roundtrip_keeps_kv_continuous_and_restores_only_saved_states( ), ) - mem.operator.offload_gpu_kv_to_cpu_cache(indexes, page_ids, torch.zeros_like(page_ids), client, req) + coordinator.offload_gpu_kv_to_cpu_cache( + indexes, page_ids, torch.zeros_like(page_ids), client, req, g_infer_context.radix_cache + ) torch.cuda.synchronize() module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) module.cpu_cache_client = client client.update_pages_status_to_ready(pages, deref=False) big.free_state_cache(list(big_states.values())) if tail_slot is not None: - requests.small_page_buffers.free_state_cache([tail_slot]) + coordinator.small_page_buffers.free_state_cache([tail_slot]) req.hybrid_len_to_big_page_id = {} req.tail_small_page_buffer_id = None req.cur_kv_len = req.hold_kv_len = 0 @@ -154,7 +166,7 @@ def allocate(req, count): module.backend = SimpleNamespace( is_master_in_dp=True, _alloc_req_kv_mem=allocate, - model=SimpleNamespace(mem_manager=mem, req_manager=requests), + model=SimpleNamespace(mem_manager=mem, req_manager=requests, cache_coordinator=coordinator), ) module.init_sync_group = None module.need_sync_compute_stream = lambda: False @@ -167,13 +179,14 @@ def allocate(req, count): assert sorted(req.hybrid_len_to_big_page_id) == list(range(128, token_num + 1, 128)) # Even the invalid page is copied normally; only restoring from it is forbidden. for end, slot in req.hybrid_len_to_big_page_id.items(): + assert big.has_checkpoint(slot, end) == (end != 128 or first_page_has_state) for actual, expected in zip(big.get_state_cache(slot), expected_big_states[end]): torch.testing.assert_close(actual, expected) - torch.testing.assert_close(requests.req_to_conv_state.buffer[:, 0].cpu(), expected_states[0]) - torch.testing.assert_close(requests.req_to_ssm_state.buffer[:, 0].cpu(), expected_states[1]) + torch.testing.assert_close(coordinator.linear_state_manager.conv_state.buffer[:, 0].cpu(), expected_states[0]) + torch.testing.assert_close(coordinator.linear_state_manager.ssm_state.buffer[:, 0].cpu(), expected_states[1]) # Loaded intermediate pages must remain unusable as radix restore points. - tree = HybridAttPagedRadixCache(384, 103, 64, 2, mem, requests.small_page_buffers) + tree = HybridAttPagedRadixCache(384, 103, 64, 2, mem, coordinator.small_page_buffers, coordinator.big_page_buffers) cache_len = max(req.hybrid_len_to_big_page_id) tree.insert( torch.arange(cache_len), @@ -186,7 +199,6 @@ def allocate(req, count): torch.arange(length), block_hashs=list(range(1, length // 64 + 1)), update_refs=True, - mm_input_spans=req.mm_input_spans, ) assert matched == expected tree.dec_node_ref_counter(node) @@ -197,3 +209,129 @@ def allocate(req, count): mm_input_spans=req.mm_input_spans, ) tree.dec_node_ref_counter(req.shared_kv_node) + + +def test_two_tail_loads_restore_before_reusing_temporary_slot(monkeypatch, linear_cache, cpu_cache_factory): + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common + from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule + + mem, requests, coordinator, config = linear_cache + client = cpu_cache_factory( + CpuKVCacheMeta(4, 1, 1, 1, config.get_cpu_cache_big_page_bytes(), torch.uint8, 0, torch.uint8), + ) + reqs, expected = [], [] + for index in range(2): + big_slot = coordinator.big_page_buffers.alloc_one_state_cache() + small_slot = coordinator.small_page_buffers.alloc_one_state_cache() + coordinator.big_page_buffers.checkpoint_lens[big_slot] = 128 + coordinator.small_page_buffers.checkpoint_lens[small_slot] = 192 + for state in coordinator.big_page_buffers.get_state_cache(big_slot): + state.fill_(index + 10) + states = coordinator.small_page_buffers.get_state_cache(small_slot) + for state in states: + state.fill_(index + 1) + expected.append([state.clone() for state in states]) + pages, _ = client.allocate_pages([100 + index * 2, 101 + index * 2], disk_offload_enable=False) + req = SimpleNamespace( + req_idx=index, + cur_kv_len=192, + hold_kv_len=192, + hybrid_cache_len=192, + hybrid_len_to_big_page_id={128: big_slot}, + tail_small_page_buffer_id=small_slot, + shared_kv_node=None, + mm_input_spans=[], + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + shm_req=SimpleNamespace( + input_len=256, + disk_prompt_cache_len=0, + cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda pages=pages: pages), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 192]), + ), + ) + indexes = torch.arange(index * 192, (index + 1) * 192, dtype=torch.int32, device="cuda") + mem.kv_buffer[:, index * 192 : (index + 1) * 192].fill_(index + 1) + page_ids = torch.tensor(pages, dtype=torch.int32, device="cuda") + coordinator.offload_gpu_kv_to_cpu_cache( + indexes, page_ids, torch.zeros_like(page_ids), client, req, g_infer_context.radix_cache + ) + torch.cuda.synchronize() + client.update_pages_status_to_ready(pages, deref=False) + coordinator.big_page_buffers.free_state_cache([big_slot]) + coordinator.small_page_buffers.free_state_cache([small_slot]) + req.hybrid_len_to_big_page_id = {} + req.tail_small_page_buffer_id = None + req.cur_kv_len = req.hold_kv_len = 0 + reqs.append(req) + + mem.kv_buffer.zero_() + + def allocate(req, count): + allocated = mem.alloc(count) + requests.req_to_token_indexs[req.req_idx, :count] = allocated.cuda() + req.hold_kv_len = count + return allocated + + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.cpu_cache_client = client + module.args = get_env_start_args() + module.backend = SimpleNamespace( + is_master_in_dp=True, + _alloc_req_kv_mem=allocate, + model=SimpleNamespace(mem_manager=mem, req_manager=requests, cache_coordinator=coordinator), + ) + module.init_sync_group = None + module.need_sync_compute_stream = lambda: False + monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 384) + monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) + events = [] + load = coordinator.load_cpu_cache_to_gpu + restore = coordinator.restore_state + + def track_load(**kwargs): + events.append(("copy", kwargs["req"].req_idx)) + assert load(**kwargs) is None + + def track_restore(req_idx, pool, slot, restored_len=None): + assert slot == coordinator.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID + events.append(("restore", req_idx)) + restore(req_idx, pool, slot, restored_len) + + monkeypatch.setattr(coordinator, "load_cpu_cache_to_gpu", track_load) + monkeypatch.setattr(coordinator, "restore_state", track_restore) + module.load_cpu_cache_to_reqs(reqs) + assert events == [("copy", 0), ("restore", 0), ("copy", 1), ("restore", 1)] + for index, req in enumerate(reqs): + state = coordinator.linear_state_manager + torch.testing.assert_close(state.conv_state.buffer[:, req.req_idx].cpu(), expected[index][0]) + torch.testing.assert_close(state.ssm_state.buffer[:, req.req_idx].cpu(), expected[index][1]) + assert req.cur_kv_len == req.hold_kv_len == 192 + assert (mem.kv_buffer[:, index * 192 : (index + 1) * 192] == index + 1).all() + + +def test_ipc_export_keeps_checkpoints_pinned_and_state_separate(linear_cache): + from multiprocessing.reduction import ForkingPickler + from lightllm.common.kv_transfer_resources import KVTransferResources, KVTransferView + + mem, requests, coordinator, _ = linear_cache + pools = (coordinator.big_page_buffers, coordinator.small_page_buffers) + tensors = [tensor for pool in pools for tensor in (pool.conv_state_cache.buffer, pool.ssm_state_cache.buffer)] + pointers = [tensor.data_ptr() for tensor in tensors] + assert all(tensor.is_pinned() for tensor in tensors) + mem.unrelated_resource = lambda: None + resources = KVTransferResources( + KVTransferView.from_mem_manager(mem), requests.req_to_token_indexs, coordinator.create_state_transfer() + ) + assert len(ForkingPickler.dumps(resources)) > 0 + assert all(tensor.is_pinned() for tensor in tensors) + assert [tensor.data_ptr() for tensor in tensors] == pointers + assert not hasattr(resources.kv, "big_page_buffers") + assert not hasattr(resources.kv, "allocator") + assert not hasattr(resources.kv, "operator") + assert not hasattr(resources.kv, "req_to_token_indexs") + assert not hasattr(resources.kv, "req_to_conv_state") + assert resources.req_to_token_indexes.data_ptr() == requests.req_to_token_indexs.data_ptr() + assert ( + resources.state_transfer.req_to_ssm_state.buffer.data_ptr() + == coordinator.linear_state_manager.ssm_state.buffer.data_ptr() + ) diff --git a/unit_tests/server/router/model_infer/test_multimodal_chunking.py b/unit_tests/server/router/model_infer/test_multimodal_chunking.py index 189101c950..247f61a384 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -1,3 +1,4 @@ +from lightllm.common.cache_coordinator import CacheCoordinator from types import SimpleNamespace import pytest @@ -13,6 +14,7 @@ def test_chunk_ids_and_lengths_use_the_same_complete_boundaries(hybrid, spans): req.args = SimpleNamespace( chunked_prefill_size=256, batch_max_tokens=1024, + disable_chunked_prefill=False, linear_att_hash_page_size=64, linear_att_page_block_num=8, ) @@ -75,17 +77,26 @@ def test_hybrid_chunks_save_state_only_at_original_boundaries( context.args = req.args context.is_hybrid_att_model = True context.radix_cache = SimpleNamespace( - big_page_buffers=SimpleNamespace(alloc_one_state_cache=lambda: next(slots)), - small_page_buffers=SimpleNamespace(alloc_one_state_cache=lambda: 7), + big_page_buffers=SimpleNamespace( + alloc_one_state_cache=lambda: next(slots), checkpoint_lens=[None] * len(page_ends) + ), + small_page_buffers=SimpleNamespace(alloc_one_state_cache=lambda: 7, checkpoint_lens=[None] * 8), free_one_small_page_buffer=lambda: None, ) - context.req_manager = SimpleNamespace( - save_big_page_states=lambda **kwargs: big_snapshots.append(req.get_chuncked_input_token_len()), - save_state=lambda **kwargs: small_snapshots.append(req.get_chuncked_input_token_len()), + from lightllm.common.cache_coordinator import HybridCacheCoordinator + + context.cache_coordinator = HybridCacheCoordinator( + state_manager=SimpleNamespace( + capture_intermediate_checkpoints=False, + save_checkpoints=lambda req_idx, reqs, slots, pool, lengths: big_snapshots.extend(lengths), + save_state=lambda req_idx, slot, pool, length: small_snapshots.append(length), + ) ) + context.cache_coordinator.big_page_buffers = context.radix_cache.big_page_buffers + context.cache_coordinator.small_page_buffers = context.radix_cache.small_page_buffers while req.cur_kv_len < req.get_cur_total_len(): - context.save_hybrid_state_to_cache(torch.tensor([0]), [req]) + context.save_cache_state(torch.tensor([0]), [req]) req.cur_kv_len = req.get_chuncked_input_token_len() assert big_snapshots == expected_big @@ -93,6 +104,10 @@ def test_hybrid_chunks_save_state_only_at_original_boundaries( assert sorted(req.hybrid_len_to_big_page_id) == page_ends assert len(set(req.hybrid_len_to_big_page_id.values())) == len(page_ends) assert big_snapshots == [end for end in page_ends if not any(start < end < stop for start, stop in spans)] + assert context.radix_cache.big_page_buffers.checkpoint_lens == [ + end if end in expected_big else None for end in page_ends + ] + assert context.radix_cache.small_page_buffers.checkpoint_lens[7] == (expected_small[0] if expected_small else None) @pytest.mark.parametrize( @@ -131,6 +146,7 @@ def test_full_attention_gpu_restore_stays_outside_image( req.shared_kv_node = None req.mm_input_spans = [(700, 1200)] req.shm_req = SimpleNamespace(input_len=len(query), shm_prompt_ids=SimpleNamespace(arr=query)) + monkeypatch.setattr(g_infer_context, "cache_coordinator", CacheCoordinator()) req._match_radix_cache() assert req.cur_kv_len == req.hold_kv_len == gpu_end assert not any(start < gpu_end < end for start, end in req.mm_input_spans) @@ -188,6 +204,7 @@ def test_mm_chunking_keeps_the_original_state_cache_position(monkeypatch, disabl linear_att_page_block_num=8, disable_chunked_mm_input=disable_chunked_mm_input, batch_max_tokens=1024, + disable_chunked_prefill=False, ) req.shm_index, req.vocab_size = 0, 100 req.multimodal_params = SimpleNamespace( @@ -269,10 +286,12 @@ def allocate(req, count): mem_manager=SimpleNamespace( operator=SimpleNamespace( load_cpu_cache_to_gpu=lambda **kwargs: loaded.append(req.cur_kv_len), + get_cpu_cache_match_page_num=lambda pages, *args: len(pages), ) ), ), ) + module.backend.model.cache_coordinator = module.backend.model.mem_manager.operator module.cpu_cache_client = client module.init_sync_group = None module.need_sync_compute_stream = lambda: False