diff --git a/docs/CN/source/framework/cache_manager_design.md b/docs/CN/source/framework/cache_manager_design.md new file mode 100644 index 0000000000..45e4112849 --- /dev/null +++ b/docs/CN/source/framework/cache_manager_design.md @@ -0,0 +1,389 @@ +# Attention Resources 架构设计 + +本文只定义 attention resource 的架构、公共接口和实现约束。具体模型参数、Tensor 形状、页大小、kernel 调优值和历史实现过程不在本文维护；这些内容属于代码和测试的实现细节。 + +## 1. 设计目标 + +`lightllm/common/attention_state/` 提供可组合的 attention 状态资源。资源层负责状态存储、布局描述、快照和跨设备数据操作；请求调度、radix cache、CPU cache、通信队列和页面生命周期由调用方负责。 + +设计必须满足： + +- 同一管理器可以同时持有 main、draft 及其他模型角色的 State。 +- full attention、bounded attention 和不同量化/编码方案通过独立 State 实现组合。 +- State 不拥有请求到物理 token 的映射，也不负责分配或回收外部页索引。 +- 共享 allocator 只管理索引和容量；State 只管理自己的 Tensor、布局和算法状态。 +- 新增算法优先新增独立 State 和创建入口，避免通过大量算法分支修改公共基类。 +- 所有跨模块数据操作都通过显式索引、长度、布局和资源引用完成，不从全局推理上下文隐式读取状态。 + +## 2. 架构和职责 + +```mermaid +flowchart TD + Model[模型装配] --> Manager[AttentionStateManager] + Manager --> Full[FullAttentionState] + Manager --> Bounded[BoundedAttentionState] + Manager --> Alloc[共享 allocator] + Request[ReqManager] --> Mapping[请求槽位与 token 映射] + Cache[Radix / CPU cache] --> Manager + Transport[DP / PD 调度] --> Manager + Forward[模型推理与 attention backend] --> Manager + Manager -.不拥有.-> Mapping +``` + +### 2.1 AttentionStateManager + +管理器负责： + +- 注册、查询和组织模型 State。 +- 按模型需要创建或绑定 full-attention、bounded-attention allocator。 +- 在容量确定后统一初始化 State 存储和大小页池。 +- 聚合多个 State 的重置、复制、保存、恢复及 CPU/DP/PD 操作。 +- 生成并缓存各 State 的传输布局。 +- 创建 PD/DP IPC 视图并重新绑定需要本地设备注册的页池。 + +管理器不负责： + +- 请求 KV 映射表、radix 树、缓存命中策略和页面引用计数。 +- 选择请求要使用的 token 或状态页索引。 +- 解释具体算法的 Tensor 字段、ring 位置或压缩规则。 +- 在 State 未提供能力时通过通用分支猜测其行为。 + +### 2.2 State + +每个 State 代表一种逻辑 attention 算法和存储编码方案。State 必须独立实现自己的配置、存储、计算视图、量化/编码、快照和恢复逻辑。 + +同一种逻辑算法的不同存储编码属于同一算法目录。例如： + +- `full_att/gqa/` 包含普通 GQA、INT4/INT8 分组量化和静态 FP8。 +- `full_att/mla/` 包含普通 MLA。 +- `full_att/packed_mla/` 包含 Packed MLA 及其 FP8/E8M0 变体。 +- `bounded_att/gdn/`、`bounded_att/kda/`、`bounded_att/sliding_window/` 和 `bounded_att/window_compressor/` 分别包含对应算法及其传输实现。 + +量化、压缩或 FP8 变体不单独建立只按精度命名的顶层目录；它们应放在所属逻辑算法目录中。 + +### 2.3 allocator、页池和请求映射 + +```text +ReqManager +├── request_slots +└── req_to_token_indexs + +AttentionStateManager +├── full_att_allocator +├── bounded_att_big_page_allocator +└── bounded_att_small_page_allocator + +State +├── runtime tensors +├── optional page pools +└── transfer layouts +``` + +- `ReqManager` 是请求槽位和 token 映射的唯一所有者。 +- `FullAttKVAllocator` 管理 full-attention token 索引。 +- `StatePageAllocator` 管理 bounded 状态页槽位。 +- 各页池只保存快照数据，不维护独立空闲列表和前缀引用关系。 +- 缓存调用方保存“前缀长度—页槽位”的关联，并在最后一个使用方结束后归还索引。 +- main/draft 可以共享 allocator；共享只表示索引域相同，不表示两个 State 的数据自动一致。 + +## 3. State 命名、注册和查询 + +### 3.1 State 类型名称 + +`get_state_type()` 返回稳定的小写 snake_case 算法和存储方案标记，不包含 `full_att`、`bounded_att`、模型角色或 `_state` 后缀。管理器生成： + +```python +state_name = f"{model_role}_{state.get_state_type()}_state" +``` + +名称必须能够区分影响布局或数据解释的方案参数，例如量化位宽、分组方式、scale 方案和编码精度。层号、容量、设备和请求信息不进入名称。 + +### 3.2 注册和查询 + +```python +manager.register_state(model_role, state) +manager.finish_state_registration() +manager.get_state( + model_role=model_role, + state_type_filter=state_type_filter, + logical_layer_id=logical_layer_id, + state_class=state_class, +) +``` + +注册完成后才允许查询。查询必须同时按模型角色、方案过滤条件、逻辑层归属和可选 State 类型筛选，最终只能匹配一个 State。`get_state` 可缓存；缓存只依赖注册期稳定信息，不依赖后续原地分配的 Tensor。 + +State 必须提供逻辑层归属查询和访问视图： + +```python +state.contains_logical_layer(logical_layer_id) +state.get_layer_view(logical_layer_id) +``` + +一个 State 可以服务多个逻辑层；共享层映射由具体配置描述，模型负责保证只有生产层写入共享数据。 + +## 4. 公共接口契约 + +以下接口是行为契约，具体配置参数和返回视图由实现定义。基类不承载具体算法字段。 + +### 4.1 FullAttentionState + +```python +class FullAttentionState: + def get_state_type(self) -> str: ... + def allocate_storage(self, token_capacity: int) -> None: ... + def get_token_bytes(self) -> int: ... + def contains_logical_layer(self, logical_layer_id: int) -> bool: ... + def get_storage_layer_index(self, logical_layer_id: int) -> int: ... + def get_layer_view(self, logical_layer_id: int) -> Any: ... + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: ... + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: ... + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: ... + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: ... + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: "FullAttentionState", + ) -> None: ... + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: ... + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: ... +``` + +所有表示 token 的索引参数由调用方准备为 CUDA Tensor。State 不在接口内部做设备、dtype 或索引格式转换，不分配请求映射。量化和 packed State 的复制必须保持其完整物理表示，包括必要的 scale、indexer 和 padding。 + +### 4.2 BoundedAttentionState + +```python +class BoundedAttentionState: + def get_state_type(self) -> str: ... + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: ... + def contains_logical_layer(self, logical_layer_id: int) -> bool: ... + def get_layer_view(self, logical_layer_id: int) -> Any: ... + def reset_bounded_att_state(self, req_indices: list[int]) -> None: ... + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: ... + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: ... + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: ... + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + def write_state_pages_to_cpu(...) -> None: ... + def read_state_pages_from_cpu(...) -> None: ... + def copy_big_pages_from_dp(...) -> None: ... + def copy_runtime_state_from_dp(...) -> None: ... + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: ... + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: ... + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: ... +``` + +`state_token_len` 表示状态已经执行到的绝对前缀长度，不表示当前搬运的 token 数。ring、窗口或游标位置必须由具体 State 根据该长度恢复。请求列表、长度列表和页槽位列表按位置一一对应，关联关系由调用方维护。 + +### 4.3 AttentionStateManager 聚合接口 + +管理器提供按资源组合调用 State 的入口，至少包括： + +```python +manager.initialize_storage(...) +manager.allocate_storage(...) +manager.get_full_attention_states(role=None) +manager.get_bounded_attention_states(role=None) +manager.get_cpu_page_layouts() +manager.get_pd_token_page_layouts() +manager.get_pd_runtime_page_layouts() +manager.copy_full_att_tokens(...) +manager.reset_bounded_att_state(...) +manager.save_bounded_att_runtime_to_pages(...) +manager.restore_bounded_att_runtime_from_pages(...) +manager.commit_bounded_att_state(...) +manager.write_req_to_cpu(...) +manager.read_req_from_cpu(...) +manager.copy_from_dp(...) +manager.write_to_pd_page(...) +manager.read_from_pd_page(...) +``` + +管理器只做聚合和公共前置检查；字段编码、ring 恢复、压缩边界和特殊传输由具体 State 实现。某种 State 不参与某条路径时返回空布局或由管理器跳过，不通过无意义的通用适配层伪造支持。 + +## 5. 布局和传输契约 + +### 5.1 传输类型 + +- `TransferRegion` 描述一个 State 字段在外部页中的范围、dtype、轴顺序、形状和 TP 映射。 +- `StateTransferLayout` 描述一次完整页面传输的统一页规格和区域集合。 +- `AxisMapping` 描述本地和外部轴的起点、长度及是否循环；循环窗口必须通过显式映射表达，不能依赖调用方猜测。 +- State 返回的布局定义必须稳定；使用 `lru_cache` 缓存，并通过独立副本返回，避免调用方修改缓存对象。 +- 存储分配或页面规格变化后，管理器负责清理相关布局缓存。 + +CPU、DP、PD 使用同一套布局语义，但入口不同： + +- CPU 传输可以处理联合页和 CPU 页就绪标记。 +- DP 传输使用源、目标 token 或页槽位的显式映射。 +- PD 传输必须接收明确的页索引；State 不分配网络页、不修改请求映射。 + +### 5.2 异步和设备约束 + +- State 的 GPU 操作提交到调用方当前流；State 不隐式同步设备或主机。 +- 传输调用方负责建立生产流、通信流、接收流和消费流之间的依赖。 +- NCCL/NIXL 的完成语义由传输层提供；读 PD 页前必须确认接收端数据已完成，不以发送端完成代替接收端完成。 +- 跨设备复制要求调用方保证 peer/IPC 可访问、源数据生命周期有效、索引和目标区域满足实现约束。 +- CPU 页的 pinned/registered 方式由缓存和进程启动策略决定，State 只使用传入的 buffer。 +- 需要 kernel 加速的聚合复制应保持异步语义，避免在公共 State 接口中加入主机同步。 + +### 5.3 数据所有权 + +- allocator 所有索引；State 所有运行态 Tensor 和页池 Tensor；ReqManager 所有请求映射；缓存流程所有引用关系。 +- 任何索引只能由其 allocator 归还一次。 +- IPC 浅拷贝不得复制 allocator；DP shared 场景必须提供页池重新注册到当前设备的显式流程。 +- 失败路径必须由调用方释放已分配索引并阻止未完成数据进入可见 State。 + +## 6. 模型初始化和装配 + +每种逻辑 attention 算法提供独立创建入口，量化和编码变体由参数选择或在同一逻辑目录中提供专用创建函数。模型基类只负责统一生命周期，不解析具体 State 的内部字段。 + +推荐流程： + +```text +模型读取配置 + -> 创建本模型拥有的 full/bounded State 列表 + -> manager.register_state(role, state) + -> 所有 main/draft 完成装配 + -> manager.finish_state_registration() + -> 绑定标定或其他外部参数 + -> manager.initialize_storage(...) + -> 模型和 backend 通过 manager.get_state(...) 获取视图 +``` + +模型可以让一个 State 服务多个逻辑层，也可以为同一模型创建多个逻辑方案。层映射必须由 State 配置完整描述，推理层不得用物理存储下标推断逻辑层归属。 + +## 7. 目录和代码组织 + +```text +lightllm/common/attention_state/ +├── manager.py +├── types.py +├── utils.py +├── allocators/ +├── full_att/ +│ ├── base_state.py +│ ├── dp_transfer.py +│ ├── gqa/ +│ ├── mla/ +│ ├── packed_mla/ +│ └── compressed_att/ +└── bounded_att/ + ├── base_state.py + ├── copy_kernel.py + ├── transfer_utils.py + ├── gdn/ + ├── kda/ + ├── sliding_window/ + └── window_compressor/ +``` + +每个逻辑 State 目录包含其配置、State、创建入口和专用传输实现；跨算法共享的基类、布局类型、allocator 和通用 kernel 保留在类别根目录。 + +代码编写约束： + +1. 公共接口放在私有辅助函数之前；类内依次放置初始化、公共属性/方法、私有辅助方法。 +2. 新增或大幅修改的函数，参数数量不少于 4 个时，声明和调用均使用逐行 `key=value` 格式；不为格式要求引入 `*args`、`**kwargs` 或转发字典。 +3. 配置字段必须有中文注释，说明含义、单位、Tensor 轴、全局/TP 范围和默认值语义。 +4. 公共接口 docstring 必须说明操作对象、索引含义、长度语义、数据范围和调用方责任。 +5. 具体 State 之间不继承实现；只继承对应类别的抽象基类，共享逻辑通过独立工具函数或明确的公共 kernel 提供。 +6. 不创建只有一次调用、只做参数转发或只包装断言的小函数；特殊算法行为放在对应 State 中。 +7. 已知输入约束在边界处使用简洁断言；不在公共接口中做隐式设备、dtype、索引或请求映射转换。 +8. 传输 kernel 的输入输出布局、stride、有效长度和循环映射必须在调用前明确；kernel 不负责外部资源分配和跨流同步。 +9. 修改现有基础流程时优先保持调用顺序和数据所有权不变；目录整理、命名调整和功能重构分开提交。 +10. 代码注释说明设计原因和调用契约，不重复逐行解释显而易见的 Python 语句，也不记录短期性能实验参数。 + +## 8. 测试和评审要求 + +新增或修改 State 至少验证： + +- 配置校验、逻辑层归属和共享层映射。 +- 存储容量、页池启用/禁用和 allocator 生命周期。 +- 运行态重置、保存、恢复及绝对状态长度语义。 +- CPU、DP、PD 的布局、索引方向和量化/编码字段完整性。 +- main/draft、多 TP、跨设备或 IPC 场景下的资源所有权。 +- 异步流依赖和失败路径，不以主机同步掩盖错误。 + +目录或导入路径调整必须同时运行： + +- `compileall` 和包导入检查。 +- 受影响 State 的单元测试。 +- 相关模型创建和资源布局测试。 +- `git diff --check`。 + +评审时优先检查：接口契约是否与基类一致、State 是否越权访问请求流程、布局是否可逆、索引是否可能重复归还、同步是否由正确的生产/消费流建立，以及量化变体是否仍归属于正确的逻辑算法目录。 + +## 9. 非目标 + +本文不规定： + +- 具体模型的 hidden size、head 数、压缩比例、窗口大小或页大小。 +- 某个 GPU、通信库或 Triton 版本的调优参数。 +- radix cache 的匹配策略、淘汰策略和外部磁盘格式。 +- 某个模型的 attention 数学公式或 kernel 内部实现。 +- 历史兼容路径和临时实验代码。 + +这些内容应在对应模型、State、kernel 或测试文件中维护，并通过公共接口满足本文契约。 diff --git a/docs/CN/source/index.rst b/docs/CN/source/index.rst index a0399f1623..f15caff1bf 100755 --- a/docs/CN/source/index.rst +++ b/docs/CN/source/index.rst @@ -81,6 +81,7 @@ Lightllm 整合了众多的开源方案的优点，包括但不限于 FasterTran :caption: 架构介绍 架构介绍 + Attention Resources 架构设计 token attention介绍 峰值显存调度器介绍 diff --git a/docs/CN/source/models/add_new_model.md b/docs/CN/source/models/add_new_model.md index 0eb8a0583c..7edb0c9701 100755 --- a/docs/CN/source/models/add_new_model.md +++ b/docs/CN/source/models/add_new_model.md @@ -109,7 +109,6 @@ class TpPartBaseModel: self._verify_must() self._verify_params() self._init_weights() - self._init_mem_manager() self._init_infer_layer() self._init_some_value() self._init_custom() @@ -124,10 +123,14 @@ class TpPartBaseModel: | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | 读取初始化模型的 config.json, 并进行一些 key 名的同名合法化操作 | | def _verify_params(self)： | 校验参数 | -| def _init_mem_manager(self): | 初始化 token attention 使用的 mem manager 对象 | +| `create_full_att_states(...)` / `create_bounded_att_states(...)` | 创建并返回当前模型的 full / bounded State 列表；默认返回空列表，公共流程统一注册 | | def _init_some_value(self): | 初始化推理框架会使用的一些成员变量的值 | | def _init_custom(self): | 一些模型自己的个性化初始化，比如 llama 初始化自己的Rotary值 | +全部 main/draft 模型构造完成后，`ModeBackend.init_model` 先按各模型 `get_req_manager_max_seq_len()` 的最大值创建一份 `ReqManager`，通过 `bind_req_manager()` 绑定给全部模型，再调用各模型的 `prepare_attention_states()` 准备自身 State；随后管理器依次完成 `finish_state_registration()`、`load_full_att_calibration(args)` 和共享资源分配。随后按 main、各 draft 的顺序逐个调用 `init_runtime()`，每次完整执行本模型的 HOLD 页和 backend 绑定、warmup、padded-request 与 CUDA Graph 初始化；main 还在该方法末尾执行最大长度检查。全部模型完成后发布就绪标志。资源准备和运行时初始化分别通过这两个单模型公共入口完成，均不接收 draft 列表。 + +公共 `prepare_attention_states()` 为 main 创建空管理器，draft 引用 main 的实例。两个 State 创建接口均为 `classmethod`，显式接收模型配置、逻辑层号及 TP 信息；bounded 接收请求容量和设备；full 工厂只创建布局，不计算标定起点或读取标定文件。接口仅返回对象列表，不注册管理器，也不依赖模型实例。CPU cache 直接复用 full/bounded 工厂生成 meta 定义，无需加载标定。公共流程统一注册；无 attention 的 draft 不创建 State。 + ## 2. 添加 bloom 模型的示例说明 具体实现在 ***lightllm/models/bloom*** 目录下，下面的代码片段请对应源码进行阅读，其中 triton_kernel 目录下为推理类使用的一些 kernel，下文中不做详细介绍，同时 bloom 模型因为不需要传递特殊状态信息使用默认的状态类即可。如想更深入的理解整个框架，可以进一步参考 llama 和 llama2 等模型的接入实现源码。 @@ -319,9 +322,6 @@ class BloomPreLayerInfer(PreLayerInferTpl): self.vob_end_id_ = tp_vocab_size_ * (self.tp_rank_ + 1) return - def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: - return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_) - def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BloomPreAndPostLayerWeight): total_token_num = infer_state.total_token_num input_ids = input_ids[0:total_token_num] @@ -346,6 +346,9 @@ class BloomPreLayerInfer(PreLayerInferTpl): dist.all_reduce(input_embdings, op=dist.ReduceOp.SUM, async_op=False) input_embdings = self._norm(input_embdings, infer_state, layer_weight) return input_embdings + + def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: + return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_) ~~~ ***transformer_layer_infer.py*** @@ -416,9 +419,10 @@ class BloomTransformerLayerInfer(TransformerLayerInferTpl): def _token_attention_kernel(self, q, infer_state:InferStateInfo, layer_weight: BloomTransformerLayerWeight)->torch.Tensor: o_tensor = torch.empty_like(q) + k_cache, v_cache = infer_state.full_att_state.get_layer_view(self.layer_num_) token_attention_fwd(q.view(-1, self.tp_q_head_num_, self.head_dim_), - infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0: self.tp_k_head_num_, :], - infer_state.mem_manager.kv_buffer[self.layer_num_][:, self.tp_k_head_num_: self.tp_k_head_num_+ self.tp_v_head_num_, :], + k_cache, + v_cache, o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_), layer_weight.tp_alibi, infer_state.b_loc, @@ -470,9 +474,6 @@ class BloomPostLayerInfer(PostLayerInferTpl): self.embed_dim_ = network_config["n_embed"] return - def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: - return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_) - def soft_max(self, data): return torch.softmax(data.permute(1, 0).float(), dim=-1) @@ -507,6 +508,9 @@ class BloomPostLayerInfer(PostLayerInferTpl): ans_logics = gather_data.permute(1, 0).float() gather_data = None return ans_logics + + def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: + return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_) ~~~ ### （3） 实现模型的框架类 diff --git a/docs/EN/source/models/add_new_model.md b/docs/EN/source/models/add_new_model.md index d9436acf08..cd71541598 100755 --- a/docs/EN/source/models/add_new_model.md +++ b/docs/EN/source/models/add_new_model.md @@ -109,7 +109,6 @@ class TpPartBaseModel: self._verify_must() self._verify_params() self._init_weights() - self._init_mem_manager() self._init_infer_layer() self._init_some_value() self._init_custom() @@ -124,10 +123,14 @@ Commonly used interfaces that need to be inherited and implemented | ---------------------------- | ------------------------------------------------------------ | | def _init_config(self)： | Read the config.json for initializing the model and perform some key name legalization operations | | def _verify_params(self)： | Validate parameters | -| def _init_mem_manager(self): | Initialize the mem manager object used by token attention | +| `create_full_att_states(...)` / `create_bounded_att_states(...)` | Return this model's full / bounded State lists; defaults return empty lists and the shared startup stage registers the results | | def _init_some_value(self): | Initialize values of some member variables that the inference framework will use | | def _init_custom(self): | Some personalized initialization of the model itself, such as llama initializing its own Rotary values | +After all main/draft models are constructed, `ModeBackend.init_model` creates one `ReqManager` using the largest `get_req_manager_max_seq_len()` returned by the models, binds it to every model through `bind_req_manager()`, then calls each model's `prepare_attention_states()`, then calls `finish_state_registration()`, `load_full_att_calibration(args)` and shared resource allocation in order. It next calls `init_runtime()` sequentially on the main model and each draft. Each call completes that model's HOLD-page and backend binding, warmup, padded-request setup and CUDA Graph initialization; the main model also checks maximum-length inference at the end of its call. Readiness is published after every model succeeds. Resource preparation and runtime initialization each use a single-model entry point; neither takes a draft-model list. + +The public `prepare_attention_states()` stage creates an empty manager for main and shares it with drafts. The two State factories are classmethods accepting model configuration, logical layer IDs and TP coordinates; bounded factories also accept request capacity and device. Full factories only define layouts and do not calculate calibration offsets or read calibration files. Factories only return lists and do not register resources or depend on model instances. CPU cache reuses the full/bounded factories to create meta definitions without loading calibration. The shared stage registers all returned States; attention-free drafts create none. + ## 2. Example of Adding Bloom Model The specific implementation is under the ***lightllm/models/bloom*** directory. Please read the source code for the code snippets below. The triton_kernel directory contains some kernels used by inference classes, which will not be introduced in detail in this article. At the same time, the bloom model uses the default state class because it doesn't need to pass special state information. For a deeper understanding of the entire framework, you can further refer to the source code implementation of llama and llama2 model integration. @@ -320,9 +323,6 @@ class BloomPreLayerInfer(PreLayerInferTpl): self.vob_end_id_ = tp_vocab_size_ * (self.tp_rank_ + 1) return - def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: - return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_) - def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BloomPreAndPostLayerWeight): total_token_num = infer_state.total_token_num input_ids = input_ids[0:total_token_num] @@ -347,6 +347,9 @@ class BloomPreLayerInfer(PreLayerInferTpl): dist.all_reduce(input_embdings, op=dist.ReduceOp.SUM, async_op=False) input_embdings = self._norm(input_embdings, infer_state, layer_weight) return input_embdings + + def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: + return layernorm_forward(input, layer_weight.pre_norm_weight_, layer_weight.pre_norm_bias_, eps=self.eps_) ~~~ ***transformer_layer_infer.py*** @@ -417,9 +420,10 @@ class BloomTransformerLayerInfer(TransformerLayerInferTpl): def _token_attention_kernel(self, q, infer_state:InferStateInfo, layer_weight: BloomTransformerLayerWeight)->torch.Tensor: o_tensor = torch.empty_like(q) + k_cache, v_cache = infer_state.full_att_state.get_layer_view(self.layer_num_) token_attention_fwd(q.view(-1, self.tp_q_head_num_, self.head_dim_), - infer_state.mem_manager.kv_buffer[self.layer_num_][:, 0: self.tp_k_head_num_, :], - infer_state.mem_manager.kv_buffer[self.layer_num_][:, self.tp_k_head_num_: self.tp_k_head_num_+ self.tp_v_head_num_, :], + k_cache, + v_cache, o_tensor.view(-1, self.tp_q_head_num_, self.head_dim_), layer_weight.tp_alibi, infer_state.b_loc, @@ -471,9 +475,6 @@ class BloomPostLayerInfer(PostLayerInferTpl): self.embed_dim_ = network_config["n_embed"] return - def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: - return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_) - def soft_max(self, data): return torch.softmax(data.permute(1, 0).float(), dim=-1) @@ -508,6 +509,9 @@ class BloomPostLayerInfer(PostLayerInferTpl): ans_logics = gather_data.permute(1, 0).float() gather_data = None return ans_logics + + def _norm(self, input, infer_state, layer_weight : BloomPreAndPostLayerWeight) -> torch.Tensor: + return layernorm_forward(input, layer_weight.final_norm_weight_, layer_weight.final_norm_bias_, eps=self.eps_) ~~~ ### （3） Add implementation model class diff --git a/lightllm/common/attention_state/__init__.py b/lightllm/common/attention_state/__init__.py new file mode 100644 index 0000000000..6c3f3107c1 --- /dev/null +++ b/lightllm/common/attention_state/__init__.py @@ -0,0 +1,5 @@ +"""Attention data resources, independent of request and cache policy.""" + +from .manager import AttentionStateManager + +__all__ = ["AttentionStateManager"] diff --git a/lightllm/common/attention_state/allocators/__init__.py b/lightllm/common/attention_state/allocators/__init__.py new file mode 100644 index 0000000000..036f71c997 --- /dev/null +++ b/lightllm/common/attention_state/allocators/__init__.py @@ -0,0 +1,4 @@ +from .state_page_allocator import StatePageAllocator +from .full_att_kv_allocator import KVCapacityReader, FullAttKVAllocator + +__all__ = ["StatePageAllocator", "KVCapacityReader", "FullAttKVAllocator"] diff --git a/lightllm/common/attention_state/allocators/full_att_kv_allocator.py b/lightllm/common/attention_state/allocators/full_att_kv_allocator.py new file mode 100644 index 0000000000..601b748553 --- /dev/null +++ b/lightllm/common/attention_state/allocators/full_att_kv_allocator.py @@ -0,0 +1,110 @@ +"""Full attention KV 分配器：接口使用 token 数量和索引，分配与回收须满足整页约束。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +import numpy as np +import torch + +from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt +from lightllm.utils.dist_utils import get_current_rank_in_node +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +if TYPE_CHECKING: + from collections.abc import Sequence + + +logger = init_logger(__name__) + + +class FullAttKVAllocator: + def __init__(self, size: int, page_size: int = 1, pin_memory: bool = True): + assert size >= 0 and page_size > 0 and size % page_size == 0, "KV capacity must contain complete pages" + assert size + page_size - 1 <= np.iinfo(np.int32).max, "KV token indices including HOLD must fit in int32" + self.size = size + self.page_size = page_size + self.mem_state = torch.arange( + end=self.size, + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) + self.mark_start = 0 + self.mark_end = self.size + self._mem_state_return = torch.empty( + size=(self.size * 3,), + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) + self._return_start = 0 + self.can_use_mem_size = self.size + rank_in_node = get_current_rank_in_node() + self.shared_can_use_token_num = SharedInt(f"kv_available_tokens_{rank_in_node}") + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + + @property + def hold_token_indices(self) -> tuple[int, ...]: + """One reserved page after the ordinary allocation domain.""" + return tuple(range(self.size, self.size + self.page_size)) + + def alloc(self, need_size: int) -> torch.Tensor: + assert need_size >= 0 and need_size % self.page_size == 0, "allocation must contain complete KV pages" + assert ( + need_size <= self.mark_end - self.mark_start + ), f"need {need_size} KV tokens, only {self.can_use_mem_size} available" + start = self.mark_start + end = start + need_size + self.mark_start = end + self.can_use_mem_size -= need_size + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + + # Reuse a rolling buffer so frees do not overwrite indices + # still being consumed by asynchronous H2D copies. + if self._return_start + need_size > self._mem_state_return.shape[0]: + self._return_start = 0 + ans = self._mem_state_return[self._return_start : self._return_start + need_size] + ans.copy_(self.mem_state[start:end]) + self._return_start += need_size + return ans + + def free(self, token_indices: Sequence[int] | torch.Tensor) -> None: + num_tokens = len(token_indices) + if num_tokens == 0: + return + assert num_tokens % self.page_size == 0, "free must contain complete KV pages" + end = self.mark_start + start = end - num_tokens + assert start >= 0, f"free {num_tokens} KV tokens exceeds allocated {end}" + if isinstance(token_indices, torch.Tensor): + self.mem_state[start:end].copy_(token_indices) + else: + self.mem_state.numpy()[start:end] = token_indices + self.mark_start = start + self.can_use_mem_size += num_tokens + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + if self.can_use_mem_size == self.size: + logger.debug(f"freed all gpu mem size {self.can_use_mem_size}") + + def free_all(self) -> None: + """Reset after all users of this allocation domain have finished.""" + self.mem_state.numpy()[:] = np.arange(self.size, dtype=np.int32) + self.mark_start = 0 + self.mark_end = self.size + self.can_use_mem_size = self.size + self.shared_can_use_token_num.set_value(self.can_use_mem_size) + + +class KVCapacityReader: + """Read each local DP group's allocator capacity without modifying it.""" + + def __init__(self): + args = get_env_start_args() + node_world_size = args.tp // args.nnodes + tp_group_size = args.tp // args.dp + self.is_multinode_tp = args.dp == 1 and args.nnodes > 1 + self.counters = [SharedInt(f"kv_available_tokens_{rank}") for rank in range(0, node_world_size, tp_group_size)] + + def get_available_tokens(self, dp_rank_in_node: int) -> int: + return self.counters[0 if self.is_multinode_tp else dp_rank_in_node].get_value() diff --git a/lightllm/common/attention_state/allocators/state_page_allocator.py b/lightllm/common/attention_state/allocators/state_page_allocator.py new file mode 100644 index 0000000000..9096bd9a96 --- /dev/null +++ b/lightllm/common/attention_state/allocators/state_page_allocator.py @@ -0,0 +1,39 @@ +"""状态页槽位分配器，仅管理槽位，不持有状态 buffer 或缓存策略。""" + +from __future__ import annotations + +from typing import List +from collections import deque + + +class StatePageAllocator: + def __init__(self, size: int, reserved_slots: int = 0): + self.size = size + self.reserved_slots = reserved_slots + if not 0 <= self.reserved_slots <= self.size: + raise ValueError("reserved_slots must be between zero and capacity") + self.capacity = self.size - self.reserved_slots + self._free = deque(range(self.capacity)) + self._allocated = bytearray(self.capacity) + + @property + def free_count(self) -> int: + return len(self._free) + + def alloc(self, num_pages: int) -> list[int] | None: + assert num_pages >= 0, "num_pages must be nonnegative" + if num_pages > self.free_count: + return None + slots = [self._free.popleft() for _ in range(num_pages)] + for slot in slots: + self._allocated[slot] = 1 + return slots + + def free(self, page_slot_ids: List[int]) -> None: + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert len(set(page_slot_ids)) == len(page_slot_ids), "cannot free duplicate slots" + assert all(0 <= slot < self.capacity for slot in page_slot_ids), "slot outside the allocatable domain" + assert all(self._allocated[slot] for slot in page_slot_ids), "cannot free an unallocated slot" + for slot in page_slot_ids: + self._allocated[slot] = 0 + self._free.extend(page_slot_ids) diff --git a/lightllm/common/attention_state/bounded_att/__init__.py b/lightllm/common/attention_state/bounded_att/__init__.py new file mode 100644 index 0000000000..55190bea6d --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/__init__.py @@ -0,0 +1,30 @@ +from .gdn.create_gdn_state import create_gdn_state +from .kda.create_kda_state import create_kda_state +from .conv_ssm_types import ConvSSMStateConfig +from .gdn.gdn_pages_pool import GDNPagesPool +from .gdn.gdn_state import GDNState +from .kda.kda_pages_pool import KDAPagesPool +from .kda.kda_state import KDAState +from .base_state import BoundedAttentionState + +from .sliding_window.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState +from .window_compressor.window_compressor_state import WindowCompressorStateConfig, WindowCompressorState +from .sliding_window.create_sliding_window_state import create_sliding_window_state +from .window_compressor.create_window_compressor_state import create_window_compressor_state + +__all__ = [ + "BoundedAttentionState", + "ConvSSMStateConfig", + "GDNState", + "GDNPagesPool", + "create_gdn_state", + "KDAState", + "KDAPagesPool", + "create_kda_state", + "FP8E8M0SlidingWindowConfig", + "FP8E8M0SlidingWindowState", + "create_sliding_window_state", + "WindowCompressorStateConfig", + "WindowCompressorState", + "create_window_compressor_state", +] diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py new file mode 100644 index 0000000000..6a96e76dc9 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/base_state.py @@ -0,0 +1,9 @@ +"""有界 attention 状态接口；具体算法实现运行态与大小页存储。""" + +from __future__ import annotations + +from .interface import CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface + + +class BoundedAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): + """单种 bounded-attention 方案的数据操作，索引和槽位由调用方分配。""" diff --git a/lightllm/common/attention_state/bounded_att/conv_ssm_types.py b/lightllm/common/attention_state/bounded_att/conv_ssm_types.py new file mode 100644 index 0000000000..aa0666eb3b --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_types.py @@ -0,0 +1,53 @@ +"""GDN/KDA 使用的 Conv/SSM 布局配置。""" + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class ConvSSMStateConfig: + """KDA/GDN 的层布局、维度、dtype 和 TP 配置；不包含模型角色或资源身份。""" + + # 本 State 管理的逻辑层编号；当前每个逻辑层独占一份运行态，列表顺序对应物理层轴，编号不能重复。 + logical_layer_ids: tuple[int, ...] + # TP 切分前的全局 K head 数，不计 TP 副本；Q 使用相同的 head 数和维度。 + global_k_heads: int + # TP 切分前的全局 V head 数；GDN 可与 K 不同，KDA 要求两者相等。 + global_v_heads: int + # 每个 Q/K head 的特征维度，也是 SSM 矩阵的 key 轴长度，不随 TP 切分。 + key_dim: int + # 每个 V head 的特征维度，也是 SSM 矩阵的 value 轴长度；KDA 要求与 key_dim 相等。 + value_dim: int + # Conv 需要保留的历史 token 数，等于卷积核宽度减 1；不含 MTP 额外候选位置。 + conv_width: int + # Conv 运行态及大小页快照的存储 dtype，不决定计算算子的中间精度。 + conv_dtype: torch.dtype + # SSM 运行态及大小页快照的存储 dtype，可与 Conv 不同，不决定算子的累加精度。 + ssm_dtype: torch.dtype + # 当前进程在本模型 TP 组中的编号，范围为 [0, tp_world_size)，用于确定本地 head 分片。 + tp_rank: int = 0 + # 本模型的 TP 并行度；head 足够时均分，少于 TP 并行度时按组复制；1 表示不切分。 + tp_world_size: int = 1 + + @property + def num_k_heads(self) -> int: + """当前 TP rank 的 K head 数；head 少于 TP 时按组复制。""" + return max(1, self.global_k_heads // self.tp_world_size) + + @property + def num_v_heads(self) -> int: + """当前 TP rank 的 V head 数，独立于 K head 的切分计算。""" + return max(1, self.global_v_heads // self.tp_world_size) + + @property + def conv_channels(self) -> int: + """当前 TP rank 的 Conv 通道总数，按 Q、K、V 三段拼接。""" + return 2 * self.num_k_heads * self.key_dim + self.num_v_heads * self.value_dim + + def __post_init__(self) -> None: + assert 0 <= self.tp_rank < self.tp_world_size + for heads in (self.global_k_heads, self.global_v_heads): + assert heads > 0 and ( + heads % self.tp_world_size == 0 or self.tp_world_size % heads == 0 + ), "Conv/SSM heads require complete TP shards or replicas" diff --git a/lightllm/common/attention_state/bounded_att/copy_kernel.py b/lightllm/common/attention_state/bounded_att/copy_kernel.py new file mode 100644 index 0000000000..b66d3eb11b --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/copy_kernel.py @@ -0,0 +1,491 @@ +"""将 bounded attention 的同一形状的多页视图聚合到 CUDA stream 拷贝。""" + +import torch +import triton +import triton.language as tl + +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +# ======================================================================== +# 一、聚合入口：只上传指针表，shape/stride 作为常量 +# ======================================================================== + + +def copy_bounded_att_tensors( + copies: list[tuple[torch.Tensor, torch.Tensor]], + grid_num: int = 24, + num_warps: int = 8, + num_stages: int = 1, +) -> None: + """在当前 CUDA stream 上聚合复制同一形状的视图，默认 24 个 program；warp 数与循环流水深度可按搬运场景覆盖。 + + num_stages 控制 tl.range 的循环流水深度，默认 1；提高该值不保证所有方向的带宽都增加。 + 支持 CUDA/IPC 存储和本进程已 pinned/registered 的 CPU 存储，两端均为 + CPU 页池时也走 GPU。整批所有源、目标的形状和 dtype 必须相同，支持最多六维的 stride。 + 同一侧的全部视图共用 stride，源与目标之间可不同；入口只核对首尾项。 + 只合并两端共同连续的尾部维度，并按对齐条件创建最大 8 字节的整数视图。 + 这些 view 不复制数据；不满足合并或宽类型条件时保留原布局，不转换数值。 + 每项表示独立复制，目标区域不得相互覆盖，也不得覆盖本批仍需读取的源数据。 + 直接使用 data_ptr()；调用方保证地址可被当前 GPU 访问，并持有存储直到拷贝完成。 + 源数据就绪，以及 CPU 读取或复用页面前的完成事件，也由调用方管理。 + """ + if not copies: + return + assert grid_num > 0, "copy grid must be positive" + # 每批使用相同 shape/dtype 和同侧 stride；只核对首尾，保留两侧的布局差异。 + source, destination = copies[0] + last_source, last_destination = copies[-1] + assert source.shape == destination.shape == last_source.shape == last_destination.shape, "copy shapes must match" + assert source.dtype == destination.dtype == last_source.dtype == last_destination.dtype, "copy dtypes must match" + assert 1 <= source.ndim <= 6, "copy tensors must have one to six dimensions" + assert source.numel() > 0, "copy tensors must be nonempty" + assert source.stride() == last_source.stride(), "source strides must match across the batch" + assert destination.stride() == last_destination.stride(), "destination strides must match across the batch" + # 阶段一：只合并两端共同连续的尾部维度，保留前面的布局。 + copies = merge_contiguous_tail_dims(copies=copies) + # 阶段二：按整批地址和行布局选择最宽的搬运单元，不转换数值。 + copies = view_max_copy_dtype(copies=copies) + # 阶段三：上传两张指针表，按最终视图维数发射聚合拷贝算子。 + _launch_copy_kernel( + copies=copies, + grid_num=grid_num, + num_warps=num_warps, + num_stages=num_stages, + ) + + +# ======================================================================== +# 二、聚合算子发射：上传指针与共享布局 +# ======================================================================== + + +def _launch_copy_kernel( + copies: list[tuple[torch.Tensor, torch.Tensor]], + grid_num: int, + num_warps: int, + num_stages: int, +) -> None: + """使用整理后的视图地址和共享布局，在当前流上发射一次聚合拷贝。""" + # 确保当前设备上下文已经建立；纯 CPU 页池之间的复制也使用这个设备和流。 + stream = torch.cuda.current_stream() + source, destination = copies[0] + source_pointers, destination_pointers = [], [] + for src_tensor, dst_tensor in copies: + source_pointers.append(src_tensor.data_ptr()) + destination_pointers.append(dst_tensor.data_ptr()) + + # 整批共享 shape/dtype 和同侧 stride，首对最终视图提供全部编译期布局参数。 + ndim = source.ndim + shape = tuple(source.shape) + src_stride = source.stride() + dst_stride = destination.stride() + unsigned_dtypes = {1: tl.uint8, 2: tl.uint16, 4: tl.uint32, 8: tl.uint64} + + kernels = { + 1: _copy_bounded_att_1d, + 2: _copy_bounded_att_2d, + 3: _copy_bounded_att_3d, + 4: _copy_bounded_att_4d, + 5: _copy_bounded_att_5d, + 6: _copy_bounded_att_6d, + } + # 源、目标地址分别上传为独立 Tensor，同一位置对应一对复制视图。 + source_pointer_tensor = torch.tensor( + data=source_pointers, + dtype=torch.int64, + device="cpu", + pin_memory=True, + ).to(device=stream.device, non_blocking=True) + destination_pointer_tensor = torch.tensor( + data=destination_pointers, + dtype=torch.int64, + device="cpu", + pin_memory=True, + ).to(device=stream.device, non_blocking=True) + # 固定 grid=24 的跨场景调优使用每轮 1024 个搬运单元，兼顾小块延迟和大块吞吐。 + kernels[ndim][(grid_num,)]( + SRC_POINTERS=source_pointer_tensor, + DST_POINTERS=destination_pointer_tensor, + total_element_count=source.numel() * len(copies), + SHAPE=shape, + SRC_STRIDES=src_stride, + DST_STRIDES=dst_stride, + DTYPE=unsigned_dtypes[source.element_size()], + BLOCK_ELEMENTS=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, + ) + + +# ======================================================================== +# 三、分维数聚合算子：计算索引与偏移 → 读取指针 → load 数据 → store 数据 +# ======================================================================== + + +@triton.jit +def _copy_bounded_att_1d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 一维直接按元素序号应用两端 stride。 + axis0 = element_indices + source_offsets = axis0 * src_stride_0 + destination_offsets = axis0 * dst_stride_0 + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_2d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 2 维坐标；两端可以具有不同的物理 stride。 + axis1 = element_indices % SHAPE[1] + remaining_indices = element_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = axis0 * src_stride_0 + axis1 * src_stride_1 + destination_offsets = axis0 * dst_stride_0 + axis1 * dst_stride_1 + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_3d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] * SHAPE[2] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + src_stride_2 = tl.cast(input=SRC_STRIDES[2], dtype=tl.int64) + dst_stride_2 = tl.cast(input=DST_STRIDES[2], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 3 维坐标；两端可以具有不同的物理 stride。 + axis2 = element_indices % SHAPE[2] + remaining_indices = element_indices // SHAPE[2] + axis1 = remaining_indices % SHAPE[1] + remaining_indices = remaining_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = axis0 * src_stride_0 + axis1 * src_stride_1 + axis2 * src_stride_2 + destination_offsets = axis0 * dst_stride_0 + axis1 * dst_stride_1 + axis2 * dst_stride_2 + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_4d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] * SHAPE[2] * SHAPE[3] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + src_stride_2 = tl.cast(input=SRC_STRIDES[2], dtype=tl.int64) + dst_stride_2 = tl.cast(input=DST_STRIDES[2], dtype=tl.int64) + src_stride_3 = tl.cast(input=SRC_STRIDES[3], dtype=tl.int64) + dst_stride_3 = tl.cast(input=DST_STRIDES[3], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 4 维坐标；两端可以具有不同的物理 stride。 + axis3 = element_indices % SHAPE[3] + remaining_indices = element_indices // SHAPE[3] + axis2 = remaining_indices % SHAPE[2] + remaining_indices = remaining_indices // SHAPE[2] + axis1 = remaining_indices % SHAPE[1] + remaining_indices = remaining_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = axis0 * src_stride_0 + axis1 * src_stride_1 + axis2 * src_stride_2 + axis3 * src_stride_3 + destination_offsets = axis0 * dst_stride_0 + axis1 * dst_stride_1 + axis2 * dst_stride_2 + axis3 * dst_stride_3 + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_5d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] * SHAPE[2] * SHAPE[3] * SHAPE[4] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + src_stride_2 = tl.cast(input=SRC_STRIDES[2], dtype=tl.int64) + dst_stride_2 = tl.cast(input=DST_STRIDES[2], dtype=tl.int64) + src_stride_3 = tl.cast(input=SRC_STRIDES[3], dtype=tl.int64) + dst_stride_3 = tl.cast(input=DST_STRIDES[3], dtype=tl.int64) + src_stride_4 = tl.cast(input=SRC_STRIDES[4], dtype=tl.int64) + dst_stride_4 = tl.cast(input=DST_STRIDES[4], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 5 维坐标；两端可以具有不同的物理 stride。 + axis4 = element_indices % SHAPE[4] + remaining_indices = element_indices // SHAPE[4] + axis3 = remaining_indices % SHAPE[3] + remaining_indices = remaining_indices // SHAPE[3] + axis2 = remaining_indices % SHAPE[2] + remaining_indices = remaining_indices // SHAPE[2] + axis1 = remaining_indices % SHAPE[1] + remaining_indices = remaining_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = ( + axis0 * src_stride_0 + + axis1 * src_stride_1 + + axis2 * src_stride_2 + + axis3 * src_stride_3 + + axis4 * src_stride_4 + ) + destination_offsets = ( + axis0 * dst_stride_0 + + axis1 * dst_stride_1 + + axis2 * dst_stride_2 + + axis3 * dst_stride_3 + + axis4 * dst_stride_4 + ) + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) + + +@triton.jit +def _copy_bounded_att_6d( + SRC_POINTERS, + DST_POINTERS, + total_element_count, + SHAPE: tl.constexpr, + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + DTYPE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + # 阶段一：grid 按整批总元素数跳步，每轮处理 BLOCK_ELEMENTS 个元素。 + # 先转 int64 再计算起点和步长，避免大批量传输时索引溢出。 + total_element_count = tl.cast(input=total_element_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK_ELEMENTS + element_lanes = tl.arange(start=0, end=BLOCK_ELEMENTS) + elements_per_copy = SHAPE[0] * SHAPE[1] * SHAPE[2] * SHAPE[3] * SHAPE[4] * SHAPE[5] + # stride 在参与乘加前显式转为 int64，避免大跨度地址计算发生 32 位溢出。 + src_stride_0 = tl.cast(input=SRC_STRIDES[0], dtype=tl.int64) + dst_stride_0 = tl.cast(input=DST_STRIDES[0], dtype=tl.int64) + src_stride_1 = tl.cast(input=SRC_STRIDES[1], dtype=tl.int64) + dst_stride_1 = tl.cast(input=DST_STRIDES[1], dtype=tl.int64) + src_stride_2 = tl.cast(input=SRC_STRIDES[2], dtype=tl.int64) + dst_stride_2 = tl.cast(input=DST_STRIDES[2], dtype=tl.int64) + src_stride_3 = tl.cast(input=SRC_STRIDES[3], dtype=tl.int64) + dst_stride_3 = tl.cast(input=DST_STRIDES[3], dtype=tl.int64) + src_stride_4 = tl.cast(input=SRC_STRIDES[4], dtype=tl.int64) + dst_stride_4 = tl.cast(input=DST_STRIDES[4], dtype=tl.int64) + src_stride_5 = tl.cast(input=SRC_STRIDES[5], dtype=tl.int64) + dst_stride_5 = tl.cast(input=DST_STRIDES[5], dtype=tl.int64) + # 循环级流水深度：允许编译器重叠相邻迭代的访存，与 launch 的 num_stages 选项不同。 + for block_start in tl.range(program_start, total_element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段二：全局索引拆成复制项编号和项内元素索引，再还原各维坐标。 + # 一个 block 可以跨越多对 Tensor；尾部 mask 同时保护指针表和数据访问。 + global_indices = block_start + element_lanes + valid_elements = global_indices < total_element_count + copy_index = global_indices // elements_per_copy + element_indices = global_indices % elements_per_copy + # 从最内层轴开始还原 6 维坐标；两端可以具有不同的物理 stride。 + axis5 = element_indices % SHAPE[5] + remaining_indices = element_indices // SHAPE[5] + axis4 = remaining_indices % SHAPE[4] + remaining_indices = remaining_indices // SHAPE[4] + axis3 = remaining_indices % SHAPE[3] + remaining_indices = remaining_indices // SHAPE[3] + axis2 = remaining_indices % SHAPE[2] + remaining_indices = remaining_indices // SHAPE[2] + axis1 = remaining_indices % SHAPE[1] + remaining_indices = remaining_indices // SHAPE[1] + axis0 = remaining_indices + source_offsets = ( + axis0 * src_stride_0 + + axis1 * src_stride_1 + + axis2 * src_stride_2 + + axis3 * src_stride_3 + + axis4 * src_stride_4 + + axis5 * src_stride_5 + ) + destination_offsets = ( + axis0 * dst_stride_0 + + axis1 * dst_stride_1 + + axis2 * dst_stride_2 + + axis3 * dst_stride_3 + + axis4 * dst_stride_4 + + axis5 * dst_stride_5 + ) + + # 阶段三：按元素所属复制项，从两张指针表读取源/目标地址，按元素宽度解释指针。 + # 使用等宽无符号类型保持原始位模式，不进行浮点数值转换。 + source_address = tl.load(pointer=SRC_POINTERS + copy_index, mask=valid_elements, other=0) + destination_address = tl.load(pointer=DST_POINTERS + copy_index, mask=valid_elements, other=0) + source_ptr = source_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + destination_ptr = destination_address.to(dtype=tl.pointer_type(element_ty=DTYPE)) + + # 阶段四：读取源数据并写入目标，末尾越界元素不参与 load/store。 + values = tl.load(pointer=source_ptr + source_offsets, mask=valid_elements, other=0) + tl.store(pointer=destination_ptr + destination_offsets, value=values, mask=valid_elements) diff --git a/lightllm/common/attention_state/bounded_att/gdn/__init__.py b/lightllm/common/attention_state/bounded_att/gdn/__init__.py new file mode 100644 index 0000000000..62dc94c11d --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/__init__.py @@ -0,0 +1,5 @@ +from .create_gdn_state import create_gdn_state +from .gdn_pages_pool import GDNPagesPool +from .gdn_state import GDNState + +__all__ = ["create_gdn_state", "GDNPagesPool", "GDNState"] diff --git a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py new file mode 100644 index 0000000000..3fe6f2317e --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py @@ -0,0 +1,23 @@ +"""GDN State 的统一创建入口。""" + +from __future__ import annotations + +import torch + +from ..conv_ssm_types import ConvSSMStateConfig +from .gdn_state import GDNState + + +def create_gdn_state( + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", +) -> GDNState: + """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" + return GDNState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_pages_pool.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_pages_pool.py new file mode 100644 index 0000000000..a54332eb84 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_pages_pool.py @@ -0,0 +1,44 @@ +"""GDN pages with separate Q/K and V dimensions; allocators own the page slots.""" + +from __future__ import annotations + + +import torch + +from ..conv_ssm_types import ConvSSMStateConfig +from ..utils import create_page_pool_tensor, register_page_pool_tensor + + +class GDNPagesPool: + def __init__( + self, + size: int, + config: ConvSSMStateConfig, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + assert size >= 0, "pool size must be nonnegative" + self.size = size + device = torch.device(device) + pin_memory = pin_memory and device.type == "cpu" + self.conv_state = create_page_pool_tensor( + size=(size, len(config.logical_layer_ids), config.conv_channels, config.conv_width), + dtype=config.conv_dtype, + device=device, + pin_memory=pin_memory, + ) + self.ssm_state = create_page_pool_tensor( + size=(size, len(config.logical_layer_ids), config.num_v_heads, config.key_dim, config.value_dim), + dtype=config.ssm_dtype, + device=device, + pin_memory=pin_memory, + ) + + # ======================================================================== + # 跨进程映射：在接收进程的当前 GPU 上注册共享页池 + # ======================================================================== + + def register_pin_memory(self) -> None: + """注册 Conv/SSM 共享存储，供本进程 GPU 使用；重复调用不会重复注册。""" + register_page_pool_tensor(tensor=self.conv_state) + register_page_pool_tensor(tensor=self.ssm_state) diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py new file mode 100644 index 0000000000..b8a641a73d --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -0,0 +1,522 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from .gdn_pages_pool import GDNPagesPool +from ..transfer_utils import bounded_att_buffer_copies +from ..copy_kernel import copy_bounded_att_tensors + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +logger = init_logger(__name__) + + +class GDNState(BoundedAttentionState): + """Gated DeltaNet 的 Conv 历史、SSM 和 MTP 游标，支持不同的 K/V head 数及维度。""" + + big_pages_pool: GDNPagesPool | None + small_pages_pool: GDNPagesPool | None + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ) -> None: + """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 + self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 + self.mtp_step = mtp_step + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) + } + storage_layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + storage_layer_count, + request_capacity * (mtp_step + 1), + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + + def get_state_type(self) -> str: + """返回 GDN 算法及 Conv/SSM 存储精度。""" + return f"gdn_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + GDNPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + GDNPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] + + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + start = req_idx * (self.mtp_step + 1) + self.ssm_state[:, start : start + self.mtp_step + 1].zero_() + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 + # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) + ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for req_idx in req_indices: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: + """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" + assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + if self.mtp_state_index is not None: + self.mtp_state_index[req_indices] = accepted_steps - 1 + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 + + region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 + 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 + """ + return self._get_transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 + pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 + # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: GDNState, + ) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 + # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def copy_runtime_state_from_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: GDNState, + ) -> None: + """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； + # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 + # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 + conv_copies, ssm_copies = [], [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + conv_copies.append( + ( + remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], + self.conv_state[:, dst_req, :, : self.config.conv_width], + ) + ) + ssm_copies.append( + ( + remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], + self.ssm_state[:, dst_req * (self.mtp_step + 1)], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for dst_req in dst_req_indices: + self.mtp_state_index[dst_req].fill_(0) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._get_transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._runtime_fields(req_idx) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._runtime_fields(req_idx) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + conv_region = TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping( + axis="window", + local_start=0, + external_start=0, + length=config.conv_width, + ), + ), + write_enabled=write_enabled, + ) + regions.append(conv_region) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=head_start, + length=config.num_v_heads, + ), + AxisMapping( + axis="key_dim", + local_start=0, + external_start=0, + length=config.key_dim, + ), + AxisMapping( + axis="value_dim", + local_start=0, + external_start=0, + length=config.value_dim, + ), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + def _runtime_fields(self, req: int): + """返回请求基准位置的 Conv/SSM 视图，供恢复和传输使用。""" + # TODO：当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0。 + # 若以后支持该场景，需要按已接受游标选取 Conv 窗口及 SSM 槽位。 + return self._fields( + self.conv_state[:, req, :, : self.config.conv_width], self.ssm_state[:, req * (self.mtp_step + 1)] + ) + + @staticmethod + def _fields(conv: torch.Tensor, ssm: torch.Tensor): + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py new file mode 100644 index 0000000000..fe1fceb3bd --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -0,0 +1,232 @@ +"""有界 attention 状态的功能接口；具体算法实现运行态与大小页存储。""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from functools import lru_cache +from typing import Any, List + +import torch + +from ..utils import deepcopy_result +from lightllm.common.attention_state.types import PageKind, StateTransferLayout, TransferRegion + + +class CoreInterface(ABC): + """有界 attention State 的资源、层视图和运行态核心能力。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + @abstractmethod + def get_state_type(self) -> str: + """返回算法与存储方案标记，用于生成资源成员名。 + + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att 类别前缀， + 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 + 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 + 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 + 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 + 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 + + KDA/GDN 格式为 {kda|gdn}_conv_{Conv精度}_ssm_{SSM精度}。 + 两份状态精度分别标明，不能只写一个 dtype，例如 kda_conv_bf16_ssm_fp32， + 对应 draft0_kda_conv_bf16_ssm_fp32_state。不同精度方案可以在同一角色下共存。 + 其他 bounded 算法使用自己的算法名及方案参数，不复用 kda/gdn 标记。 + """ + ... + + @abstractmethod + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建大小页池；容量为零表示不启用该页池。""" + ... + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + @abstractmethod + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" + ... + + @abstractmethod + def get_layer_view(self, logical_layer_id: int) -> Any: + """按逻辑层号返回运行态和算法控制信息的计算视图。""" + ... + + @abstractmethod + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """将指定槽位的运行态及控制信息恢复到长度 0，保留大小页快照。""" + ... + + @abstractmethod + def save_bounded_att_runtime_to_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """按各请求已执行的绝对 token 长度保存基准运行态，三个列表逐项对应。 + + 长度不包含尚未执行的采样 token；ring 状态据此确定窗口和槽位位置。 + 前缀与页槽位的关联仍由调用方维护，不要求页池保存长度元数据。 + """ + ... + + @abstractmethod + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """将页快照恢复到对应请求的绝对 token 长度，保留源快照数据。 + + req_indices、state_token_lens、page_slot_ids 逐项对应；长度用于恢复 ring 位置等控制信息， + 不表示本批搬运的 token 数，调用方保证与所选快照的前缀一致。 + """ + ... + + @abstractmethod + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: + """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" + ... + + +class CpuCacheInterface(ABC): + # ======================================================================== + # 三、CPU Cache 数据交换：读写大小页快照 + # ======================================================================== + + @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回一份 CPU 快照的字段区间、全局形状与 TP 切片。 + + byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； + 参与 CPU 页时，管理器按字段最大终点计算占用，不乘页内 token 数。 + 在合法缓存边界无需保留历史的 State 返回空元组，不占用 CPU 页面。 + 具体实现使用 lru_cache 缓存，外层 deepcopy_result 每次返回独立副本。 + 管理器分配后统一清理，子类覆盖时须保留缓存和副本语义。 + """ + ... + + @abstractmethod + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页；跳过 cpu_page_ready_flags 中为 True 的页面。""" + ... + + @abstractmethod + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" + ... + + +class DpCopyInterface(ABC): + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + @abstractmethod + def copy_big_pages_from_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: CoreInterface, + ) -> None: + """将其他 DP 同一 TP 分片的大页快照复制到本地大页，不修改运行态或分配页槽位。""" + ... + + @abstractmethod + def copy_runtime_state_from_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: CoreInterface, + ) -> None: + """复制其他 DP 同一 TP 分片的 prefill 基准运行态并重置目标游标，不修改页池。 + + 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 + """ + ... + + +class PdInterface(ABC): + # ======================================================================== + # 五、PD 分离传输：读写基准运行态 + # ======================================================================== + + @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 需要的字段、全局形状和 TP 切片，区域偏移相对本 State 起点。 + + 内层 lru_cache 缓存区域定义，外层 deepcopy_result 每次返回独立副本。 + 管理器分配存储后统一清理，子类覆盖时须保留缓存和副本语义。 + """ + ... + + @abstractmethod + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """将指定槽位在 state_token_len 绝对长度处的 prefill 基准运行态写入 PD 页。 + + 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 + """ + ... + + @abstractmethod + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按 state_token_len 恢复请求运行态及 ring 位置等控制信息，不写入大小页池。""" + ... + + +__all__ = [ + "CoreInterface", + "CpuCacheInterface", + "DpCopyInterface", + "PdInterface", +] diff --git a/lightllm/common/attention_state/bounded_att/kda/__init__.py b/lightllm/common/attention_state/bounded_att/kda/__init__.py new file mode 100644 index 0000000000..7154892202 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/__init__.py @@ -0,0 +1,5 @@ +from .create_kda_state import create_kda_state +from .kda_pages_pool import KDAPagesPool +from .kda_state import KDAState + +__all__ = ["create_kda_state", "KDAPagesPool", "KDAState"] diff --git a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py new file mode 100644 index 0000000000..49c9ad0ce5 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py @@ -0,0 +1,23 @@ +"""KDA State 的统一创建入口。""" + +from __future__ import annotations + +import torch + +from ..conv_ssm_types import ConvSSMStateConfig +from .kda_state import KDAState + + +def create_kda_state( + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", +) -> KDAState: + """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" + return KDAState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py b/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py new file mode 100644 index 0000000000..df17447cc4 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/kda_pages_pool.py @@ -0,0 +1,48 @@ +"""KDA pages: Conv history and square recurrent matrices; allocators own the page slots.""" + +from __future__ import annotations + + +import torch + +from ..conv_ssm_types import ConvSSMStateConfig +from ..utils import create_page_pool_tensor, register_page_pool_tensor + + +class KDAPagesPool: + def __init__( + self, + size: int, + config: ConvSSMStateConfig, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + assert size >= 0, "pool size must be nonnegative" + assert ( + config.num_k_heads == config.num_v_heads and config.key_dim == config.value_dim + ), "KDA pages require equal Q/K/V heads and dimensions" + self.size = size + device = torch.device(device) + pin_memory = pin_memory and device.type == "cpu" + storage_layer_count, heads, dim = len(config.logical_layer_ids), config.num_k_heads, config.key_dim + self.conv_state = create_page_pool_tensor( + size=(size, storage_layer_count, 3 * heads * dim, config.conv_width), + dtype=config.conv_dtype, + device=device, + pin_memory=pin_memory, + ) + self.ssm_state = create_page_pool_tensor( + size=(size, storage_layer_count, heads, dim, dim), + dtype=config.ssm_dtype, + device=device, + pin_memory=pin_memory, + ) + + # ======================================================================== + # 跨进程映射：在接收进程的当前 GPU 上注册共享页池 + # ======================================================================== + + def register_pin_memory(self) -> None: + """注册 Conv/SSM 共享存储，供本进程 GPU 使用；重复调用不会重复注册。""" + register_page_pool_tensor(tensor=self.conv_state) + register_page_pool_tensor(tensor=self.ssm_state) diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py new file mode 100644 index 0000000000..e5480a3ab0 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -0,0 +1,522 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, List +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from .kda_pages_pool import KDAPagesPool +from ..transfer_utils import bounded_att_buffer_copies +from ..copy_kernel import copy_bounded_att_tensors + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +logger = init_logger(__name__) + + +class KDAState(BoundedAttentionState): + """KDA 的 Conv 历史、方形 SSM 和 MTP 游标，要求 Q/K/V head 数及维度相同。""" + + big_pages_pool: KDAPagesPool | None + small_pages_pool: KDAPagesPool | None + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ) -> None: + """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 + self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 + self.mtp_step = mtp_step + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) + } + storage_layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + storage_layer_count, + request_capacity * (mtp_step + 1), + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + + def get_state_type(self) -> str: + """返回 KDA 算法及 Conv/SSM 存储精度。""" + return f"kda_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + KDAPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + KDAPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] + + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + start = req_idx * (self.mtp_step + 1) + self.ssm_state[:, start : start + self.mtp_step + 1].zero_() + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 + # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) + ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for req_idx in req_indices: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + ) -> None: + """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" + assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + if self.mtp_state_index is not None: + self.mtp_state_index[req_indices] = accepted_steps - 1 + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 + + region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 + 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 + """ + return self._get_transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 + pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 + # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: KDAState, + ) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 + # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def copy_runtime_state_from_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: KDAState, + ) -> None: + """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； + # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 + # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 + conv_copies, ssm_copies = [], [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + conv_copies.append( + ( + remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], + self.conv_state[:, dst_req, :, : self.config.conv_width], + ) + ) + ssm_copies.append( + ( + remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], + self.ssm_state[:, dst_req * (self.mtp_step + 1)], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for dst_req in dst_req_indices: + self.mtp_state_index[dst_req].fill_(0) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._get_transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._runtime_fields(req_idx) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._runtime_fields(req_idx) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + conv_region = TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping( + axis="window", + local_start=0, + external_start=0, + length=config.conv_width, + ), + ), + write_enabled=write_enabled, + ) + regions.append(conv_region) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=head_start, + length=config.num_v_heads, + ), + AxisMapping( + axis="key_dim", + local_start=0, + external_start=0, + length=config.key_dim, + ), + AxisMapping( + axis="value_dim", + local_start=0, + external_start=0, + length=config.value_dim, + ), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + def _runtime_fields(self, req: int): + """返回请求基准位置的 Conv/SSM 视图，供恢复和传输使用。""" + # TODO：当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0。 + # 若以后支持该场景，需要按已接受游标选取 Conv 窗口及 SSM 槽位。 + return self._fields( + self.conv_state[:, req, :, : self.config.conv_width], self.ssm_state[:, req * (self.mtp_step + 1)] + ) + + @staticmethod + def _fields(conv: torch.Tensor, ssm: torch.Tensor): + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/__init__.py b/lightllm/common/attention_state/bounded_att/sliding_window/__init__.py new file mode 100644 index 0000000000..06a759d916 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/sliding_window/__init__.py @@ -0,0 +1,4 @@ +from .create_sliding_window_state import create_sliding_window_state +from .fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState + +__all__ = ["create_sliding_window_state", "FP8E8M0SlidingWindowConfig", "FP8E8M0SlidingWindowState"] diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/create_sliding_window_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/create_sliding_window_state.py new file mode 100644 index 0000000000..c9f3af28bb --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/sliding_window/create_sliding_window_state.py @@ -0,0 +1,19 @@ +"""FP8E8M0SlidingWindowState 的独立创建入口。""" + +import torch + +from .fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig, FP8E8M0SlidingWindowState + + +def create_sliding_window_state( + config: FP8E8M0SlidingWindowConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", +) -> FP8E8M0SlidingWindowState: + return FP8E8M0SlidingWindowState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py new file mode 100644 index 0000000000..a1d72102ac --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py @@ -0,0 +1,367 @@ +"""单个共享 KV head 的 FP8/E8M0 滑动窗口运行态，不依赖模型或 compressor。""" + +from __future__ import annotations + +from functools import lru_cache + +from dataclasses import dataclass +from math import prod + +import torch + +from ...utils import deepcopy_result +from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from ..window_pages_pool import WindowPagesPool +from .fp8_e8m0_swa_transfer import copy_swa_window + + +@dataclass(frozen=True) +class FP8E8M0SlidingWindowConfig: + # 本 State 服务的逻辑层，顺序对应物理 layer 轴，各层拥有独立窗口。 + logical_layer_ids: tuple[int, ...] + # 每个 query 可见的近期 token 数，也是大小页和 PD 快照的逻辑窗口长度。 + window_size: int + # 未编码 KV 向量维度，包含末尾保留 BF16 的 RoPE 部分。 + kv_dim: int + # RoPE 维度；其余维度按 group64 FP8 编码，scale 使用 E8M0。 + rope_dim: int + # packed 物理页内的 token 数；页内 KV 数据与 scale 分区存放。 + page_token_size: int = 128 + # 各 TP rank 完整复制 KV；仅 rank 0 写出共享 CPU/PD 区域。 + tp_rank: int = 0 + tp_world_size: int = 1 + + @property + def kv_data_bytes(self) -> int: + return self.kv_dim + self.rope_dim + + @property + def kv_bytes(self) -> int: + scales = (self.kv_dim - self.rope_dim) // 64 + return self.kv_data_bytes + (scales + 7) // 8 * 8 + + def get_packed_kv_bytes(self, token_count: int) -> int: + """返回 packed 页字节数，包含数据行对齐。""" + return (token_count * self.kv_bytes + self.kv_data_bytes - 1) // self.kv_data_bytes * self.kv_data_bytes + + def get_attention_view(self, packed_kv: torch.Tensor) -> torch.Tensor: + """将运行态或 workspace 解释为 [page, token, head, byte]，不复制或重排数据。""" + pages = packed_kv.view(-1, self.get_packed_kv_bytes(self.page_token_size)) + return pages[:, : self.page_token_size * self.kv_bytes].view(-1, self.page_token_size, 1, self.kv_bytes) + + def __post_init__(self) -> None: + assert self.window_size > 0 and self.page_token_size > 0 + assert self.kv_dim > self.rope_dim >= 0 and (self.kv_dim - self.rope_dim) % 64 == 0 + assert self.logical_layer_ids and len(set(self.logical_layer_ids)) == len(self.logical_layer_ids) + assert 0 <= self.tp_rank < self.tp_world_size + assert self.get_packed_kv_bytes(self.page_token_size) % self.page_token_size == 0 + + +class FP8E8M0SlidingWindowState(BoundedAttentionState): + """仅保存近期原始 KV；prefill workspace 由计算 backend 管理。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: FP8E8M0SlidingWindowConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ) -> None: + self.config = config + self.state_name: str | None = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self._logical_to_storage_layer = {layer: index for index, layer in enumerate(config.logical_layer_ids)} + page_tokens = config.page_token_size + self.swa_capacity_tokens = (config.window_size + 2 * mtp_step + page_tokens - 1) // page_tokens * page_tokens + self.HOLD_REQUEST_ID = request_capacity - 1 + self.HOLD_SWA_SLOT = self.HOLD_REQUEST_ID * self.swa_capacity_tokens + self.swa_buffer = torch.zeros( + size=( + len(config.logical_layer_ids), + request_capacity, + self.swa_capacity_tokens, + config.get_packed_kv_bytes(page_tokens) // page_tokens, + ), + dtype=torch.uint8, + device=device, + ) + self.big_pages_pool: WindowPagesPool | None = None + self.small_pages_pool: WindowPagesPool | None = None + + def get_state_type(self) -> str: + """窗口长度参与方案名称，允许同一模型组合不同滑窗。""" + return f"swa_w{self.config.window_size}_fp8_group64_e8m0" + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """大小页只保存时间有序 KV，不保存 ring padding 或 MTP 候选位置。""" + shape = (len(self.config.logical_layer_ids), self.config.window_size, self.config.kv_bytes) + self.big_pages_pool = ( + WindowPagesPool( + size=big_page_capacity, + snapshot_shape=shape, + dtype=torch.uint8, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + WindowPagesPool( + size=small_page_capacity, + snapshot_shape=shape, + dtype=torch.uint8, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、运行态：层视图、重置与大小页保存恢复 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + return logical_layer_id in self._logical_to_storage_layer + + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: + """返回该层的 [request, ring_token, packed_byte] 视图。""" + return self.swa_buffer[self._logical_to_storage_layer[logical_layer_id]] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + for req_idx in req_indices: + self.swa_buffer[:, req_idx].zero_() + + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + pool = self._pool(page_kind=page_kind) + for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): + self._copy_window( + req_idx=req_idx, + external=pool.buffer[slot], + state_token_len=length, + write=True, + ) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + pool = self._pool(page_kind=page_kind) + for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): + self._copy_window( + req_idx=req_idx, + external=pool.buffer[slot], + state_token_len=length, + write=False, + ) + + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """ring 按绝对 token 位置寻址，无需提交额外游标。""" + pass + + # ======================================================================== + # 三、CPU Cache：仅搬运时间有序的大小页快照 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """CPU 快照保存完整近期窗口。""" + return self._get_transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定快照写入 CPU 页，已就绪页保持不变。""" + pool = self._pool(page_kind=page_kind) + fields, pages = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + fields.append({"swa": TensorField(tensor=pool.buffer[slot], axes=("layer", "window", "channel"))}) + pages.append(page) + for copies in bounded_att_buffer_copies( + page_fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=pages, + write=True, + ): + if copies: + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """CPU 页先恢复到大页池，随后再按前缀长度恢复请求 ring。""" + fields = [ + {"swa": TensorField(tensor=self.big_pages_pool.buffer[slot], axes=("layer", "window", "channel"))} + for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + if copies: + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP：相同方案的快照及请求运行态复制 + # ======================================================================== + + def copy_big_pages_from_dp( + self, + src_page_ids: list[int], + dst_page_ids: list[int], + remote_dp_bounded_state: FP8E8M0SlidingWindowState, + ) -> None: + """复制已保存的大页，调用方保持槽位及 stream 依赖一致。""" + assert self.config == remote_dp_bounded_state.config + if self.big_pages_pool is not None: + copy_bounded_att_tensors( + copies=[ + (remote_dp_bounded_state.big_pages_pool.buffer[src], self.big_pages_pool.buffer[dst]) + for src, dst in zip(src_page_ids, dst_page_ids) + ] + ) + + def copy_runtime_state_from_dp( + self, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: FP8E8M0SlidingWindowState, + ) -> None: + """直接复制相同配置和 MTP 容量的物理 ring，不创建快照。""" + assert self.config == remote_dp_bounded_state.config and self.mtp_step == remote_dp_bounded_state.mtp_step + copy_bounded_att_tensors( + copies=[ + (remote_dp_bounded_state.swa_buffer[:, src], self.swa_buffer[:, dst]) + for src, dst in zip(src_req_indices, dst_req_indices) + ] + ) + + # ======================================================================== + # 五、PD：按绝对前缀展开完整逻辑窗口，兼容不同 MTP 容量 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回完整逻辑窗口，不包含本地 MTP 候选区。""" + return self._get_transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按绝对长度将请求 ring 写入选定 PD 页。""" + page = pd_page_buffers[pd_page_index].view(torch.uint8).view(-1) + for region in layout.regions: + if region.write_enabled: + external = page[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + self._copy_window( + req_idx=req_idx, + external=external, + state_token_len=state_token_len, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按接收端 ring 容量还原逻辑窗口，不写入大小页池。""" + page = pd_page_buffers[pd_page_index].view(torch.uint8).view(-1) + for region in layout.regions: + external = page[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + self._copy_window( + req_idx=req_idx, + external=external, + state_token_len=state_token_len, + write=False, + ) + + # ======================================================================== + # 六、内部页池和窗口搬运 + # ======================================================================== + + def _pool(self, page_kind: PageKind) -> WindowPagesPool: + assert page_kind in ("big", "small") + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled" + return pool + + def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: + shape = (len(self.config.logical_layer_ids), self.config.window_size, self.config.kv_bytes) + return ( + TransferRegion( + field="swa", + byte_start=0, + byte_end=prod(shape), + dtype=torch.uint8, + axes=("layer", "window", "channel"), + shape=shape, + axis_mappings=(), + write_enabled=self.config.tp_rank == 0, + ), + ) + + def _copy_window( + self, + req_idx: int, + external: torch.Tensor, + state_token_len: int, + write: bool, + ) -> None: + copy_swa_window( + local=self.swa_buffer[:, req_idx], + external=external, + state_token_len=state_token_len, + config=self.config, + write=write, + ) diff --git a/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_transfer.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_transfer.py new file mode 100644 index 0000000000..49774a5c87 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_transfer.py @@ -0,0 +1,97 @@ +"""FP8/E8M0 SWA packed ring 与时间有序快照的 GPU 搬运。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import triton +import triton.language as tl + +if TYPE_CHECKING: + from .fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig + + +def copy_swa_window( + local: torch.Tensor, + external: torch.Tensor, + state_token_len: int, + config: FP8E8M0SlidingWindowConfig, + write: bool, +) -> None: + """按绝对长度在 packed SWA ring 和 [layer, window, byte] 快照间双向复制。 + + local 是选定请求的视图。external 可为已 pinned/registered 的大小页或 CUDA PD 页； + 直接由 GPU 读写，不分配 gather Tensor，不在主机等待拷贝完成。 + SWA 的数据和 scale 分区存储，不能直接按 token 行 memcpy。 + """ + elements = external.numel() + _copy_swa_window[(min(24, triton.cdiv(elements, 1024)),)]( + LOCAL=local, + EXTERNAL=external, + window_start=state_token_len - external.shape[1], + LOCAL_LAYER_STRIDE=local.stride(0), + EXTERNAL_STRIDES=external.stride(), + LAYERS=local.shape[0], + RING_TOKENS=local.shape[1], + WINDOW_TOKENS=external.shape[1], + KV_BYTES=config.kv_bytes, + KV_DATA_BYTES=config.kv_data_bytes, + PAGE_TOKENS=config.page_token_size, + PAGE_BYTES=config.get_packed_kv_bytes(config.page_token_size), + WRITE=write, + BLOCK=1024, + num_warps=4, + ) + + +@triton.jit(do_not_specialize=["window_start"]) +def _copy_swa_window( + LOCAL, + EXTERNAL, + window_start, + LOCAL_LAYER_STRIDE: tl.constexpr, + EXTERNAL_STRIDES: tl.constexpr, + LAYERS: tl.constexpr, + RING_TOKENS: tl.constexpr, + WINDOW_TOKENS: tl.constexpr, + KV_BYTES: tl.constexpr, + KV_DATA_BYTES: tl.constexpr, + PAGE_TOKENS: tl.constexpr, + PAGE_BYTES: tl.constexpr, + WRITE: tl.constexpr, + BLOCK: tl.constexpr, +): + # 阶段一：按线性快照拆出 layer/window/byte，绝对窗口起点保持运行时参数。 + lanes = tl.arange(0, BLOCK).to(tl.int64) + count = LAYERS * WINDOW_TOKENS * KV_BYTES + window_start = window_start.to(tl.int64) + for start in range(tl.program_id(0).to(tl.int64) * BLOCK, count, tl.num_programs(0).to(tl.int64) * BLOCK): + index = start + lanes + byte = index % KV_BYTES + window = index // KV_BYTES % WINDOW_TOKENS + layer = index // (KV_BYTES * WINDOW_TOKENS) + position = window_start + window + slot = position % RING_TOKENS + + # 阶段二：SWA 每页先存全部 KV 数据，再存全部 scale，末尾保留物理 padding。 + packed_offset = slot // PAGE_TOKENS * PAGE_BYTES + tl.where( + byte < KV_DATA_BYTES, + slot % PAGE_TOKENS * KV_DATA_BYTES + byte, + PAGE_TOKENS * KV_DATA_BYTES + slot % PAGE_TOKENS * (KV_BYTES - KV_DATA_BYTES) + byte - KV_DATA_BYTES, + ) + local_offset = layer * tl.cast(LOCAL_LAYER_STRIDE, tl.int64) + packed_offset + external_offset = ( + layer * tl.cast(EXTERNAL_STRIDES[0], tl.int64) + + window * tl.cast(EXTERNAL_STRIDES[1], tl.int64) + + byte * tl.cast(EXTERNAL_STRIDES[2], tl.int64) + ) + valid = index < count + has_history = valid & (position >= 0) + # TODO：按数据/scale 两个分区扩大搬运单位，融合宽类型访问，保持 packed padding 不变。 + if WRITE: + values = tl.load(LOCAL + local_offset, mask=has_history, other=0) + tl.store(EXTERNAL + external_offset, values, mask=valid) + else: + values = tl.load(EXTERNAL + external_offset, mask=has_history, other=0) + tl.store(LOCAL + local_offset, values, mask=has_history) diff --git a/lightllm/common/attention_state/bounded_att/transfer_utils.py b/lightllm/common/attention_state/bounded_att/transfer_utils.py new file mode 100644 index 0000000000..c5123600d1 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/transfer_utils.py @@ -0,0 +1,69 @@ +"""Bounded attention 的 CPU/PD 页传输：复制一份完整状态快照。""" + +from collections.abc import Mapping + +import torch + +from ..types import StateTransferLayout, TensorField + + +# ======================================================================== +# 一、CPU/PD 页交换：按状态字段生成复制视图 +# ======================================================================== + + +def bounded_att_buffer_copies( + page_fields: list[Mapping[str, TensorField]], + page_buffers: torch.Tensor, + layout: StateTransferLayout, + page_indices: list[int], + write: bool, +) -> list[list[tuple[torch.Tensor, torch.Tensor]]]: + """按传输区域收集所有页面的复制视图，每个返回批次只包含一种 shape。 + + page_fields 与 page_indices 一一对应，调用方已排除 ready 页。 + CPU 路径的 page_fields 来自本地大小页池，PD 路径来自请求运行态；均为本地 TP 分片。 + layout.regions 描述外部页中的全局字段，以及本地 TP 分片在该字段中的位置。 + axis_mappings 的 local_start 固定指本地，external_start 固定指外部页，读回时也不交换含义。 + CPU 页和大小页快照的窗口均按逻辑时间排列，此处只处理普通轴切片，不展开或恢复 ring。 + Conv 的 Q/K/V 切片与 SSM 分别成批，不将不同字段或不同 TP 切片混入同一次拷贝。 + page_buffers 的第一维为物理页；这里只创建视图，由调用方逐批提交 Triton 算子。 + """ + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + batches = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + copies = [] + for local_fields, page_index in zip(page_fields, page_indices): + # 阶段一：region.field 选择本地 Conv/SSM；CPU 读回时这里就是目标大页槽位。 + local_field = local_fields[region.field] + local_tensor = local_field.tensor + + # 阶段二：region 的字节区间和 shape 定位外部页中的全局字段，再对齐轴顺序。 + page_bytes = page_buffers[page_index].view(torch.uint8).view(-1) + external_tensor = page_bytes[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + assert local_tensor.dtype == external_tensor.dtype, "local and external field dtypes must match" + external_tensor = external_tensor.permute(tuple(region.axes.index(axis) for axis in local_field.axes)) + + # 阶段三：region 的两侧切片定位同一个 TP 分片，不代表本次复制的源和目标。 + # 例如 rank 1 的本地 head[0:4] 对应外部页 head[4:8]，写出和读回使用相同映射。 + local_slices = [slice(None)] * local_tensor.ndim + external_slices = [slice(None)] * external_tensor.ndim + for axis_mapping in region.axis_mappings: + axis = local_field.axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) + local_view = local_tensor[tuple(local_slices)] + external_view = external_tensor[tuple(external_slices)] + assert local_view.shape == external_view.shape, "local and external copy views must have the same shape" + + # 阶段四：只在此确定搬运方向；CPU 读回为外部页 → 本地大页池，不直接恢复请求。 + if write: + copies.append((local_view, external_view)) + else: + copies.append((external_view, local_view)) + batches.append(copies) + return batches diff --git a/lightllm/common/attention_state/bounded_att/utils.py b/lightllm/common/attention_state/bounded_att/utils.py new file mode 100644 index 0000000000..93f5b2ebe3 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/utils.py @@ -0,0 +1,125 @@ +"""Bounded attention 页池的存储分配策略。""" + +from __future__ import annotations + +from math import prod +from weakref import finalize + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + + +logger = init_logger(__name__) + +# Portable：本进程内其他 CUDA context 也将该内存识别为 pinned；不代替跨进程注册。 +CUDA_HOST_REGISTER_PORTABLE = 1 +# Mapped：建立设备可访问映射，供 CUDA/Triton kernel 直接访问主机存储。 +CUDA_HOST_REGISTER_MAPPED = 2 + + +# ======================================================================== +# 一、页池分配：普通存储或 DP fetch 所需的共享 pinned 存储 +# ======================================================================== + + +def create_page_pool_tensor( + size: tuple[int, ...], + dtype: torch.dtype, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, +) -> torch.Tensor: + """按 DP fetch 开关创建页池 Tensor，保留调用方指定的设备及 pin_memory 行为。 + + CPU 页池在启用 DP fetch 时直接使用共享 Storage，避免普通 pinned Tensor + 经 Queue 导出时重新分配共享存储、复制整块页池，并丢失原地址的 CUDA 注册。 + 接收进程仍需为自己的共享映射地址注册 CUDA；注册状态不随 Tensor 传递。 + """ + device = torch.device(device) + pin_memory = pin_memory and device.type == "cpu" + shared = device.type == "cpu" and get_env_start_args().enable_dp_prompt_cache_fetch + nbytes = prod(size) * dtype.itemsize + + # 阶段一：未启用 DP fetch 时沿用 torch 分配；GPU 页池也保持原设备分配方式。 + if not shared: + tensor = torch.empty( + size=size, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + reason = "GPU device allocation" if device.type != "cpu" else "DP prompt cache fetch disabled" + logger.info( + f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, " + f"device={device}, shared=False, pin_memory={pin_memory}, allocator=torch.empty, reason={reason}" + ) + return tensor + + # 阶段二：直接申请共享字节存储，再绑定无数据区的 Tensor，不经过 share_memory_ 的整块复制。 + # _new_shared 是 PyTorch 私有入口，集中在此使用，便于版本升级时验证或替换。 + storage = torch.UntypedStorage._new_shared(nbytes, device="cpu") + tensor = torch.empty(size=(0,), dtype=dtype, device="cpu").set_( + source=storage, + storage_offset=0, + size=size, + ) + + # 阶段三：在当前进程原地注册；Portable | Mapped 允许 CUDA 异步搬运及设备映射。 + # 不调用 tensor.pin_memory()，否则会另建私有 pinned 存储，破坏共享关系。 + # 空页池没有数据，无需注册；显式 pin_memory=False 时仅创建共享 CPU 存储。 + if pin_memory: + register_page_pool_tensor(tensor=tensor) + logger.info( + f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, " + f"device=cpu, shared=True, pin_memory={pin_memory}, reason=DP prompt cache fetch; " + "direct shared allocation avoids a full pool copy during IPC export" + ) + return tensor + + +# ======================================================================== +# 二、本进程注册：首次创建和跨进程接收复用同一入口 +# ======================================================================== + + +def register_page_pool_tensor(tensor: torch.Tensor) -> None: + """将共享 CPU 存储注册到当前进程，已 pinned、空存储或 GPU Tensor 无需处理。 + + 接收端必须先选定自己的 CUDA device。按整个 Storage 注册，避免多个视图重复注册 + 重叠区间；注册和注销信息只保存在本进程，不依赖发送方的对象属性或默认 GPU。 + """ + if tensor.device.type != "cpu" or tensor.numel() == 0 or tensor.is_pinned(): + return + assert tensor.is_shared(), "page pool registration requires shared CPU storage" + storage = tensor.untyped_storage() + ptr = storage.data_ptr() + nbytes = storage.nbytes() + device_id = torch.cuda.current_device() + flags = CUDA_HOST_REGISTER_PORTABLE | CUDA_HOST_REGISTER_MAPPED + torch.cuda.check_error(torch.cuda.cudart().cudaHostRegister(ptr, nbytes, flags)) + # finalizer 绑定 Storage，所有本地 Tensor 及视图释放后才注销，不能绑在单个 Tensor 上。 + # 正常释放时等待 GPU 完成使用；解释器退出时交给进程清理，避免 CUDA 已关闭后再调用。 + cleanup = finalize( + storage, + _unregister_page_pool_storage, + ptr, + device_id, + ) + cleanup.atexit = False + logger.info( + f"Registered shared bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " + "flags=Portable|Mapped; registration applies to this process's mapping" + ) + + +# ======================================================================== +# 三、私有辅助：最后一个本地存储引用释放时注销 CUDA 注册 +# ======================================================================== + + +def _unregister_page_pool_storage(ptr: int, device_id: int) -> None: + """等待当前 GPU 使用结束，再注销本进程的映射；共享内存本身由 Storage 管理。""" + torch.cuda.synchronize(device=device_id) + torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr)) + logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, ptr={ptr}") diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/__init__.py b/lightllm/common/attention_state/bounded_att/window_compressor/__init__.py new file mode 100644 index 0000000000..41294a3751 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_compressor/__init__.py @@ -0,0 +1,4 @@ +from .create_window_compressor_state import create_window_compressor_state +from .window_compressor_state import WindowCompressorState, WindowCompressorStateConfig + +__all__ = ["create_window_compressor_state", "WindowCompressorState", "WindowCompressorStateConfig"] diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/create_window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/create_window_compressor_state.py new file mode 100644 index 0000000000..e80d802925 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_compressor/create_window_compressor_state.py @@ -0,0 +1,19 @@ +"""WindowCompressorState 的独立创建入口。""" + +import torch + +from .window_compressor_state import WindowCompressorStateConfig, WindowCompressorState + + +def create_window_compressor_state( + config: WindowCompressorStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", +) -> WindowCompressorState: + return WindowCompressorState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ) diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py new file mode 100644 index 0000000000..cea0a49af2 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_state.py @@ -0,0 +1,410 @@ +"""固定窗口、固定产出步长的 attention 压缩器候选/gate 运行态。""" + +from __future__ import annotations + +from functools import lru_cache + +from dataclasses import dataclass +from math import prod +from typing import Literal + +import torch + +from lightllm.utils.envs_utils import get_env_start_args + +from ...utils import deepcopy_result +from ...types import PageKind, StateTransferLayout, TensorField, TransferRegion +from ..base_state import BoundedAttentionState +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from ..window_pages_pool import WindowPagesPool +from .window_compressor_transfer import copy_compressor_window + + +@dataclass(frozen=True) +class WindowCompressorStateConfig: + """按固定窗口和产出步长保存候选向量及 gate，供模型的压缩或池化算子使用。 + + 这里只描述运行态的存储和历史保留方式，具体聚合、归一化及量化由模型算子完成。 + window_size / compress_stride 分别表示窗口覆盖的原始 token 数和相邻产出的间隔： + - GLM K-pool：W4/S4，每四个 key/gate 形成一组，无跨组重叠历史。 + - V4 C4：W8/S4，每四步产出一次，下一组仍需保留四个 token 的历史。 + - V4 C128：W128/S128，每 128 步产出一次，无跨组重叠历史。 + MTP 的额外候选槽位由 State.mtp_step 扩展，不计入窗口或产出步长。 + """ + + # 本实例服务的逻辑层，顺序对应运行态 Tensor 的 layer 轴。 + # attention KV 候选与索引 key 候选分别创建实例，可使用不同窗口、维度和精度。 + logical_layer_ids: tuple[int, ...] + # 单次聚合覆盖的原始 token 窗口，不包含 MTP 额外槽位。 + # 例如 K-pool 为 4，V4 C4 为 8；压缩比例同为 4 不代表窗口长度相同。 + window_size: int + # 相邻产出之间经过的原始 token 数；K-pool / V4 C4 为 4，V4 C128 为 128。 + # 等于 window_size 时无重叠，小于时保留 window_size - compress_stride 个历史位置。 + # 大小页快照在该步长的边界保存；PD 可在任意前缀终点传输完整逻辑窗口。 + compress_stride: int + # 每个原始 token 保存的候选向量宽度；gate 具有相同形状和存储精度。 + # K-pool 中等于 indexer key 的 head_dim；V4 W8/S4 为前后半窗口分别保存候选， + # 使用 2 * 输出向量维度。因此它不一定等于最终压缩条目的维度，也不包含 gate 的宽度。 + candidate_dim: int + # 候选数据的用途，仅用于区分规范 State 成员名，不选择聚合算法或改变存储和恢复行为。 + # kv 表示 attention KV 候选；indexer 表示索引 key 候选及其 gate，不代表整个 indexer。 + target: Literal["kv", "indexer"] = "kv" + # 候选与 gate 使用相同精度，沿最后一轴拼接为 candidate|gate，共用一个 Tensor。 + # 此精度描述压缩器输入历史，不是产出 KV / indexer key 的量化编码； + # GLM K-pool 随模型精度配置（通常 BF16），V4 的 KV/indexer 压缩器当前使用 FP32。 + # 当前接口不支持候选与 gate 使用不同 dtype。 + dtype: torch.dtype = torch.float32 + # 所属模型 TP 组内的 rank；运行态完整复制，仅 rank 0 向共享外部页写出。 + tp_rank: int = 0 + # 保存相同运行态副本的 TP rank 数，不用于切分 candidate_dim 或窗口长度。 + tp_world_size: int = 1 + + @property + def overlap_tokens(self) -> int: + """组完成时仍需跨步保留的重叠历史：W8/S4 为 4，W128/S128 为 0。""" + return self.window_size - self.compress_stride + + def __post_init__(self) -> None: + assert self.window_size >= self.compress_stride > 0 and self.candidate_dim > 0 + assert self.logical_layer_ids and len(set(self.logical_layer_ids)) == len(self.logical_layer_ids) + assert self.target in ("kv", "indexer") and self.dtype in (torch.float32, torch.float16, torch.bfloat16) + assert 0 <= self.tp_rank < self.tp_world_size + + +class WindowCompressorState(BoundedAttentionState): + """只保存窗口候选/gate，不保存压缩后的 KV 或 SWA。 + + 大小页在产出步长的边界保存重叠历史，PD 则保存完整窗口以支持任意前缀长度。 + 无重叠的窗口在组边界不需要快照，因此不创建大小页池，也不占 CPU 联合页字节。 + """ + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: WindowCompressorStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ) -> None: + self.config = config + self.state_name: str | None = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self._logical_to_storage_layer = {layer: index for index, layer in enumerate(config.logical_layer_ids)} + self.buffer = torch.zeros( + size=( + len(config.logical_layer_ids), + request_capacity, + config.window_size + mtp_step, + 2 * config.candidate_dim, + ), + dtype=config.dtype, + device=device, + ) + # HOLD 不得成为有效候选；正常请求在 reset 时初始化候选及 gate。 + self.buffer[:, request_capacity - 1, :, config.candidate_dim :].fill_(float("-inf")) + self.big_pages_pool: WindowPagesPool | None = None + self.small_pages_pool: WindowPagesPool | None = None + + def get_state_type(self) -> str: + """用途、窗口、步长和精度共同标识可复用的压缩器方案。""" + precision = {torch.float32: "fp32", torch.float16: "fp16", torch.bfloat16: "bf16"}[self.config.dtype] + return ( + f"compressor_{self.config.target}_w{self.config.window_size}" f"_s{self.config.compress_stride}_{precision}" + ) + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """只为跨步重叠历史创建快照池；无重叠时保持 None。""" + if not self.config.overlap_tokens: + return + shape = (len(self.config.logical_layer_ids), self.config.overlap_tokens, 2 * self.config.candidate_dim) + self.big_pages_pool = ( + WindowPagesPool( + size=big_page_capacity, + snapshot_shape=shape, + dtype=self.config.dtype, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + WindowPagesPool( + size=small_page_capacity, + snapshot_shape=shape, + dtype=self.config.dtype, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、运行态：层视图、重置与大小页保存恢复 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + return logical_layer_id in self._logical_to_storage_layer + + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: + """返回该层 [request, window+MTP, candidate|gate] 的计算视图。""" + return self.buffer[self._logical_to_storage_layer[logical_layer_id]] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + for req_idx in req_indices: + self.buffer[:, req_idx].zero_() + self.buffer[:, req_idx, :, self.config.candidate_dim :].fill_(float("-inf")) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + """边界必须完成一个产出步长；仅保留下一步仍需读取的窗口重叠部分。""" + assert all( + length % self.config.compress_stride == 0 for length in state_token_lens + ), "prefix must close a compressor stride" + if not self.config.overlap_tokens: + return + pool = self._pool(page_kind=page_kind) + for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): + self._copy_window( + req_idx=req_idx, + external=pool.buffer[slot], + state_token_len=length, + write=True, + ) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + assert all( + length % self.config.compress_stride == 0 for length in state_token_lens + ), "prefix must close a compressor stride" + if not self.config.overlap_tokens: + return + pool = self._pool(page_kind=page_kind) + for req_idx, length, slot in zip(req_indices, state_token_lens, page_slot_ids): + self._copy_window( + req_idx=req_idx, + external=pool.buffer[slot], + state_token_len=length, + write=False, + ) + + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """按绝对位置寻址；接受后继续覆盖 ring，不维护额外 MTP 游标。""" + pass + + # ======================================================================== + # 三、CPU Cache：仅搬运时间有序的大小页快照 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """校验缓存页对齐产出步长；无重叠历史的窗口不占 CPU 页，也不创建大小页池。""" + args = get_env_start_args() + # 大页是小页的整数倍；小页对齐产出步长后，大小页终点均不会落在未完成的组内。 + assert ( + args.linear_att_hash_page_size % self.config.compress_stride == 0 + ), "state cache pages must align to the compressor stride" + # CPU 服务只查询 meta 布局，也要在忽略无重叠窗口之前拒绝不安全的恢复边界。 + if args.cpu_cache_token_page_size is not None: + assert ( + args.cpu_cache_token_page_size % self.config.compress_stride == 0 + ), "CPU cache pages must align to the compressor stride" + return self._get_transfer_regions(runtime=False) + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定快照写入 CPU 页，已就绪页保持不变。""" + pool = self._pool(page_kind=page_kind) + fields, pages = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + fields.append({"compressor": TensorField(tensor=pool.buffer[slot], axes=("layer", "window", "channel"))}) + pages.append(page) + for copies in bounded_att_buffer_copies( + page_fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=pages, + write=True, + ): + if copies: + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """CPU 页先恢复到大页池，随后再按前缀长度恢复请求 ring。""" + fields = [ + {"compressor": TensorField(tensor=self.big_pages_pool.buffer[slot], axes=("layer", "window", "channel"))} + for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + if copies: + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP：相同方案的快照及请求运行态复制 + # ======================================================================== + + def copy_big_pages_from_dp( + self, + src_page_ids: list[int], + dst_page_ids: list[int], + remote_dp_bounded_state: WindowCompressorState, + ) -> None: + """复制已保存的大页，调用方保持槽位及 stream 依赖一致。""" + assert self.config == remote_dp_bounded_state.config + if self.big_pages_pool is not None: + copy_bounded_att_tensors( + copies=[ + (remote_dp_bounded_state.big_pages_pool.buffer[src], self.big_pages_pool.buffer[dst]) + for src, dst in zip(src_page_ids, dst_page_ids) + ] + ) + + def copy_runtime_state_from_dp( + self, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: WindowCompressorState, + ) -> None: + """直接复制相同配置和 MTP 容量的物理 ring，不创建快照。""" + assert self.config == remote_dp_bounded_state.config and self.mtp_step == remote_dp_bounded_state.mtp_step + copy_bounded_att_tensors( + copies=[ + (remote_dp_bounded_state.buffer[:, src], self.buffer[:, dst]) + for src, dst in zip(src_req_indices, dst_req_indices) + ] + ) + + # ======================================================================== + # 五、PD：按绝对前缀展开完整逻辑窗口，兼容不同 MTP 容量 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回完整逻辑窗口，不包含本地 MTP 候选区。""" + return self._get_transfer_regions(runtime=True) + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按绝对长度将请求 ring 写入选定 PD 页。""" + page = pd_page_buffers[pd_page_index].view(torch.uint8).view(-1) + for region in layout.regions: + if region.write_enabled: + external = page[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + self._copy_window( + req_idx=req_idx, + external=external, + state_token_len=state_token_len, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """按接收端 ring 容量还原逻辑窗口，不写入大小页池。""" + page = pd_page_buffers[pd_page_index].view(torch.uint8).view(-1) + for region in layout.regions: + external = page[region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + self._copy_window( + req_idx=req_idx, + external=external, + state_token_len=state_token_len, + write=False, + ) + + # ======================================================================== + # 六、内部页池和窗口搬运 + # ======================================================================== + + def _pool(self, page_kind: PageKind) -> WindowPagesPool: + assert page_kind in ("big", "small") + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled" + return pool + + def _get_transfer_regions(self, runtime: bool) -> tuple[TransferRegion, ...]: + window = self.config.window_size if runtime else self.config.overlap_tokens + if not window: + return () + shape = (len(self.config.logical_layer_ids), window, 2 * self.config.candidate_dim) + return ( + TransferRegion( + field="compressor", + byte_start=0, + byte_end=prod(shape) * self.config.dtype.itemsize, + dtype=self.config.dtype, + axes=("layer", "window", "channel"), + shape=shape, + axis_mappings=(), + write_enabled=self.config.tp_rank == 0, + ), + ) + + def _copy_window( + self, + req_idx: int, + external: torch.Tensor, + state_token_len: int, + write: bool, + ) -> None: + copy_compressor_window( + local=self.buffer[:, req_idx], + external=external, + state_token_len=state_token_len, + write=write, + ) diff --git a/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_transfer.py b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_transfer.py new file mode 100644 index 0000000000..dc985d68e6 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_compressor/window_compressor_transfer.py @@ -0,0 +1,116 @@ +"""窗口压缩器候选/gate ring 与时间有序快照的 GPU 搬运。""" + +import torch +import triton +import triton.language as tl + +from ...utils import view_max_copy_dtype + + +def copy_compressor_window( + local: torch.Tensor, + external: torch.Tensor, + state_token_len: int, + write: bool, +) -> None: + """在候选/gate ring 与时间有序快照之间复制，窗口长度由 external 的形状决定。 + + 大小页可以只保留跨步重叠历史，PD 保存完整窗口；两者均不包含 MTP 候选区。 + """ + if not local.numel(): + return + local_view, external_view = view_max_copy_dtype(copies=[(local, external)])[0] + window = external.shape[1] + elements = local_view.shape[0] * window * local_view.shape[2] + _copy_compressor_window[(min(24, triton.cdiv(elements, 1024)),)]( + LOCAL=local_view, + EXTERNAL=external_view, + window_start=state_token_len - window, + external_window_start=0, + LOCAL_STRIDES=local_view.stride(), + EXTERNAL_STRIDES=external_view.stride(), + LAYERS=local_view.shape[0], + LOCAL_RING_SIZE=local.shape[1], + EXTERNAL_RING_SIZE=window, + WINDOW_SIZE=window, + CHANNELS=local_view.shape[2], + ELEMENT_BYTES=local_view.element_size(), + LOCAL_CIRCULAR=True, + EXTERNAL_CIRCULAR=False, + WRITE=write, + BLOCK_ELEMENTS=1024, + num_warps=4, + ) + + +@triton.jit(do_not_specialize=["window_start", "external_window_start"]) +def _copy_compressor_window( + LOCAL, # 当前请求的本地 ring：[layer, ring_slot, channel]。 + EXTERNAL, # 已选定页中的候选/gate 拼接字段：[layer, window, channel]。 + window_start, # 未取模的绝对窗口起点；短前缀时可为负数，保持为运行时参数。 + external_window_start, # 本次窗口在外部字段中的起点，通常为 0。 + LOCAL_STRIDES: tl.constexpr, + EXTERNAL_STRIDES: tl.constexpr, + LAYERS: tl.constexpr, + LOCAL_RING_SIZE: tl.constexpr, + EXTERNAL_RING_SIZE: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + CHANNELS: tl.constexpr, + ELEMENT_BYTES: tl.constexpr, + LOCAL_CIRCULAR: tl.constexpr, + EXTERNAL_CIRCULAR: tl.constexpr, + WRITE: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, +): + if ELEMENT_BYTES == 8: + copy_dtype = tl.uint64 + elif ELEMENT_BYTES == 4: + copy_dtype = tl.uint32 + elif ELEMENT_BYTES == 2: + copy_dtype = tl.uint16 + else: + copy_dtype = tl.uint8 + LOCAL = LOCAL.to(tl.pointer_type(copy_dtype)) + EXTERNAL = EXTERNAL.to(tl.pointer_type(copy_dtype)) + window_start = tl.cast(window_start, tl.int64) + external_window_start = tl.cast(external_window_start, tl.int64) + element_count = LAYERS * WINDOW_SIZE * CHANNELS + program_start = tl.program_id(0).to(tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(0).to(tl.int64) * BLOCK_ELEMENTS + lanes = tl.arange(0, BLOCK_ELEMENTS).to(tl.int64) + + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=1): + # 两个方向共用逻辑窗口坐标，只在 load/store 时交换数据流向。 + element = block_start + lanes + valid = element < element_count + channel = element % CHANNELS + window_index = element // CHANNELS % WINDOW_SIZE + layer = element // (CHANNELS * WINDOW_SIZE) + # 先用绝对位置判断历史是否存在，再映射到物理 ring，避免取模丢失负数信息。 + position = window_start + window_index + has_history = position >= 0 + local_slot = position + if LOCAL_CIRCULAR: + # 窗口不超过 ring 容量；先加容量，使短前缀读回补零也落在非负槽位。 + local_slot = (position + LOCAL_RING_SIZE) % LOCAL_RING_SIZE + external_slot = external_window_start + window_index + if EXTERNAL_CIRCULAR: + external_slot %= EXTERNAL_RING_SIZE + + # 所有 stride 在参与地址乘加前转为 int64，支持请求切片产生的大 layer stride。 + local_offset = ( + layer * tl.cast(LOCAL_STRIDES[0], tl.int64) + + local_slot * tl.cast(LOCAL_STRIDES[1], tl.int64) + + channel * tl.cast(LOCAL_STRIDES[2], tl.int64) + ) + external_offset = ( + layer * tl.cast(EXTERNAL_STRIDES[0], tl.int64) + + external_slot * tl.cast(EXTERNAL_STRIDES[1], tl.int64) + + channel * tl.cast(EXTERNAL_STRIDES[2], tl.int64) + ) + if WRITE: + value = tl.load(LOCAL + local_offset, mask=valid & has_history, other=0) + tl.store(EXTERNAL + external_offset, value, mask=valid) + else: + value = tl.load(EXTERNAL + external_offset, mask=valid & has_history, other=0) + tl.store(LOCAL + local_offset, value, mask=valid) diff --git a/lightllm/common/attention_state/bounded_att/window_pages_pool.py b/lightllm/common/attention_state/bounded_att/window_pages_pool.py new file mode 100644 index 0000000000..957eef0613 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/window_pages_pool.py @@ -0,0 +1,27 @@ +"""单种时间有序窗口快照的页池，槽位由 manager 的共享 allocator 管理。""" + +import torch + +from .utils import create_page_pool_tensor, register_page_pool_tensor + + +class WindowPagesPool: + def __init__( + self, + size: int, + snapshot_shape: tuple[int, int, int], + dtype: torch.dtype, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + # [页槽位, 物理层, 时间窗口, 字段宽度]；只保存快照，不保存请求或长度元数据。 + self.buffer = create_page_pool_tensor( + size=(size, *snapshot_shape), + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + + def register_pin_memory(self) -> None: + """跨 DP 导入后，在接收进程当前设备上注册共享页池。""" + register_page_pool_tensor(tensor=self.buffer) diff --git a/lightllm/common/attention_state/cpu_cache_operator.py b/lightllm/common/attention_state/cpu_cache_operator.py new file mode 100644 index 0000000000..814d340481 --- /dev/null +++ b/lightllm/common/attention_state/cpu_cache_operator.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import replace +from functools import lru_cache +from typing import TYPE_CHECKING, List + +import torch + +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.utils.envs_utils import get_env_start_args + +if TYPE_CHECKING: + from lightllm.common.attention_state.manager import AttentionStateManager + + +class CpuCacheOperator: + """实现 CPU cache 的联合布局和 State 搬运。 + + AttentionStateManager 通过该 Operator 提供 CPU 页操作；这里集中处理 CPU 页如何组织， + 以及 full-attention token 与 bounded 运行态快照如何写回/读回。 + """ + + def __init__(self, owner: AttentionStateManager): + self.owner = owner + + @lru_cache(maxsize=None) + def get_page_layouts(self) -> Mapping[str, StateTransferLayout]: + page_tokens = get_env_start_args().cpu_cache_token_page_size + assert page_tokens > 0, "CPU pages require a positive token capacity" + states = (*self.owner.get_full_attention_states(), *self.owner.get_bounded_attention_states()) + assert states, "a CPU page requires participating storage" + + layouts, offset = {}, 0 + for state in states: + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_cpu_transfer_regions() + ) + # 没有跨步历史的 State 不占 CPU 页；其余 State 仍按联合页顺序排列。 + if not regions: + continue + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=page_tokens, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + @lru_cache(maxsize=None) + def get_page_bytes(self) -> int: + layouts = self.get_page_layouts() + page_bytes = max((layout.byte_end for layout in layouts.values()), default=0) + return (page_bytes + 15) // 16 * 16 + + def clear_layout_cache(self) -> None: + self.get_page_layouts.cache_clear() + self.get_page_bytes.cache_clear() + + def write_req( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layouts: Mapping[str, StateTransferLayout], + cpu_page_ready_flags: List[bool] | None = None, + big_page_slot_ids: List[int] | None = None, + small_page_slot_ids: List[int] | None = None, + ) -> None: + """写出单个请求的 full-attention token 和 bounded 页快照。""" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + if all(cpu_page_ready_flags): + return + + full_states = self.owner.get_full_attention_states() + bounded_states = self.owner.get_bounded_attention_states() + if bounded_states: + big_page_slot_ids = [] if big_page_slot_ids is None else big_page_slot_ids + small_page_slot_ids = [] if small_page_slot_ids is None else small_page_slot_ids + big_page_count = len(big_page_slot_ids) + assert big_page_count + len(small_page_slot_ids) == len( + cpu_page_indices + ), "one state slot is required per CPU page" + big_cpu_page_indices = cpu_page_indices[:big_page_count] + small_cpu_page_indices = cpu_page_indices[big_page_count:] + big_page_readies = cpu_page_ready_flags[:big_page_count] if cpu_page_ready_flags is not None else None + small_page_readies = cpu_page_ready_flags[big_page_count:] if cpu_page_ready_flags is not None else None + + for state in full_states: + state.write_to_cpu( + token_indices=token_indices, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + layout=layouts[state.state_name], + cpu_page_ready_flags=cpu_page_ready_flags, + ) + + for state in bounded_states: + layout = layouts.get(state.state_name) + if layout is None: + # 某些压缩窗口在合法缓存边界不需要快照，因此可以不进入 CPU cache。 + assert not state.get_cpu_transfer_regions(), f"missing CPU layout for {state.state_name}" + continue + if big_page_slot_ids: + state.write_state_pages_to_cpu( + page_kind="big", + page_slot_ids=big_page_slot_ids, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=big_cpu_page_indices, + layout=layout, + cpu_page_ready_flags=big_page_readies, + ) + if small_page_slot_ids: + state.write_state_pages_to_cpu( + page_kind="small", + page_slot_ids=small_page_slot_ids, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=small_cpu_page_indices, + layout=layout, + cpu_page_ready_flags=small_page_readies, + ) + + def read_req( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layouts: Mapping[str, StateTransferLayout], + big_page_slot_ids: List[int] | None = None, + req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: + """读回请求页面，并从末页恢复 bounded 运行态。""" + full_states = self.owner.get_full_attention_states() + bounded_states = self.owner.get_bounded_attention_states() + if bounded_states: + assert req_idx is not None and cpu_page_indices, "state recovery requires a target request and prefix pages" + assert state_token_len is not None, "bounded state recovery requires its absolute prefix length" + page_tokens = get_env_start_args().cpu_cache_token_page_size + has_partial_page = token_indices.numel() % page_tokens != 0 + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert len(big_page_slot_ids) + int(has_partial_page) == len( + cpu_page_indices + ), "slots must cover complete CPU pages only" + load_big_page_slot_ids = big_page_slot_ids.copy() + if has_partial_page: + load_big_page_slot_ids.append(self.owner.bounded_att_cpu_load_big_page_slot) + + for state in full_states: + state.read_from_cpu( + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + token_indices=token_indices, + layout=layouts[state.state_name], + ) + + for state in bounded_states: + if state.state_name not in layouts: + assert not state.get_cpu_transfer_regions(), f"missing CPU layout for {state.state_name}" + continue + state.read_state_pages_from_cpu( + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=cpu_page_indices, + big_page_slot_ids=load_big_page_slot_ids, + layout=layouts[state.state_name], + ) + + if bounded_states: + self.owner.restore_bounded_att_runtime_from_pages( + req_indices=[req_idx], + state_token_lens=[state_token_len], + page_kind="big", + page_slot_ids=[load_big_page_slot_ids[-1]], + ) diff --git a/lightllm/common/attention_state/dp_shared_operator.py b/lightllm/common/attention_state/dp_shared_operator.py new file mode 100644 index 0000000000..5f0eb7fa0d --- /dev/null +++ b/lightllm/common/attention_state/dp_shared_operator.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from copy import copy +from typing import TYPE_CHECKING, List + +import torch + +if TYPE_CHECKING: + from lightllm.common.attention_state.manager import AttentionStateManager + + +class DPSharedOperator: + """封装 DP shared 的跨副本复制、IPC 导出和共享页注册。""" + + def __init__(self, owner: AttentionStateManager): + self.owner = owner + + def create_ipc_copy(self) -> AttentionStateManager: + """保留共享运行态和共享 CPU 页池，排除本地 allocator。""" + for state in self.owner.get_bounded_attention_states(): + for pool in (state.big_pages_pool, state.small_pages_pool): + if pool is not None: + for tensor in vars(pool).values(): + if isinstance(tensor, torch.Tensor): + assert ( + tensor.device.type == "cpu" and tensor.is_shared() + ), "DP IPC requires shared CPU page pools" + ipc_copy = copy(self.owner) + ipc_copy.full_att_allocator = None + ipc_copy.bounded_att_big_page_allocator = None + ipc_copy.bounded_att_small_page_allocator = None + ipc_copy._reset_operators() + return ipc_copy + + def register_page_pools_pin_memory(self) -> None: + """在当前 GPU 映射所有共享 bounded 页池。""" + for state in self.owner.get_bounded_attention_states(): + if state.big_pages_pool is not None: + state.big_pages_pool.register_pin_memory() + if state.small_pages_pool is not None: + state.small_pages_pool.register_pin_memory() + + def copy_from_dp( + self, + remote_att_manager: AttentionStateManager, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + src_big_page_ids: List[int], + dst_big_page_ids: List[int], + src_req_idx: int, + dst_req_idx: int, + ) -> None: + """复制其他 DP 的 token、大小页和 prefill 基准运行态。""" + assert type(remote_att_manager) is type(self.owner), "DP replicas must use the same model resource manager" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + assert len(src_big_page_ids) == len(dst_big_page_ids), "source and destination page counts differ" + if src_token_indices.numel(): + for state in self.owner.get_full_attention_states(): + state.copy_tokens_from_dp( + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + remote_dp_full_att_state=getattr(remote_att_manager, state.state_name), + ) + if src_big_page_ids: + for state in self.owner.get_bounded_attention_states(): + state.copy_big_pages_from_dp( + src_page_ids=src_big_page_ids, + dst_page_ids=dst_big_page_ids, + remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), + ) + for state in self.owner.get_bounded_attention_states(): + state.copy_runtime_state_from_dp( + src_req_indices=[src_req_idx], + dst_req_indices=[dst_req_idx], + remote_dp_bounded_state=getattr(remote_att_manager, state.state_name), + ) diff --git a/lightllm/common/attention_state/full_att/__init__.py b/lightllm/common/attention_state/full_att/__init__.py new file mode 100644 index 0000000000..848fe47a51 --- /dev/null +++ b/lightllm/common/attention_state/full_att/__init__.py @@ -0,0 +1,59 @@ +from .mla.create_mla_state import create_mla_state +from .packed_mla.create_packed_mla_state import create_packed_mla_state +from .gqa.create_gqa_state import create_gqa_state +from .base_state import FullAttentionState +from .gqa.gqa_state import GQAAttentionConfig, GQAFullAttentionState +from .gqa import gqa_transfer +from .mla.mla_state import MLAConfig, MLAFullAttentionState + +__all__ = [ + "create_gqa_state", + "create_mla_state", + "create_packed_mla_state", + "FullAttentionState", + "GQAFullAttentionState", + "GQAAttentionConfig", + "MLAConfig", + "MLAFullAttentionState", +] + +from .gqa.group_quantized_state import ( + GroupQuantizedFullAttentionState, + GroupQuantizedAttentionConfig, +) +from .gqa.static_fp8_state import StaticFP8FullAttentionState, StaticFP8AttentionConfig +from .packed_mla.packed_mla_state import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, + PackedMLALayerView, +) + +__all__ += [ + "GroupQuantizedFullAttentionState", + "GroupQuantizedAttentionConfig", + "StaticFP8FullAttentionState", + "StaticFP8AttentionConfig", + "PackedMLAFullAttentionState", + "PackedMLAAttentionConfig", + "PackedMLALogicalConfig", + "PackedMLALayerView", +] + +from .compressed_att.compressed_att_state import ( + CompressedAttentionConfig, + CompressedAttentionState, + CompressedAttentionLayerView, +) +from .compressed_att.create_compressed_att_state import create_compressed_att_state + +__all__ += [ + "create_compressed_att_state", + "CompressedAttentionConfig", + "CompressedAttentionState", + "CompressedAttentionLayerView", +] + +from .packed_mla.fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState + +__all__ += ["FP8E8M0PackedMLAConfig", "FP8E8M0PackedMLAState"] diff --git a/lightllm/common/attention_state/full_att/base_state.py b/lightllm/common/attention_state/full_att/base_state.py new file mode 100644 index 0000000000..d3f00430b0 --- /dev/null +++ b/lightllm/common/attention_state/full_att/base_state.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from .interface import CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface + + +class FullAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): + """单种 full-attention 存储方案的数据接口，不分配 token 索引或维护请求映射。 + + 所有表示 token 索引的参数都须为调用方准备的 CUDA Tensor，位于对应存储设备上； + State 不转换设备、dtype 或连续布局，int32/int64 由具体算子支持，DP 目标索引须为 int64。 + CPU 页位于主机内存不改变 token 索引的 CUDA 契约。 + """ diff --git a/lightllm/common/attention_state/full_att/compressed_att/__init__.py b/lightllm/common/attention_state/full_att/compressed_att/__init__.py new file mode 100644 index 0000000000..4cbe7773f0 --- /dev/null +++ b/lightllm/common/attention_state/full_att/compressed_att/__init__.py @@ -0,0 +1,9 @@ +from .compressed_att_state import CompressedAttentionConfig, CompressedAttentionLayerView, CompressedAttentionState +from .create_compressed_att_state import create_compressed_att_state + +__all__ = [ + "CompressedAttentionConfig", + "CompressedAttentionLayerView", + "CompressedAttentionState", + "create_compressed_att_state", +] diff --git a/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py b/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py new file mode 100644 index 0000000000..21ce083697 --- /dev/null +++ b/lightllm/common/attention_state/full_att/compressed_att/compressed_att_state.py @@ -0,0 +1,427 @@ +"""序列压缩后的历史 KV，采用 FP8/E8M0、BF16 RoPE 与可选 indexer 的 packed 编码。""" + +from __future__ import annotations + +from functools import lru_cache + +from typing import TYPE_CHECKING, List +from dataclasses import dataclass +from math import lcm, prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args + +from ...utils import deepcopy_result +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState +from ..packed_mla.packed_transfer import transfer_packed_pages +from ..dp_transfer import copy_packed_tokens + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import StateTransferLayout + + +@dataclass(frozen=True) +class CompressedAttentionConfig: + # 外层顺序是物理 layer 轴；内层列出共享整份压缩 KV/indexer 存储的逻辑层。 + # 例如 ((0, 4), (8,)) 分配两份物理层存储，逻辑层 0/4 访问同一份数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # 以原始 token 为单位的压缩比例，不改变外部请求映射的索引域。 + compress_ratio: int + # 每个压缩 KV 条目编码前的向量维度，包含末尾的 RoPE 部分；不是存储字节数。 + kv_dim: int + # 保留 BF16 的 RoPE 维度；其余维度按每 64 个元素一组编码为 FP8，scale 使用 E8M0。 + rope_dim: int + # 每个物理块覆盖的原始 token 数，必须能被 compress_ratio 整除。 + # 块内先保存全部 KV 数据，再保存 scale 和可选 indexer；不能按单个压缩条目拆开搬运。 + block_token_size: int + # 可选 indexer key 的 FP8 维度，每个 key 另有一个 FP32 scale；0 表示不保存 indexer。 + indexer_dim: int = 0 + # 当前 State 所属 TP 组内的 rank；当前实现是单个共享 KV head，在各 rank 完整复制。 + # rank 0 写出共享 CPU/PD 页，每个 rank 均可读回，不按 TP 切分压缩条目的向量。 + tp_rank: int = 0 + # 当前模型 TP 组的 rank 数，只描述复制关系，不影响单个 rank 的存储字节数。 + tp_world_size: int = 1 + + @property + def kv_entry_bytes(self) -> int: + """单条压缩 KV 的数据与 scale 字节数，不包含块末尾 padding 或 indexer。""" + scale_count = (self.kv_dim - self.rope_dim) // 64 + return self.kv_dim + self.rope_dim + (scale_count + 7) // 8 * 8 + + @property + def kv_block_bytes(self) -> int: + """块内 KV 数据及 E8M0 scale 区的字节跨度，包含数据行对齐，不包含 indexer。""" + entries = self.block_token_size // self.compress_ratio + # FP8 占一字节，BF16 RoPE 占两字节；每条目的 scale 区补齐到八字节。 + data_bytes = self.kv_dim + self.rope_dim + return (entries * self.kv_entry_bytes + data_bytes - 1) // data_bytes * data_bytes + + @property + def token_slot_bytes(self) -> int: + """每个原始 token 索引摊销的物理字节数，包含 KV、scale、indexer 和块内 padding。""" + block_bytes = self.kv_block_bytes + if self.indexer_dim: + entries = self.block_token_size // self.compress_ratio + block_bytes += entries * (self.indexer_dim + torch.float32.itemsize) + # 块既要能均摊到原始 token 索引域，也要保留 packed KV 数据行的对齐。 + alignment = lcm(self.block_token_size, self.kv_dim + self.rope_dim) + return (block_bytes + alignment - 1) // alignment * alignment // self.block_token_size + + def get_attention_view(self, packed_kv: torch.Tensor) -> torch.Tensor: + """返回算子使用的 [block, entry, head, byte] 视图，保留数据/scale 分区和块 stride。""" + entries = self.block_token_size // self.compress_ratio + return packed_kv[:, : entries * self.kv_entry_bytes].view(-1, entries, 1, self.kv_entry_bytes) + + def __post_init__(self) -> None: + assert ( + self.compress_ratio > 0 and self.block_token_size > 0 and self.block_token_size % self.compress_ratio == 0 + ) + assert self.kv_dim > self.rope_dim >= 0 and (self.kv_dim - self.rope_dim) % 64 == 0 + assert self.indexer_dim >= 0 and 0 <= self.tp_rank < self.tp_world_size + assert self.logical_layer_groups and all(self.logical_layer_groups), "storage layer groups must be nonempty" + logical_layer_ids = [layer for group in self.logical_layer_groups for layer in group] + assert len(logical_layer_ids) == len( + set(logical_layer_ids) + ), "logical layers cannot repeat across storage groups" + + +@dataclass(frozen=True) +class CompressedAttentionLayerView: + # [压缩块, packed KV 字节]，保留整个块中数据与 scale 的分区格式。 + kv: torch.Tensor + # 可选的 [压缩块, key/scale 字节]；未配置 indexer 时返回 None。 + indexer_k: torch.Tensor | None + + +class CompressedAttentionState(FullAttentionState): + """保存一种序列压缩比例的历史 KV 和可选 indexer，不依赖具体模型配置。 + + 本 State 不执行压缩，不保存 SWA 或压缩器运行态；这些运行态由 bounded State 管理。 + 不同压缩方案各自持有存储，沿用原始 token 索引及共享分配器。 + 当前实现固定为单个共享 KV head，各 TP rank 完整复制；不执行 CSA/HCA 的历史选择。 + """ + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: CompressedAttentionConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """只准备压缩 KV 配置和层映射，存储由统一容量分配阶段创建。""" + self.config = layout_config + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.storage: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回序列压缩比例、FP8 分组及 E8M0 scale 编码对应的方案标记。""" + return f"compressed_c{self.config.compress_ratio}_fp8_group64_e8m0" + + def allocate_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" + # 外部 allocator 的每个物理页必须包含完整压缩块，不能拆开块内的数据和 scale。 + assert ( + self.page_size > 0 and self.page_size % self.config.block_token_size == 0 + ), "allocator page_size must be a positive multiple of block_token_size" + assert ( + token_capacity % self.config.block_token_size == 0 + ), "token capacity must contain complete compressed blocks" + self.capacity = token_capacity + # token 轴仍使用原始索引域；每 block_token_size 个槽位组成一个完整 packed 压缩块。 + self.storage = torch.zeros( + size=( + len(self.config.logical_layer_groups), + token_capacity + self.page_size, + 1, + self.config.token_slot_bytes, + ), + dtype=torch.uint8, + device=self.device, + ) + + def get_token_bytes(self) -> int: + """返回每个原始 token 索引的存储开销，包含全部物理层及压缩块 padding。""" + config = self.config + return len(config.logical_layer_groups) * config.token_slot_bytes + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int) -> CompressedAttentionLayerView: + """从同一物理压缩块取 KV 和 indexer 视图，不重排数据或创建副本。""" + config = self.config + layer = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] + blocks = layer.view(-1, config.block_token_size * config.token_slot_bytes) + entries = config.block_token_size // config.compress_ratio + kv_bytes = config.kv_block_bytes + indexer = ( + blocks[:, kv_bytes : kv_bytes + entries * (config.indexer_dim + torch.float32.itemsize)] + if config.indexer_dim + else None + ) + return CompressedAttentionLayerView(kv=blocks[:, :kv_bytes], indexer_k=indexer) + + def write_full_att_kv( + self, + logical_layer_id: int, + token_indices: torch.Tensor, + kv: torch.Tensor, + ) -> None: + """压缩器通过计算视图融合写入；普通 KV 写入路径不应调用此方法。""" + raise NotImplementedError("compressed KV is written by the fused compressor") + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整压缩块；调用方保证块内索引完整、顺序对应且源目标不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + # 数据、scale 和 indexer 在块内分区保存，不能只复制部分原始 token 槽位。 + assert ( + src_indices.numel() % self.config.block_token_size == 0 + ), "token copies require complete compressed blocks" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" + assert self.storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + copy_packed_tokens( + src_storage=self.storage, + dst_storage=self.storage, + src_token_indices=src_indices, + dst_token_indices=dst_indices, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小描述完整压缩块，包含 KV、scale、indexer 和块内 padding。""" + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + page_tokens = layout.tokens_per_page + assert ( + page_tokens is not None and page_tokens > 0 and page_tokens % self.config.block_token_size == 0 + ), "CPU page token capacity must be a positive multiple of block_token_size" + # 联合页可包含其他 State；只校验本 State 的区域完整且不越过单页末尾。 + page_bytes = prod(cpu_page_buffers.shape[1:]) * cpu_page_buffers.element_size() + required_bytes = self.get_token_bytes() * page_tokens + assert ( + 0 <= layout.byte_start < layout.byte_start + required_bytes <= layout.byte_end <= page_bytes + ), "compressed State byte range must fit within one CPU page" + assert all( + 0 <= page < cpu_page_buffers.shape[0] for page in cpu_page_indices + ), "CPU page index is out of bounds" + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + page_tokens = layout.tokens_per_page + assert ( + page_tokens is not None and page_tokens > 0 and page_tokens % self.config.block_token_size == 0 + ), "CPU page token capacity must be a positive multiple of block_token_size" + # 读回也允许联合页中的非零偏移；页面必须完整容纳当前 State 的全部压缩块。 + page_bytes = prod(cpu_page_buffers.shape[1:]) * cpu_page_buffers.element_size() + required_bytes = self.get_token_bytes() * page_tokens + assert ( + 0 <= layout.byte_start < layout.byte_start + required_bytes <= layout.byte_end <= page_bytes + ), "compressed State byte range must fit within one CPU page" + assert all( + 0 <= page < cpu_page_buffers.shape[0] for page in cpu_page_indices + ), "CPU page index is out of bounds" + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: CompressedAttentionState, + ) -> None: + """复制其他 DP 的完整压缩块；调用方保证块内索引完整且顺序对应。""" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + assert ( + src_token_indices.numel() % self.config.block_token_size == 0 + ), "DP token copies require complete compressed blocks" + copy_packed_tokens( + src_storage=remote_dp_full_att_state.storage, + dst_storage=self.storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + # packed scale/indexer 可位于末尾有效 token 之后，PD 必须传完整的末压缩块。 + # TODO：将末块索引扩展融合到传输 kernel，省去这份很小的索引临时区。 + alignment = self.config.block_token_size + padding = (-token_indices.numel()) % alignment + if token_indices.numel() and padding: + tail = token_indices[-1:] + torch.arange( + start=1, + end=padding + 1, + dtype=token_indices.dtype, + device=token_indices.device, + ) + token_indices = torch.cat(tensors=(token_indices, tail)) + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + # packed scale/indexer 可位于末尾有效 token 之后，PD 必须传完整的末压缩块。 + # TODO：将末块索引扩展融合到传输 kernel，省去这份很小的索引临时区。 + alignment = self.config.block_token_size + padding = (-token_indices.numel()) % alignment + if token_indices.numel() and padding: + tail = token_indices[-1:] + torch.arange( + start=1, + end=padding + 1, + dtype=token_indices.dtype, + device=token_indices.device, + ) + token_indices = torch.cat(tensors=(token_indices, tail)) + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """整个物理压缩块按原始 token 槽位传输，KV、scale、indexer 和 padding 一起保留。""" + config = self.config + assert ( + page_tokens > 0 and page_tokens % config.block_token_size == 0 + ), "CPU/PD page token capacity must be a positive multiple of block_token_size" + shape = (len(config.logical_layer_groups), page_tokens, 1, config.token_slot_bytes) + return ( + TransferRegion( + field="kv", + byte_start=0, + byte_end=prod(shape), + dtype=torch.uint8, + axes=("layer", "token", "head", "dim"), + shape=shape, + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=0, + length=1, + ), + AxisMapping( + axis="dim", + local_start=0, + external_start=0, + length=config.token_slot_bytes, + ), + ), + write_enabled=config.tp_rank == 0, + ), + ) diff --git a/lightllm/common/attention_state/full_att/compressed_att/create_compressed_att_state.py b/lightllm/common/attention_state/full_att/compressed_att/create_compressed_att_state.py new file mode 100644 index 0000000000..d01c080ae6 --- /dev/null +++ b/lightllm/common/attention_state/full_att/compressed_att/create_compressed_att_state.py @@ -0,0 +1,14 @@ +"""序列压缩 attention 历史 KV 的独立创建入口。""" + +import torch + +from .compressed_att_state import CompressedAttentionConfig, CompressedAttentionState + + +def create_compressed_att_state( + config: CompressedAttentionConfig, + page_size: int = 256, + device: str | torch.device = "cuda", +) -> CompressedAttentionState: + """根据压缩比例创建历史 KV 定义；不分配存储，也不创建 SWA/compressor 运行态。""" + return CompressedAttentionState(layout_config=config, page_size=page_size, device=device) diff --git a/lightllm/common/attention_state/full_att/dp_transfer.py b/lightllm/common/attention_state/full_att/dp_transfer.py new file mode 100644 index 0000000000..edcb8c2b5d --- /dev/null +++ b/lightllm/common/attention_state/full_att/dp_transfer.py @@ -0,0 +1,183 @@ +"""Full-attention DP 复制：显式传入各方案的存储，一次发射直接搬运完整 token。""" + +import torch +import triton +import triton.language as tl + +from ..utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +def copy_gqa_tokens( + src_k_storage: torch.Tensor, + src_v_storage: torch.Tensor, + dst_k_storage: torch.Tensor, + dst_v_storage: torch.Tensor, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, +) -> None: + """融合 GQA/静态 FP8 的 K/V 复制；静态标定参数不参与搬运。""" + _copy_token_storages( + src_storages=(src_k_storage, src_v_storage), + dst_storages=(dst_k_storage, dst_v_storage), + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + +def copy_group_quantized_tokens( + src_k_storage: torch.Tensor, + src_v_storage: torch.Tensor, + src_k_scales: torch.Tensor, + src_v_scales: torch.Tensor, + dst_k_storage: torch.Tensor, + dst_v_storage: torch.Tensor, + dst_k_scales: torch.Tensor, + dst_v_scales: torch.Tensor, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, +) -> None: + """融合 INT4/INT8 的 K/V payload 和动态 scale，分别保留字段 dtype。""" + _copy_token_storages( + src_storages=(src_k_storage, src_v_storage, src_k_scales, src_v_scales), + dst_storages=(dst_k_storage, dst_v_storage, dst_k_scales, dst_v_scales), + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + +def copy_packed_tokens( + src_storage: torch.Tensor, + dst_storage: torch.Tensor, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, +) -> None: + """复制单字段存储的完整 token 槽位，包含编码数据、scale、indexer 和 padding。 + + 压缩历史等方案的完整块约束由具体 State 在调用前检查，本函数不识别压缩比例。 + """ + _copy_token_storages( + src_storages=(src_storage,), + dst_storages=(dst_storage,), + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + +def _copy_token_storages( + src_storages: tuple[torch.Tensor, ...], + dst_storages: tuple[torch.Tensor, ...], + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, +) -> None: + """共享的四维 token 布局搬运，不创建 payload 或索引的中间 Tensor。 + + 在目标设备当前流发射；调用方保证远端存储及索引已可通过 peer/IPC 访问， + 并负责源数据就绪、存储生命周期和完成事件。目标 token 不重复，源目标区域不重叠。 + """ + token_count = src_token_indices.numel() + if token_count == 0: + return + + # 阶段一：先合并两侧连续的 head/width，保留各自 layer/token 容量和 stride。 + # head 轴保留为 1，沿用四维 token kernel；不连续时保持原布局。 + # 再按 8/4/2 字节搬运，保留 FP8、NaN、动态 scale 等原始位模式。 + sources, destinations = [], [] + for source, destination in zip(src_storages, dst_storages): + source, destination = merge_contiguous_tail_dims( + copies=[(source, destination)], + start_dim=-2, + keepdim=True, + )[0] + source, destination = view_max_copy_dtype(copies=[(source, destination)])[0] + sources.append(source) + destinations.append(destination) + + # 阶段二：Tensor tuple 在编译时展开为独立指针参数，不上传指针表。 + # 一次 launch 完成当前方案全部字段；每个 program 用跳步循环覆盖任意 token 数。 + # 保持 24 个 program，使用 4096 个搬运单元提高跨卡读取吞吐，循环流水深度为 1。 + _copy_indexed_token_storages[(24,)]( + SOURCES=tuple(sources), + DESTINATIONS=tuple(destinations), + SRC_TOKEN_IDS=src_token_indices, + DST_TOKEN_IDS=dst_token_indices, + token_count=token_count, + SRC_INDEX_STRIDE=src_token_indices.stride(0), + DST_INDEX_STRIDE=dst_token_indices.stride(0), + SHAPES=tuple((tensor.shape[0], tensor.shape[2], tensor.shape[3]) for tensor in sources), + SRC_STRIDES=tuple(tensor.stride() for tensor in sources), + DST_STRIDES=tuple(tensor.stride() for tensor in destinations), + ELEMENT_BYTES=tuple(tensor.element_size() for tensor in sources), + BLOCK_ELEMENTS=4096, + NUM_STAGES=1, + num_warps=8, + ) + + +@triton.jit +def _copy_indexed_token_storages( + # 各字段的四维存储指针；K/V、scale 或 packed 行在一次 launch 中处理。 + SOURCES, + DESTINATIONS, + SRC_TOKEN_IDS, + DST_TOKEN_IDS, + token_count, + # 索引 Tensor 可以非连续，直接使用调用方提供的 int32/int64 索引。 + SRC_INDEX_STRIDE: tl.constexpr, + DST_INDEX_STRIDE: tl.constexpr, + SHAPES: tl.constexpr, # 各字段的 (layer_count, head_count, width)。 + SRC_STRIDES: tl.constexpr, + DST_STRIDES: tl.constexpr, + ELEMENT_BYTES: tl.constexpr, + BLOCK_ELEMENTS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + token_count = tl.cast(token_count, tl.int64) + program_start = tl.program_id(0).to(tl.int64) * BLOCK_ELEMENTS + grid_stride = tl.num_programs(0).to(tl.int64) * BLOCK_ELEMENTS + lane = tl.arange(0, BLOCK_ELEMENTS).to(tl.int64) + + for field in tl.static_range(len(SHAPES)): + # 每个字段独立计算地址；不同容量下 layer stride 不同，不能复用源端 stride。 + layers = SHAPES[field][0] + heads = SHAPES[field][1] + width = SHAPES[field][2] + if ELEMENT_BYTES[field] == 8: + copy_dtype = tl.uint64 + elif ELEMENT_BYTES[field] == 4: + copy_dtype = tl.uint32 + elif ELEMENT_BYTES[field] == 2: + copy_dtype = tl.uint16 + else: + copy_dtype = tl.uint8 + source = SOURCES[field].to(tl.pointer_type(copy_dtype)) + destination = DESTINATIONS[field].to(tl.pointer_type(copy_dtype)) + element_count = token_count * layers * heads * width + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): + # 阶段一：按 [token, layer, head, width] 展开；除数均为布局常量， + # 避免按动态 token_count 做逐元素的 int64 除法。 + element = block_start + lane + valid = element < element_count + column = element % width + head = element // width % heads + layer = element // (width * heads) % layers + token = element // (width * heads * layers) + + # 阶段二：读取两端物理槽位；索引值及后续地址乘法使用 int64 防止溢出。 + src_token = tl.load(SRC_TOKEN_IDS + token * SRC_INDEX_STRIDE, mask=valid, other=0).to(tl.int64) + dst_token = tl.load(DST_TOKEN_IDS + token * DST_INDEX_STRIDE, mask=valid, other=0).to(tl.int64) + src_offset = ( + layer * tl.cast(SRC_STRIDES[field][0], tl.int64) + + src_token * tl.cast(SRC_STRIDES[field][1], tl.int64) + + head * tl.cast(SRC_STRIDES[field][2], tl.int64) + + column * tl.cast(SRC_STRIDES[field][3], tl.int64) + ) + dst_offset = ( + layer * tl.cast(DST_STRIDES[field][0], tl.int64) + + dst_token * tl.cast(DST_STRIDES[field][1], tl.int64) + + head * tl.cast(DST_STRIDES[field][2], tl.int64) + + column * tl.cast(DST_STRIDES[field][3], tl.int64) + ) + + # 阶段三：直接从远端存储读出并写入目标槽位，无 gather 或跨设备中转存储。 + value = tl.load(source + src_offset, mask=valid, other=0) + tl.store(destination + dst_offset, value, mask=valid) diff --git a/lightllm/common/attention_state/full_att/gqa/__init__.py b/lightllm/common/attention_state/full_att/gqa/__init__.py new file mode 100644 index 0000000000..5ef32afe8e --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/__init__.py @@ -0,0 +1,14 @@ +from .create_gqa_state import create_gqa_state +from .gqa_state import GQAAttentionConfig, GQAFullAttentionState +from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState +from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState + +__all__ = [ + "create_gqa_state", + "GQAAttentionConfig", + "GQAFullAttentionState", + "GroupQuantizedAttentionConfig", + "GroupQuantizedFullAttentionState", + "StaticFP8AttentionConfig", + "StaticFP8FullAttentionState", +] diff --git a/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py b/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py new file mode 100644 index 0000000000..d25a7c37fc --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/create_gqa_state.py @@ -0,0 +1,56 @@ +"""GQA（含 MHA/MQA）的统一创建入口，选择普通精度或量化 State。""" + +from __future__ import annotations + +from typing import Literal + +import torch + +from .gqa_state import GQAAttentionConfig, GQAFullAttentionState +from .group_quantized_state import GroupQuantizedAttentionConfig, GroupQuantizedFullAttentionState +from .static_fp8_state import StaticFP8AttentionConfig, StaticFP8FullAttentionState + + +def create_gqa_state( + config: GQAAttentionConfig, + kv_type: Literal["None", "int4kv", "int8kv", "fp8kv_spt", "fp8kv_sph"] = "None", + page_size: int = 1, + device: str | torch.device = "cuda", + group_size: int = 8, +) -> GQAFullAttentionState | GroupQuantizedFullAttentionState | StaticFP8FullAttentionState: + """按编码参数创建 GQA（含 MHA/MQA）State，普通 token 容量由后续流程分配。 + + config 描述逻辑 KV 维度和 TP 信息；kv_type 沿用启动参数中的编码名称。 + group_size 仅用于 INT4/INT8。静态 FP8 只创建布局，标定参数必须通过后续 + bind_calibration 绑定；本入口不读取配置文件或绑定模型角色。 + """ + if kv_type == "None": + return GQAFullAttentionState(layout_config=config, page_size=page_size, device=device) + if kv_type in ("int4kv", "int8kv"): + quant_config = GroupQuantizedAttentionConfig( + logical_layer_groups=config.logical_layer_groups, + global_head_num=config.global_head_num, + head_dim=config.head_dim, + scale_dtype=config.dtype, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + replicated=config.replicated, + bits=4 if kv_type == "int4kv" else 8, + group_size=group_size, + ) + return GroupQuantizedFullAttentionState(layout_config=quant_config, page_size=page_size, device=device) + assert kv_type in ("fp8kv_spt", "fp8kv_sph"), f"unsupported GQA encoding: {kv_type}" + quant_config = StaticFP8AttentionConfig( + logical_layer_groups=config.logical_layer_groups, + global_head_num=config.global_head_num, + head_dim=config.head_dim, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + replicated=config.replicated, + scale_mode="per_head" if kv_type == "fp8kv_sph" else "per_tensor", + ) + return StaticFP8FullAttentionState( + layout_config=quant_config, + page_size=page_size, + device=device, + ) diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_state.py b/lightllm/common/attention_state/full_att/gqa/gqa_state.py new file mode 100644 index 0000000000..64cc0e384b --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -0,0 +1,410 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, List +from math import prod +from dataclasses import dataclass + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState +from .gqa_transfer import transfer_gqa_pages +from ..dp_transfer import copy_gqa_tokens + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +@dataclass(frozen=True) +class GQAAttentionConfig: + """普通 GQA（含 MHA/MQA）的逻辑 KV 布局；不描述量化编码或物理行 padding。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # TP 切分前单份 K 或 V 的 head 数，不将 K/V 相加，也不是 query head 数。 + global_head_num: int + # 每个 K/V head 编码前的特征维度，量化位宽不改变此逻辑维度。 + head_dim: int + # K/V 浮点存储精度。 + dtype: torch.dtype + # 所属模型 TP 组内的 rank，不是全局 rank 或 DP rank。 + tp_rank: int = 0 + # 所属模型 TP 组大小，用于确定分片或副本的访问关系。 + tp_world_size: int = 1 + # True 时各 rank 保存全部 heads；False 时切分，heads 少于 ranks 时相邻 ranks 复制同一 head。 + replicated: bool = False + + @property + def head_num(self) -> int: + """当前 TP rank 保存的单份 K 或 V 的 head 数。""" + return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) + + def __post_init__(self) -> None: + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all( + isinstance(layer, int) and layer >= 0 for layer in layers + ), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + assert ( + self.replicated + or self.global_head_num % self.tp_world_size == 0 + or self.tp_world_size % self.global_head_num == 0 + ), "KV heads require complete TP shards or replicas" + + +class GQAFullAttentionState(FullAttentionState): + """GQA 的 full-attention KV 存储，MHA 和 MQA 均为其特例。 + + TP 切分前，H_q 个 query head 分为 H_kv 组，每组共享一对 K/V head， + 要求 H_q 能被 H_kv 整除。MHA 对应 H_kv == H_q；MQA 对应 H_kv == 1； + 其余 GQA 对应 1 < H_kv < H_q。 + + config.head_num 表示 TP 切分或复制后的本地 KV head 数。 + K/V 分别为 [layers, token_slots, local_kv_heads, head_dim]，两者独立分配。 + query 与 KV 的分组计算由 attention backend 处理，因此三种情况共用此 State。 + """ + + # 保留既有 K/V 编码的稳定传输布局标识。 + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: GQAAttentionConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """只记录布局配置，KV 与 HOLD 页由后续统一容量初始化分配。""" + self.config = layout_config + self.storage_dtype = layout_config.dtype # 实际 KV 编码类型；量化方案显式设置，区别于 scale 精度。 + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.k_storage: torch.Tensor | None = None + self.v_storage: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回 GQA KV 精度对应的存储方案标记。""" + return "gqa_" + dtype_name(dtype=self.config.dtype) + + def allocate_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" + self.capacity = token_capacity + shape = self._storage_shape() + self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + + def get_token_bytes(self) -> int: + """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" + config = self.config + return len(config.logical_layer_groups) * 2 * config.head_num * config.head_dim * self.storage_dtype.itemsize + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) + return self.k_storage[storage_layer_index], self.v_storage[storage_layer_index] + + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """通过 CUDA 算子将形状匹配的 K/V 数据写入指定物理 token 位置。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) + if not len(token_indices): + return + assert self.k_storage.is_cuda, "KV writes require CUDA storage" + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。 + destindex_copy_kv( + K=kv[:, : self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index] + ) + destindex_copy_kv( + K=kv[:, self.config.head_num :], DestLoc=token_indices, Out=self.v_storage[storage_layer_index] + ) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" + assert self.k_storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + # TODO：将 K/V 的同索引 token 复制融合为一次 kernel 发射， + # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.k_storage, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.v_storage, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局将 token 数据写到指定 CPU 页，覆盖本 State 负责的字段和 TP 分片。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_gqa_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """从指定 CPU 页读回 token 数据，保留目标中由其他资源或 TP 分片持有的区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_gqa_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: GQAFullAttentionState, + ) -> None: + """复制其他 DP 同一 TP 分片的完整 token 数据，层顺序保持一致。""" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + copy_gqa_tokens( + src_k_storage=remote_dp_full_att_state.k_storage, + src_v_storage=remote_dp_full_att_state.v_storage, + dst_k_storage=self.k_storage, + dst_v_storage=self.v_storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按 token 索引将 K/V 直接写入指定 PD 页，网络发送与同步由调用方负责。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_gqa_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """将指定 PD 页直接恢复到 K/V 物理槽位，不分配中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_gqa_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.state_name, + config.tp_rank, + config.tp_world_size, + page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + field_specs = ( + ("k", self.storage_dtype, config.head_dim, "dim"), + ("v", self.storage_dtype, config.head_dim, "dim"), + ) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) + global_shape = list(local_shape) + global_shape[axes.index("token")] = page_tokens + head_axis = axes.index("head") + # K/V 各自是独立字段：这里读取的是单个字段的本地 head 数。 + local_head_count = global_shape[head_axis] + # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 + # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 + # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 + global_shape[head_axis] = config.global_head_num + + # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 + # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 + replica_count = config.tp_world_size * local_head_count // config.global_head_num + # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 + head_start = config.tp_rank // replica_count * local_head_count + # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 + write_enabled = config.tp_rank % replica_count == 0 + axis_mappings = tuple( + AxisMapping( + axis=axis, + local_start=0, + external_start=head_start if axis == "head" else 0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_mappings=axis_mappings, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", + self.state_name, + field_name, + head_start, + local_head_count, + replica_count, + write_enabled, + field_region.byte_start, + field_region.byte_end, + ) + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", + self.state_name, + len(regions), + byte_offset, + ) + return tuple(regions) + + def _storage_shape(self) -> tuple[int, ...]: + return ( + len(self.config.logical_layer_groups), + self.capacity + self.page_size, + self.config.head_num, + self.config.head_dim, + ) diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py b/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py new file mode 100644 index 0000000000..516c947864 --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/gqa_transfer.py @@ -0,0 +1,254 @@ +"""GQA 的统一 CPU/PD 页传输：按 token 和页索引直接交换 K/V 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_gqa_pages( + k_storage: torch.Tensor, + v_storage: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + layout: StateTransferLayout, + write: bool, + grid_num: int = 24, + page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, +) -> None: + """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 + """ + assert grid_num > 0, "grid_num must be positive" + assert k_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 K/V。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 K/V 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + axes = ("layer", "token", "head", "dim") + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "k": + local = k_storage + elif region.field == "v": + local = v_storage + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, dim]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], + external.shape[3], + external.shape[4], + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个 K/V 字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 K/V 的逐字段 launch 融合， + # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 + for local, external in copy_views: + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_gqa_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_gqa_pages( + # 数据指针：一次处理 K 或 V；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, dim]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 每个 head 参与拷贝的元素数。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + storage_layer_index = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 K/V 字段的大小。 + local_offset = ( + storage_layer_index * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + storage_layer_index * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py new file mode 100644 index 0000000000..28b84c9551 --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/group_quantized_state.py @@ -0,0 +1,478 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, List +from dataclasses import dataclass +from math import prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState +from .group_quantized_transfer import transfer_group_quantized_pages +from ..dp_transfer import copy_group_quantized_tokens + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +@dataclass(frozen=True) +class GroupQuantizedAttentionConfig: + """INT4/INT8 GQA 的独立布局；KV 编码为整数，动态 scale 使用单独精度。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # TP 切分前单份 K 或 V 的 head 数，不将 K/V 相加，也不是 query head 数。 + global_head_num: int + # 每个 K/V head 编码前的特征维度，量化位宽不改变此逻辑维度。 + head_dim: int + # 每个 token、每组 K/V 量化 scale 的浮点精度，不是 KV 编码类型。 + scale_dtype: torch.dtype + # 所属模型 TP 组内的 rank，不是全局 rank 或 DP rank。 + tp_rank: int = 0 + # 所属模型 TP 组大小，用于确定分片或副本的访问关系。 + tp_world_size: int = 1 + # True 时各 rank 保存全部 heads；False 时切分，heads 少于 ranks 时相邻 ranks 复制同一 head。 + replicated: bool = False + # KV 量化位数；INT4 将两个编码值打包进一个字节。 + bits: int = 8 + # 沿 head_dim 分组，每组共用一个动态 scale；必须为二次幂且整除 head_dim。 + group_size: int = 8 + + @property + def head_num(self) -> int: + """当前 TP rank 保存的单份 K 或 V 的 head 数。""" + return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) + + def __post_init__(self) -> None: + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all( + isinstance(layer, int) and layer >= 0 for layer in layers + ), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + assert ( + self.replicated + or self.global_head_num % self.tp_world_size == 0 + or self.tp_world_size % self.global_head_num == 0 + ), "KV heads require complete TP shards or replicas" + assert ( + self.bits in (4, 8) and self.group_size >= 2 and not (self.group_size & (self.group_size - 1)) + ), "INT4/INT8 quantization requires a power-of-two group size" + assert self.head_dim % self.group_size == 0 and self.scale_dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ), "group scales require a floating dtype and complete groups" + + +class GroupQuantizedFullAttentionState(FullAttentionState): + """管理 INT4/INT8 KV 与逐 token 分组 scale，复制和传输始终覆盖二者。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: GroupQuantizedAttentionConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """按分组量化配置准备 KV 与动态 scale 存储。""" + self.config = layout_config + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.k_storage: torch.Tensor | None = None + self.v_storage: torch.Tensor | None = None + self.storage_dtype = torch.int8 + self.k_scales: torch.Tensor | None = None + self.v_scales: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回量化位数、分组大小及 scale 精度对应的存储方案标记。""" + return ( + f"gqa_int{self.config.bits}_group{self.config.group_size}_{dtype_name(dtype=self.config.scale_dtype)}_scale" + ) + + def allocate_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配独立的量化 KV 与配套动态 scale。""" + self.capacity = token_capacity + shape = self._storage_shape() + self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + shape = ( + len(self.config.logical_layer_groups), + token_capacity + self.page_size, + self.config.head_num, + self.config.head_dim // self.config.group_size, + ) + self.k_scales = torch.empty(size=shape, dtype=self.config.scale_dtype, device=self.device) + self.v_scales = torch.empty(size=shape, dtype=self.config.scale_dtype, device=self.device) + + def get_token_bytes(self) -> int: + """累计量化 KV 编码与逐 token 分组 scale 的开销，无需分配存储。""" + config = self.config + head_bytes = config.head_dim * config.bits // 8 + head_bytes += config.head_dim // config.group_size * config.scale_dtype.itemsize + return len(config.logical_layer_groups) * 2 * config.head_num * head_bytes + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view( + self, logical_layer_id: int + ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: + """返回该层 K/V 编码数据及各自的分组 scale 视图。""" + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) + return (self.k_storage[storage_layer_index], self.k_scales[storage_layer_index]), ( + self.v_storage[storage_layer_index], + self.v_scales[storage_layer_index], + ) + + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """量化输入浮点 K/V，并将编码数据和分组 scale 写入相同 token 位置。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) + if not len(token_indices): + return + assert self.k_storage.is_cuda, "KV writes require CUDA storage" + # TODO：融合 K/V 两次量化写入，同时输出各自编码和动态 scale，保持分组及舍入语义。 + # K/V 分别写入独立存储，各自携带同一 token 的动态 scale。 + for offset, storage, scales_buffer in ( + (0, self.k_storage, self.k_scales), + (self.config.head_num, self.v_storage, self.v_scales), + ): + values = kv[:, offset : offset + self.config.head_num] + if self.config.bits == 4: + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import destindex_copy_int4kv + + destindex_copy_int4kv( + KV=values, + DestLoc=token_indices, + KV_buffer=storage[storage_layer_index], + KV_scale_buffer=scales_buffer[storage_layer_index], + quant_group_size=self.config.group_size, + ) + else: + from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import ( + destindex_copy_quantize_kv, + ) + + destindex_copy_quantize_kv( + K=values, + DestLoc=token_indices, + Out=storage[storage_layer_index], + Out_scale=scales_buffer[storage_layer_index], + quant_group_dim=self.config.group_size, + ) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" + assert self.k_storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + # TODO：将 K/V 编码及各自的动态 scale 的同索引 token 复制融合为一次 kernel 发射， + # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.k_storage, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.v_storage, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.k_scales, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.v_scales, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_group_quantized_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + k_scales=self.k_scales, + v_scales=self.v_scales, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_group_quantized_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + k_scales=self.k_scales, + v_scales=self.v_scales, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: GroupQuantizedFullAttentionState, + ) -> None: + """校验量化位数和分组大小，再复制 KV 与动态 scale。""" + assert ( + remote_dp_full_att_state.config.bits, + remote_dp_full_att_state.config.group_size, + ) == (self.config.bits, self.config.group_size), "incompatible group quantization" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + copy_group_quantized_tokens( + src_k_storage=remote_dp_full_att_state.k_storage, + src_v_storage=remote_dp_full_att_state.v_storage, + src_k_scales=remote_dp_full_att_state.k_scales, + src_v_scales=remote_dp_full_att_state.v_scales, + dst_k_storage=self.k_storage, + dst_v_storage=self.v_storage, + dst_k_scales=self.k_scales, + dst_v_scales=self.v_scales, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_group_quantized_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + k_scales=self.k_scales, + v_scales=self.v_scales, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_group_quantized_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + k_scales=self.k_scales, + v_scales=self.v_scales, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.state_name, + config.tp_rank, + config.tp_world_size, + page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + field_specs = ( + ("k", torch.int8, config.head_dim * config.bits // 8, "dim"), + ("v", torch.int8, config.head_dim * config.bits // 8, "dim"), + ("k_scale", config.scale_dtype, config.head_dim // config.group_size, "group"), + ("v_scale", config.scale_dtype, config.head_dim // config.group_size, "group"), + ) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) + global_shape = list(local_shape) + global_shape[axes.index("token")] = page_tokens + head_axis = axes.index("head") + # K/V 各自是独立字段：这里读取的是单个字段的本地 head 数。 + local_head_count = global_shape[head_axis] + # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 + # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 + # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 + global_shape[head_axis] = config.global_head_num + + # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 + # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 + replica_count = config.tp_world_size * local_head_count // config.global_head_num + # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 + head_start = config.tp_rank // replica_count * local_head_count + # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 + write_enabled = config.tp_rank % replica_count == 0 + axis_mappings = tuple( + AxisMapping( + axis=axis, + local_start=0, + external_start=head_start if axis == "head" else 0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_mappings=axis_mappings, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", + self.state_name, + field_name, + head_start, + local_head_count, + replica_count, + write_enabled, + field_region.byte_start, + field_region.byte_end, + ) + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", + self.state_name, + len(regions), + byte_offset, + ) + return tuple(regions) + + def _storage_shape(self) -> tuple[int, ...]: + return ( + len(self.config.logical_layer_groups), + self.capacity + self.page_size, + self.config.head_num, + self.config.head_dim * self.config.bits // 8, + ) diff --git a/lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py b/lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py new file mode 100644 index 0000000000..72eea0ba6b --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/group_quantized_transfer.py @@ -0,0 +1,264 @@ +"""INT4/INT8 分组量化 的统一 CPU/PD 页传输：按 token 和页索引直接交换 K/V、动态 scale 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_group_quantized_pages( + k_storage: torch.Tensor, + v_storage: torch.Tensor, + k_scales: torch.Tensor, + v_scales: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + layout: StateTransferLayout, + write: bool, + grid_num: int = 24, + page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, +) -> None: + """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 + """ + assert grid_num > 0, "grid_num must be positive" + assert k_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 K/V。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 K/V 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "k": + local = k_storage + axes = ("layer", "token", "head", "dim") + elif region.field == "v": + local = v_storage + axes = ("layer", "token", "head", "dim") + elif region.field == "k_scale": + # 动态 scale 与 K/V 使用相同的 token/page 索引。 + local = k_scales + axes = ("layer", "token", "head", "group") + elif region.field == "v_scale": + local = v_scales + axes = ("layer", "token", "head", "group") + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], + external.shape[3], + external.shape[4], + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个传输字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 K/V 及动态 scale 的逐字段 launch 融合， + # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 + for local, external in copy_views: + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_group_quantized_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_group_quantized_pages( + # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 字段末轴参与拷贝的元素数，如 dim、group 或 scale。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + storage_layer_index = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 + local_offset = ( + storage_layer_index * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + storage_layer_index * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py new file mode 100644 index 0000000000..20a79bd917 --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/static_fp8_state.py @@ -0,0 +1,487 @@ +from __future__ import annotations + +from functools import lru_cache + +from typing import TYPE_CHECKING, List +from collections.abc import Mapping +from dataclasses import dataclass +from math import prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState +from .static_fp8_transfer import transfer_static_fp8_pages +from ..dp_transfer import copy_gqa_tokens + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +@dataclass(frozen=True) +class StaticFP8AttentionConfig: + """静态 FP8 GQA 的逻辑布局；KV 固定为 FP8 编码，标定 Tensor 由后续流程绑定。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # TP 切分前单份 K 或 V 的 head 数，不将 K/V 相加，也不是 query head 数。 + global_head_num: int + # 每个 K/V head 编码前的特征维度，量化位宽不改变此逻辑维度。 + head_dim: int + # 所属模型 TP 组内的 rank，不是全局 rank 或 DP rank。 + tp_rank: int = 0 + # 所属模型 TP 组大小，用于确定分片或副本的访问关系。 + tp_world_size: int = 1 + # True 时各 rank 保存全部 heads；False 时切分，heads 少于 ranks 时相邻 ranks 复制同一 head。 + replicated: bool = False + # 每层 K/V 各一个 scale，或每个本地 K/V head 各一个；不占用逐 token 页空间。 + scale_mode: str = "per_tensor" + + @property + def head_num(self) -> int: + """当前 TP rank 保存的单份 K 或 V 的 head 数。""" + return self.global_head_num if self.replicated else max(1, self.global_head_num // self.tp_world_size) + + def __post_init__(self) -> None: + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all( + isinstance(layer, int) and layer >= 0 for layer in layers + ), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.global_head_num > 0 and self.head_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + assert ( + self.replicated + or self.global_head_num % self.tp_world_size == 0 + or self.tp_world_size % self.global_head_num == 0 + ), "KV heads require complete TP shards or replicas" + assert self.scale_mode in ("per_tensor", "per_head"), "invalid static FP8 scale mode" + + +class StaticFP8FullAttentionState(FullAttentionState): + """使用静态标定参数的 FP8 KV 存储，固定 scale 随 State 初始化，不逐 token 搬运。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: StaticFP8AttentionConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """只准备静态 FP8 布局；标定参数由后续 bind_calibration 统一绑定。""" + self.scales: torch.Tensor | None = None + self.q_scales: torch.Tensor | None = None + self.cpu_scales = None # 实际存储初始化时准备计算 backend 使用的 CPU 标定值。 + self.config = layout_config + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.k_storage: torch.Tensor | None = None + self.v_storage: torch.Tensor | None = None + self.storage_dtype = torch.uint8 + + def bind_calibration( + self, + calibration: Mapping[str, object] | None, + storage_layer_start: int, + ) -> None: + """在分配 KV 前绑定标定表中本 State 的物理层区间，并切出当前 TP 的 head。 + + 行起点由管理器按 main、draft0..draft35 的完整 State 集合确定，不是逻辑层号。 + 同组逻辑层共用一行；State 内的行顺序与 logical_layer_groups 的物理层顺序一致。 + 无文件时仅 per_tensor 可使用单位 scale；per_head 必须有 K/V 标定，Q 标定可选。 + 未提供 Q 标定时由 attention backend 动态量化 Q。 + """ + assert ( + self.k_storage is None and self.scales is None + ), "calibration must be bound once before storage allocation" + config = self.config + storage_layer_count = len(config.logical_layer_groups) + if calibration is None: + assert config.scale_mode == "per_tensor", "per-head FP8 requires a KV calibration file" + scales = torch.ones(size=(storage_layer_count, 2), dtype=torch.float32) + q_scales = None + else: + # 阶段一：检查文件编码和物理行覆盖范围。 + calibrated_layer_count = calibration["num_layers"] + storage_layer_end = storage_layer_start + storage_layer_count + assert ( + 0 <= storage_layer_start < storage_layer_end <= calibrated_layer_count + ), "calibration does not cover this State's physical layers" + assert ( + calibration.get("quant_type", config.scale_mode) == config.scale_mode + ), "calibration quantization mode does not match the State" + limits = torch.finfo(torch.float8_e4m3fn) + assert ( + calibration.get("qmin", limits.min) == limits.min and calibration.get("qmax", limits.max) == limits.max + ), "calibration range does not match float8_e4m3fn" + + def scale_tensor(values: Mapping[str, object], shape: tuple[int, ...]) -> torch.Tensor: + tensor = torch.tensor(data=values["scales"], dtype=torch.float32) + assert tuple(values["scales_shape"]) == tensor.shape == shape, "invalid calibration scale shape" + assert bool( + (torch.isfinite(tensor) & (tensor > 0)).all() + ), "calibration scales must be finite and positive" + return tensor[storage_layer_start:storage_layer_end] + + def shard(tensor: torch.Tensor, calibrated_heads: int) -> torch.Tensor: + heads = config.head_num + assert ( + calibrated_heads > 0 and (heads * config.tp_world_size) % calibrated_heads == 0 + ), "calibration heads do not cover the model's KV heads" + expanded = tensor.repeat_interleave(repeats=heads * config.tp_world_size // calibrated_heads, dim=-1) + return expanded[:, config.tp_rank * heads : (config.tp_rank + 1) * heads].contiguous() + + # 阶段二：先按物理层取行，再分别切分 K、V、Q 的 head；标定 Tensor 暂留 CPU。 + if config.scale_mode == "per_head": + calibrated_heads = calibration["num_head"] + all_scales = scale_tensor(calibration, (calibrated_layer_count, 2 * calibrated_heads)) + k, v = all_scales.split(split_size=calibrated_heads, dim=-1) + scales = torch.cat(tensors=(shard(k, calibrated_heads), shard(v, calibrated_heads)), dim=-1) + q_config = calibration.get("q_calibration") + q_scales = None + if q_config is not None: + q_scales = shard( + scale_tensor(q_config, (calibrated_layer_count, q_config["num_head"])), + q_config["num_head"], + ) + else: + scales = scale_tensor(calibration, (calibrated_layer_count, 2)) + q_scales = None + + # 阶段三：绑定标定数值；后续 allocate_storage 将其迁移到实际设备。 + self.scales = scales + self.q_scales = q_scales + + def get_state_type(self) -> str: + """返回包含 per-tensor 或 per-head 标定方式的 FP8 方案标记。""" + return "gqa_fp8_static_" + self.config.scale_mode + + def allocate_storage(self, token_capacity: int) -> None: + """绑定 KV 与真实标定参数，并预备 FlashInfer 使用的 CPU scale。""" + assert self.scales is not None, "static FP8 calibration must be bound before allocation" + self.capacity = token_capacity + shape = self._storage_shape() + self.k_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + self.v_storage = torch.empty(size=shape, dtype=self.storage_dtype, device=self.device) + # 统一标定阶段保留 CPU 数值，实际分配时再迁移到 KV 所在设备。 + if self.k_storage.device.type != "meta": + self.scales = self.scales.to(device=self.k_storage.device) + if self.q_scales is not None: + self.q_scales = self.q_scales.to(device=self.k_storage.device) + # 提前准备 CPU scale，避免计算热路径同步设备。 + self.cpu_scales = self.scales.detach().cpu().numpy() + + def get_token_bytes(self) -> int: + """由配置累计当前 TP rank 的 K/V 单 token 开销，无需分配存储。""" + config = self.config + return len(config.logical_layer_groups) * 2 * config.head_num * config.head_dim * self.storage_dtype.itemsize + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询 KV 和标定参数共用的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号返回 K、V 视图，head 数使用当前 TP 分片配置。""" + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) + return self.k_storage[storage_layer_index], self.v_storage[storage_layer_index] + + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """按固定 per-tensor 或 per-head scale 量化输入 K/V，写入 FP8 编码。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) + if not len(token_indices): + return + assert self.k_storage.is_cuda, "KV writes require CUDA storage" + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8 + + destindex_copy_kv_fp8( + KV=kv, + DestLoc=token_indices, + scales=self.scales[storage_layer_index], + Out=self.k_storage[storage_layer_index].view(torch.float8_e4m3fn), + is_per_tensor_quant=self.config.scale_mode == "per_tensor", + Out_v=self.v_storage[storage_layer_index].view(torch.float8_e4m3fn), + ) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" + assert self.k_storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + # TODO：将 FP8 K/V 原始编码 的同索引 token 复制融合为一次 kernel 发射， + # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.k_storage, + ) + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.v_storage, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_static_fp8_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_static_fp8_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: StaticFP8FullAttentionState, + ) -> None: + """复制相同配置的 FP8 KV 编码；各 DP 使用相同的静态标定参数。""" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + copy_gqa_tokens( + src_k_storage=remote_dp_full_att_state.k_storage, + src_v_storage=remote_dp_full_att_state.v_storage, + dst_k_storage=self.k_storage, + dst_v_storage=self.v_storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_static_fp8_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_static_fp8_pages( + k_storage=self.k_storage, + v_storage=self.v_storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.state_name, + config.tp_rank, + config.tp_world_size, + page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + field_specs = ( + ("k", self.storage_dtype, config.head_dim, "dim"), + ("v", self.storage_dtype, config.head_dim, "dim"), + ) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.logical_layer_groups), page_tokens, config.head_num, field_width) + global_shape = list(local_shape) + global_shape[axes.index("token")] = page_tokens + head_axis = axes.index("head") + # K/V 各自是独立字段：这里读取的是单个字段的本地 head 数。 + local_head_count = global_shape[head_axis] + # 外部页保存该字段未做 TP 切分的完整 heads，TP 副本只占一份区域。 + # K/V 各生成一个 region 并分别累加字节数，因此这里不能再乘 2；动态 scale 同理。 + # 例如全局 H=8、TP=2：本地 k/v 各 4 heads，外部 k/v 区域各 8 heads。 + global_shape[head_axis] = config.global_head_num + + # 阶段二：确定本地 heads 的全局起点；同组副本只有首 rank 写出，所有副本都可读回。 + # replica_count 表示同一份 head 分片被多少个 ranks 复制：纯切分时为 1，全复制时为 TP。 + replica_count = config.tp_world_size * local_head_count // config.global_head_num + # 将组内 rank 换算为去重后的分片编号，再定位当前字段在全局 head 轴上的起点。 + head_start = config.tp_rank // replica_count * local_head_count + # 副本共享同一外部区域，只让组内第一个 rank 写出，避免重复写入；读取不受此标记限制。 + write_enabled = config.tp_rank % replica_count == 0 + axis_mappings = tuple( + AxisMapping( + axis=axis, + local_start=0, + external_start=head_start if axis == "head" else 0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段三：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_mappings=axis_mappings, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s field %s: head_start=%d, heads=%d, replicas=%d, write_enabled=%s, bytes=[%d, %d)", + self.state_name, + field_name, + head_start, + local_head_count, + replica_count, + write_enabled, + field_region.byte_start, + field_region.byte_end, + ) + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", + self.state_name, + len(regions), + byte_offset, + ) + return tuple(regions) + + def _storage_shape(self) -> tuple[int, ...]: + return ( + len(self.config.logical_layer_groups), + self.capacity + self.page_size, + self.config.head_num, + self.config.head_dim, + ) diff --git a/lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py b/lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py new file mode 100644 index 0000000000..95db4d79d4 --- /dev/null +++ b/lightllm/common/attention_state/full_att/gqa/static_fp8_transfer.py @@ -0,0 +1,256 @@ +"""静态 FP8 的统一 CPU/PD 页传输：按 token 和页索引直接交换 K/V 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_static_fp8_pages( + k_storage: torch.Tensor, + v_storage: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + layout: StateTransferLayout, + write: bool, + grid_num: int = 24, + page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, +) -> None: + """直接搬运 K/V；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 + """ + assert grid_num > 0, "grid_num must be positive" + assert k_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 K/V。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 K/V 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + # 静态 KV/Q 标定参数属于模型配置，不随 token 页面传输。 + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "k": + local = k_storage + axes = ("layer", "token", "head", "dim") + elif region.field == "v": + local = v_storage + axes = ("layer", "token", "head", "dim") + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], + external.shape[3], + external.shape[4], + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个传输字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 FP8 K/V 的逐字段 launch 融合， + # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 + for local, external in copy_views: + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_static_fp8_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_static_fp8_pages( + # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 字段末轴参与拷贝的元素数，如 dim、group 或 scale。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + storage_layer_index = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 + local_offset = ( + storage_layer_index * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + storage_layer_index * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/attention_state/full_att/interface.py b/lightllm/common/attention_state/full_att/interface.py new file mode 100644 index 0000000000..7240bf9f79 --- /dev/null +++ b/lightllm/common/attention_state/full_att/interface.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from functools import lru_cache +from typing import Any, List + +import torch + +from ..utils import deepcopy_result +from lightllm.common.attention_state.types import StateTransferLayout, TransferRegion + + +class CoreInterface(ABC): + """Full-attention 的资源定义、层视图和 KV 存储核心接口。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + # 具体存储和量化方案标识，用于 main/draft State 成员命名。 + @abstractmethod + def get_state_type(self) -> str: + """返回算法与存储方案标记，用于生成资源成员名。 + + 返回小写 snake_case 的“算法＋存储方案”，不带 full_att/bounded_att 类别前缀， + 也不带 main/draftN 角色或最终的 _state；类别由 Python State 类型表达。 + 管理器统一生成 state_name = f"{model_role}_{get_state_type()}_state"，与成员名完全一致。 + 同角色下不同算法、精度或量化方案必须使用不同标记，同方案的全部逻辑层放在同一实例。 + 层号、容量、设备、TP rank 和请求信息不进入名称；完整布局仍由 config 描述。 + 注册后返回值保持稳定，实际分配与 CPU/DP/PD 传输均使用相同标记。 + + - GQA 浮点 KV：gqa_{bf16|fp16|fp32}。 + - GQA 分组量化：gqa_int{4|8}_group{分组大小}_{scale精度}_scale。 + - GQA 静态 FP8：gqa_fp8_static_{per_tensor|per_head}。 + - MLA 浮点：mla_{精度}；KV/indexer 联合存储：mla_packed_{精度}。 + - MLA packed FP8：mla_packed_fp8_group128，或 mla_packed_fp8_group64_e8m0。 + - 序列压缩历史：compressed_c{压缩比例}_fp8_group64_e8m0。 + 例如 gqa_int8_group128_bf16_scale 对应 main_gqa_int8_group128_bf16_scale_state。 + 新量化实现必须区分影响方案选择的位宽、分组、scale 方式及精度，不编码标定数值。 + """ + ... + + @abstractmethod + def allocate_storage(self, token_capacity: int) -> None: + """按 token 容量和 logical_layer_groups 的物理层数创建存储，沿用已设置的页大小、设备及配套参数。""" + raise NotImplementedError() + + @abstractmethod + def get_token_bytes(self) -> int: + """累计当前 TP rank 各物理层的单 token 字节数，包含动态量化参数和行 padding；共享逻辑层不重复计数。""" + raise NotImplementedError() + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + @abstractmethod + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，不依赖存储分配。""" + ... + + @abstractmethod + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + ... + + @abstractmethod + def get_layer_view(self, logical_layer_id: int) -> Any: + """按逻辑层号返回存储视图；同组逻辑层返回同一物理层的数据，具体字段由该布局定义。""" + raise NotImplementedError() + + @abstractmethod + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: Any) -> None: + """向已分配位置写入一层 KV，并完成该布局要求的编码、量化和关联参数更新。 + + 融合压缩等专用写入路径可直接使用计算视图，不调用本接口。 + """ + raise NotImplementedError() + + @abstractmethod + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制全部所辖层的完整 token 数据，包含动态 scale、indexer 字段，保持原存储表示。 + + 源和目标索引为存储设备上的等长一维 CUDA 整数 Tensor，位置不重叠，直接使用原 dtype 和 stride。""" + raise NotImplementedError() + + +class CpuCacheInterface(ABC): + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 cpu_cache_token_page_size 返回 CPU 字段、全局形状和 TP 切片。 + + byte_start / byte_end 相对本 State 起点，字段间的必要 padding 由区间间隔体现； + 管理器按字段最大终点计算该 State 的 CPU 页占用。 + 本次有效 token 数由读写流程处理，不改变布局。 + 具体实现使用 lru_cache 缓存，外层 deepcopy_result 每次返回独立副本。 + 首次查询后配置与全局 CPU 页规格保持不变。 + 管理器实际分配存储后统一清除缓存，以兼容物理页规格变化；子类覆盖时须保留缓存能力。 + """ + ... + + @abstractmethod + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局写出 token 数据及配套字段；cpu_page_ready_flags 中为 True 的页跳过写入。""" + raise NotImplementedError() + + @abstractmethod + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """从指定 CPU 页读回 token 数据，保留目标中由其他资源或 TP 分片持有的区域。""" + raise NotImplementedError() + + +class DpCopyInterface(ABC): + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + @abstractmethod + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: CoreInterface, + ) -> None: + """复制其他 DP 同一 TP 分片的完整 token 数据，包含量化参数及 indexer 字段。 + + 源、目标索引须已位于各自存储的 CUDA 设备，可为 int32/int64；接口不转换设备或 dtype。 + 调用方保证 peer/IPC 可访问、源数据就绪、目标索引唯一，且源目标数据区域不重叠。 + """ + raise NotImplementedError() + + +class PdInterface(ABC): + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @abstractmethod + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 pd_kv_page_size 返回 PD 字段、全局形状和 TP 切片，区域偏移相对本 State 起点。 + + 内层 lru_cache 缓存区域定义，外层 deepcopy_result 每次返回独立副本。 + 首次查询后配置与全局 PD 页规格保持不变；管理器分配存储后统一清除缓存。 + 子类覆盖时须保留缓存和副本语义。 + """ + ... + + @abstractmethod + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """将 token 数据写入 PD 页集合中的指定发送页，网络发送与同步由调用方负责。""" + raise NotImplementedError() + + @abstractmethod + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """从 PD 页集合中的指定接收页还原 token 数据，不分配索引或修改请求映射。""" + raise NotImplementedError() + + +__all__ = [ + "CoreInterface", + "CpuCacheInterface", + "DpCopyInterface", + "PdInterface", +] diff --git a/lightllm/common/attention_state/full_att/mla/__init__.py b/lightllm/common/attention_state/full_att/mla/__init__.py new file mode 100644 index 0000000000..557ad284a2 --- /dev/null +++ b/lightllm/common/attention_state/full_att/mla/__init__.py @@ -0,0 +1,4 @@ +from .create_mla_state import create_mla_state +from .mla_state import MLAConfig, MLAFullAttentionState + +__all__ = ["create_mla_state", "MLAConfig", "MLAFullAttentionState"] diff --git a/lightllm/common/attention_state/full_att/mla/create_mla_state.py b/lightllm/common/attention_state/full_att/mla/create_mla_state.py new file mode 100644 index 0000000000..f93c9775af --- /dev/null +++ b/lightllm/common/attention_state/full_att/mla/create_mla_state.py @@ -0,0 +1,18 @@ +"""Independent creation entry points for normal and packed MLA States.""" + +from __future__ import annotations + +import torch + +from ..base_state import FullAttentionState +from .mla_state import MLAConfig, MLAFullAttentionState + + +def create_mla_state( + config: MLAConfig, + page_size: int = 1, + device: str | torch.device = "cuda", +) -> FullAttentionState: + """Create a normal latent MLA State without indexer fields.""" + assert type(config) is MLAConfig, "normal MLA creation requires MLAConfig" + return MLAFullAttentionState(layout_config=config, page_size=page_size, device=device) diff --git a/lightllm/common/attention_state/full_att/mla/mla_state.py b/lightllm/common/attention_state/full_att/mla/mla_state.py new file mode 100644 index 0000000000..a818f5af3a --- /dev/null +++ b/lightllm/common/attention_state/full_att/mla/mla_state.py @@ -0,0 +1,296 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, List +from math import prod +from dataclasses import dataclass + +import torch + +from lightllm.utils.envs_utils import get_env_start_args + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState +from ..packed_mla.packed_transfer import transfer_packed_pages +from ..dp_transfer import copy_packed_tokens + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import StateTransferLayout + + +@dataclass(frozen=True) +class MLAConfig: + """普通 latent MLA 的逻辑配置，不包含 indexer 或 packed 物理字段。""" + + # 外层顺序对应物理存储层，组内逻辑层共享该方案的全部字段；各组非空，逻辑层不能重复。 + # 例如 ((0, 4), (8,)) 分配两层，逻辑层 0/4 共用第一层的全部字段。 + # 分组不决定执行顺序或写入者；模型须让共享消费者跳过写入，避免覆盖生产层数据。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # 编码前的完整 latent 宽度，包含模型需要的 RoPE 部分，不含 indexer 或 padding。 + latent_dim: int + # 未压缩 latent 在运行态和存储中的浮点精度。 + latent_dtype: torch.dtype + # 所属模型 TP 组内的 rank；各 rank 保存相同 latent，仅 rank 0 写出共享外部页。 + tp_rank: int = 0 + # 保存完整 latent 副本的 TP rank 数，不改变 latent 的维度或物理行宽。 + tp_world_size: int = 1 + + def __post_init__(self) -> None: + assert isinstance(self.logical_layer_groups, tuple) and all( + isinstance(group, tuple) and group for group in self.logical_layer_groups + ), "layer groups must be a tuple of nonempty tuples" + layers = [layer for group in self.logical_layer_groups for layer in group] + assert all( + isinstance(layer, int) and layer >= 0 for layer in layers + ), "logical layer IDs must be nonnegative integers" + assert len(layers) == len(set(layers)), "a logical layer must belong to exactly one storage group" + assert self.latent_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + + +class MLAFullAttentionState(FullAttentionState): + """存储复制到各 TP rank 的 MLA 压缩 latent，使用单个 KV 字段。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: MLAConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """按 MLA 布局配置准备压缩 KV 存储。""" + self.config = layout_config + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.storage: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回 MLA 压缩 KV 精度对应的存储方案标记。""" + return "mla_" + dtype_name(dtype=self.config.latent_dtype) + + def allocate_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" + self.capacity = token_capacity + self.storage = torch.empty( + size=( + len(self.config.logical_layer_groups), + token_capacity + self.page_size, + 1, + self.config.latent_dim, + ), + dtype=self.config.latent_dtype, + device=self.device, + ) + + def get_token_bytes(self) -> int: + """由配置计算当前 TP rank 的单 token latent 字节数，无需分配存储。""" + config = self.config + return len(config.logical_layer_groups) * config.latent_dim * config.latent_dtype.itemsize + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int) -> torch.Tensor: + """按逻辑层号返回压缩 latent 视图，不拆分为 GQA 的 K/V head。""" + return self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] + + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """通过 CUDA 算子将形状匹配的 MLA latent 数据写入指定物理 token 位置。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) + if not len(token_indices): + return + assert self.storage.is_cuda, "KV writes require CUDA storage" + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + destindex_copy_kv(K=kv, DestLoc=token_indices, Out=self.storage[storage_layer_index]) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" + assert self.storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.storage, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: MLAFullAttentionState, + ) -> None: + """复制其他 DP 同一 TP 分片的完整 token 数据，层顺序保持一致。""" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + copy_packed_tokens( + src_storage=remote_dp_full_att_state.storage, + dst_storage=self.storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """各 rank 复制同一 latent，仅 rank 0 写出共享外部页，所有 rank 均可读回。""" + config = self.config + shape = (len(config.logical_layer_groups), page_tokens, 1, config.latent_dim) + return ( + TransferRegion( + field="kv", + byte_start=0, + byte_end=prod(shape) * config.latent_dtype.itemsize, + dtype=config.latent_dtype, + axes=("layer", "token", "head", "dim"), + shape=shape, + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=0, length=1), + AxisMapping(axis="dim", local_start=0, external_start=0, length=config.latent_dim), + ), + write_enabled=config.tp_rank == 0, + ), + ) diff --git a/lightllm/common/attention_state/full_att/packed_mla/__init__.py b/lightllm/common/attention_state/full_att/packed_mla/__init__.py new file mode 100644 index 0000000000..8447751eae --- /dev/null +++ b/lightllm/common/attention_state/full_att/packed_mla/__init__.py @@ -0,0 +1,12 @@ +from .create_packed_mla_state import create_packed_mla_state +from .packed_mla_state import PackedMLAAttentionConfig, PackedMLALogicalConfig, PackedMLAFullAttentionState +from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState + +__all__ = [ + "create_packed_mla_state", + "PackedMLAAttentionConfig", + "PackedMLALogicalConfig", + "PackedMLAFullAttentionState", + "FP8E8M0PackedMLAConfig", + "FP8E8M0PackedMLAState", +] diff --git a/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py new file mode 100644 index 0000000000..f41646f119 --- /dev/null +++ b/lightllm/common/attention_state/full_att/packed_mla/create_packed_mla_state.py @@ -0,0 +1,36 @@ +"""Creation entry point for packed MLA States.""" + +from __future__ import annotations + +import torch + +from ..base_state import FullAttentionState +from .packed_mla_state import PackedMLAAttentionConfig, PackedMLALogicalConfig, PackedMLAFullAttentionState +from .fp8_e8m0_packed_mla_state import FP8E8M0PackedMLAConfig, FP8E8M0PackedMLAState + + +def create_packed_mla_state( + config: PackedMLALogicalConfig, + page_size: int = 1, + device: str | torch.device = "cuda", +) -> FullAttentionState: + """Create a packed MLA State; packed formats always include an indexer field.""" + assert type(config) is PackedMLALogicalConfig, "packed MLA creation requires PackedMLALogicalConfig" + if config.format == "packed_group64_e8m0": + assert config.latent_dim == FP8E8M0PackedMLAConfig.LATENT_DIM, "FP8/E8M0 MLA requires a 512-dimensional latent" + return FP8E8M0PackedMLAState( + layout_config=FP8E8M0PackedMLAConfig( + logical_layer_groups=config.logical_layer_groups, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + indexer_dim=config.indexer_dim, + ), + page_size=page_size, + device=device, + ) + + return PackedMLAFullAttentionState( + layout_config=PackedMLAAttentionConfig(logical_config=config), + page_size=page_size, + device=device, + ) diff --git a/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py new file mode 100644 index 0000000000..4f018ef215 --- /dev/null +++ b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_state.py @@ -0,0 +1,301 @@ +"""FP8/E8M0 packed MLA：448 维 FP8、64 维 BF16，并在同一行保存 indexer。""" + +from __future__ import annotations + +from functools import lru_cache + +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from lightllm.utils.envs_utils import get_env_start_args + +from ...utils import deepcopy_result +from ...types import AxisMapping, StateTransferLayout, TransferRegion +from ..base_state import FullAttentionState +from ..dp_transfer import copy_packed_tokens +from ..packed_mla.packed_transfer import transfer_packed_pages +from ..packed_mla.packed_mla_state import PackedMLALayerView + + +@dataclass(frozen=True) +class FP8E8M0PackedMLAConfig: + """固定 group64/E8M0 编码，只记录调用方能够选择的原子布局信息。""" + + # 外层为物理存储层，内层为共享该层全部 KV/indexer 的逻辑层号。 + logical_layer_groups: tuple[tuple[int, ...], ...] + # 每个 token 的 indexer key 维度；其后另保存一个 FP32 scale，不包含在该维度中。 + indexer_dim: int = 128 + # 完整 latent 在各 TP rank 复制；rank 0 写出共享外部页，各 rank 均可读回。 + tp_rank: int = 0 + # 所属 TP 组大小，不增加物理 KV head，也不改变单 token 的存储字节数。 + tp_world_size: int = 1 + + # 算子输入的完整 latent 维度，不是量化后的字节数，也不包含 indexer。 + LATENT_DIM: ClassVar[int] = 512 + # 固定编码：448 个 FP8 元素 + 64 个 BF16 元素 + 7 个 E8M0 scale + 1 字节对齐。 + # 末尾 64 维保留 BF16 不代表启用了 RoPE，GLM 使用 NoPE attention。 + KV_BYTES: ClassVar[int] = 584 + + @property + def payload_bytes(self) -> int: + """一行有效数据的字节数：KV 编码、indexer key 及其 FP32 scale，不包含行尾 padding。""" + return self.KV_BYTES + self.indexer_dim + torch.float32.itemsize + + @property + def token_bytes(self) -> int: + """每个 token 的物理行字节数，包含末尾 16 字节对齐 padding。""" + return (self.payload_bytes + 15) // 16 * 16 + + def __post_init__(self) -> None: + assert self.logical_layer_groups and all(self.logical_layer_groups) + layers = [layer for group in self.logical_layer_groups for layer in group] + assert len(layers) == len(set(layers)), "logical layers cannot repeat across storage groups" + assert self.indexer_dim > 0 and 0 <= self.tp_rank < self.tp_world_size + + +class FP8E8M0PackedMLAState(FullAttentionState): + """独立的 group64/E8M0 MLA 实现，不继承其他具体 State 或动态选择编码。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: FP8E8M0PackedMLAConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """只准备布局和逻辑层映射，实际存储由管理器确定 token 容量后统一分配。""" + self.config = layout_config + self.state_name: str | None = None + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.storage: torch.Tensor | None = None + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in group + } + + def get_state_type(self) -> str: + """返回固定算法、量化分组与 scale 编码标记，用于 main/draft 的规范成员名。""" + return "mla_packed_fp8_group64_e8m0" + + def allocate_storage(self, token_capacity: int) -> None: + """分配普通 token 及 HOLD 页；完整行包含 KV、indexer 和 16 字节对齐 padding。""" + self.capacity = token_capacity + self.storage = torch.empty( + size=(len(self.config.logical_layer_groups), token_capacity + self.page_size, 1, self.config.token_bytes), + dtype=torch.uint8, + device=self.device, + ) + + def get_token_bytes(self) -> int: + """返回当前 TP rank 的单 token 开销，包含全部物理层和行尾 padding。""" + return len(self.config.logical_layer_groups) * self.config.token_bytes + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断该逻辑层是否使用当前 State，共享逻辑层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """将逻辑层映射为物理存储行；同组逻辑层返回相同下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: + """返回固定 584 字节 KV 和 indexer key/scale 视图，保留物理行 stride。""" + row = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] + return PackedMLALayerView( + kv=row[..., : self.config.KV_BYTES], + indexer_k=row[..., self.config.KV_BYTES : self.config.payload_bytes], + ) + + def write_full_att_kv( + self, + logical_layer_id: int, + token_indices: torch.Tensor, + kv: torch.Tensor, + ) -> None: + """使用固定 group64/E8M0 算子量化并写入 KV，保留 indexer 和行尾 padding。""" + from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import ( + write_fp8_e8m0_packed_kv, + ) + + assert token_indices.is_cuda, "token indices must reside on CUDA" + if token_indices.numel() == 0: + return + write_fp8_e8m0_packed_kv( + kv=kv, + token_indices=token_indices, + packed_kv=self.get_layer_view(logical_layer_id=logical_layer_id).kv, + ) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """按 CUDA 索引复制完整物理行，包含 indexer、scale 和 padding；两组索引不得重叠。""" + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + if src_indices.numel() == 0: + return + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.storage, + ) + + # ======================================================================== + # 三、CPU Cache：按 token 搬运整份有效编码 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 CPU 页中的完整有效 token 编码，包含 indexer，不包含 GPU 行尾 padding。""" + return self._get_transfer_regions(page_tokens=get_env_start_args().cpu_cache_token_page_size) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """直接写出 KV、E8M0 scale 和 indexer；已就绪的 CPU 页面保持不变。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_pages( + kv_storage=self.storage[..., : self.config.payload_bytes], + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """从 CPU 页恢复完整有效 token 编码，支持末页只恢复部分 token。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_pages( + kv_storage=self.storage[..., : self.config.payload_bytes], + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移：复制相同编码的完整物理行 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: FP8E8M0PackedMLAState, + ) -> None: + """从其他 DP 的同一 TP 分片复制完整 token 数据，由调用方提供可访问的 CUDA 索引。""" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + copy_packed_tokens( + src_storage=remote_dp_full_att_state.storage, + dst_storage=self.storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离：传输固定的完整 token 编码 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回固定 FP8/E8M0 格式的 PD 页描述，由管理器安排联合页偏移。""" + return self._get_transfer_regions(page_tokens=get_env_start_args().pd_kv_page_size) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """将完整有效 token 编码写入指定 PD 页，网络发送与同步由调用方负责。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_pages( + kv_storage=self.storage[..., : self.config.payload_bytes], + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """从指定 PD 页恢复完整有效 token 编码，保留其他 token 与行尾 padding。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_pages( + kv_storage=self.storage[..., : self.config.payload_bytes], + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + write=False, + ) + + # ======================================================================== + # 六、私有辅助：固定 packed 编码作为一个区域传输 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + # 有效字节在每行内连续，KV、scale 和 indexer 无需分别发射 kernel。 + # GPU 行有 padding，通过实际 stride 跳过；外部页只保存 payload_bytes。 + return ( + TransferRegion( + field="kv", + byte_start=0, + byte_end=len(self.config.logical_layer_groups) * page_tokens * self.config.payload_bytes, + dtype=torch.uint8, + axes=("layer", "token", "head", "dim"), + shape=(len(self.config.logical_layer_groups), page_tokens, 1, self.config.payload_bytes), + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=0, + length=1, + ), + AxisMapping( + axis="dim", + local_start=0, + external_start=0, + length=self.config.payload_bytes, + ), + ), + write_enabled=self.config.tp_rank == 0, + ), + ) diff --git a/lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_write.py similarity index 64% rename from lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py rename to lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_write.py index 7a28bc6614..495bbbf7cb 100644 --- a/lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py +++ b/lightllm/common/attention_state/full_att/packed_mla/fp8_e8m0_packed_mla_write.py @@ -1,22 +1,45 @@ +"""group64/E8M0 packed MLA 的定制 KV 量化写入，不依赖具体模型。""" + import torch import triton import triton.language as tl +@torch.no_grad() +def write_fp8_e8m0_packed_kv( + kv: torch.Tensor, + token_indices: torch.Tensor, + packed_kv: torch.Tensor, +) -> None: + assert kv.shape[1:] == (1, 512) + assert packed_kv.dtype == torch.uint8 and packed_kv.shape[1] == 1 + assert packed_kv.shape[2] >= 584 and packed_kv.stride(2) == 1 + _write_fp8_e8m0_packed_kv[(token_indices.numel(),)]( + KV=kv, + TOKEN_INDICES=token_indices.contiguous(), + PACKED_KV=packed_kv, + kv_stride_s=kv.stride(0), + kv_stride_d=kv.stride(2), + packed_stride_s=packed_kv.stride(0), + num_warps=4, + num_stages=1, + ) + + @triton.jit -def _fwd_kernel_destindex_copy_kv_flashmla_fp8( +def _write_fp8_e8m0_packed_kv( KV, - MemIndex, - PackedKV, + TOKEN_INDICES, + PACKED_KV, kv_stride_s, kv_stride_d, packed_stride_s, ): token = tl.program_id(0) - mem_index = tl.load(MemIndex + token).to(tl.int64) - packed_row = PackedKV + mem_index * packed_stride_s + mem_index = tl.load(TOKEN_INDICES + token).to(tl.int64) + packed_row = PACKED_KV + mem_index * packed_stride_s - # Eight rows share one CTA; seven quantized groups are processed in parallel. + # 每个 token 按 group64 并行量化七组特征；末尾 64 维单独保留 BF16。 groups = tl.arange(0, 8) offsets = groups[:, None] * 64 + tl.arange(0, 64)[None, :] group_mask = groups < 7 @@ -33,20 +56,3 @@ def _fwd_kernel_destindex_copy_kv_flashmla_fp8( tail_values = tl.load(KV + token * kv_stride_s + (448 + tail_offsets) * kv_stride_d) tl.store((packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + tail_offsets, tail_values) tl.store(packed_row + 583, 0) - - -@torch.no_grad() -def destindex_copy_kv_flashmla_fp8(kv: torch.Tensor, mem_index: torch.Tensor, packed_kv: torch.Tensor): - assert kv.shape[1:] == (1, 512) - assert packed_kv.dtype == torch.uint8 and packed_kv.shape[1] == 1 - assert packed_kv.shape[2] >= 584 and packed_kv.stride(2) == 1 - _fwd_kernel_destindex_copy_kv_flashmla_fp8[(mem_index.numel(),)]( - kv, - mem_index.contiguous(), - packed_kv, - kv.stride(0), - kv.stride(2), - packed_kv.stride(0), - num_warps=4, - num_stages=1, - ) diff --git a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py new file mode 100644 index 0000000000..99c920fe26 --- /dev/null +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_state.py @@ -0,0 +1,495 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, List, Literal +from dataclasses import dataclass +from math import prod + +import torch + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from ...utils import deepcopy_result, dtype_name +from ...types import AxisMapping, TransferRegion +from ..base_state import FullAttentionState +from .packed_mla_transfer import transfer_packed_mla_pages +from ..dp_transfer import copy_packed_tokens + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import StateTransferLayout + + +logger = init_logger(name=__name__) + + +@dataclass(frozen=True) +class PackedMLALogicalConfig: + """带 indexer 的 packed MLA 逻辑配置；不描述实际字节布局。""" + + logical_layer_groups: tuple[tuple[int, ...], ...] + # 编码前的完整 latent 宽度；具体分块、scale 和 padding 由工厂转换。 + latent_dim: int + # packed MLA 的逻辑格式；所有选项都包含 indexer。 + format: Literal["packed_fp16", "packed_bf16", "packed_group128_fp8", "packed_group64_e8m0"] + # indexer key 的逻辑宽度。 + indexer_dim: int = 128 + # latent 中保留原始 RoPE 精度的宽度；NoPE 宽度由 latent_dim - rope_dim 得到。 + rope_dim: int = 64 + tp_rank: int = 0 + tp_world_size: int = 1 + + def __post_init__(self) -> None: + assert self.logical_layer_groups and all(self.logical_layer_groups) + layers = [layer for group in self.logical_layer_groups for layer in group] + assert len(layers) == len(set(layers)), "logical layers cannot repeat across storage groups" + assert ( + self.latent_dim > 0 + and 0 <= self.rope_dim <= self.latent_dim + and self.indexer_dim > 0 + and 0 <= self.tp_rank < self.tp_world_size + ) + assert self.format in ( + "packed_fp16", + "packed_bf16", + "packed_group128_fp8", + "packed_group64_e8m0", + ), "unknown packed MLA logical format" + + +@dataclass(frozen=True) +class PackedMLAAttentionConfig: + """由 packed 逻辑配置解析出的物理布局，不重复保存逻辑字段。""" + + logical_config: PackedMLALogicalConfig + quant_group_size: int = 128 + + @property + def logical_layer_groups(self): + return self.logical_config.logical_layer_groups + + @property + def latent_dim(self): + return self.logical_config.latent_dim + + @property + def kv_encoding(self): + return { + "packed_fp16": "fp16", + "packed_bf16": "bf16", + "packed_group128_fp8": "fp8_group128", + }[self.logical_config.format] + + @property + def indexer_dim(self): + return self.logical_config.indexer_dim + + @property + def tp_rank(self): + return self.logical_config.tp_rank + + @property + def tp_world_size(self): + return self.logical_config.tp_world_size + + @property + def rope_dim(self): + return self.logical_config.rope_dim + + @property + def dtype(self) -> torch.dtype: + """返回编码对应的逻辑存储 dtype；FP8 的物理承载仍由 State 内部使用 uint8。""" + if self.kv_encoding == "fp16": + return torch.float16 + if self.kv_encoding == "bf16": + return torch.bfloat16 + return torch.float8_e4m3fn + + @property + def kv_bytes(self) -> int: + """按逻辑 NoPE/RoPE 宽度和 scale 数量计算单行 KV 字节数。""" + if self.kv_encoding == "fp8_group128": + nope_dim = self.latent_dim - self.rope_dim + scale_count = nope_dim // self.quant_group_size + return nope_dim + scale_count * torch.float32.itemsize + self.rope_dim * torch.bfloat16.itemsize + return self.latent_dim * self.dtype.itemsize + + @property + def token_bytes(self) -> int: + """KV、indexer key/scale 及整行 16 字节对齐所需的物理字节数。""" + return (self.kv_bytes + self.indexer_dim + 4 + 15) // 16 * 16 + + @property + def indexer_offset(self) -> int: + """保持既有格式：FP8 紧接 KV，浮点格式将对齐间隙放在 KV 与 indexer 之间。""" + return self.kv_bytes if self.kv_encoding == "fp8_group128" else self.token_bytes - self.indexer_dim - 4 + + def __post_init__(self) -> None: + assert self.kv_encoding in ("fp16", "bf16", "fp8_group128"), "unknown packed MLA KV encoding" + if self.kv_encoding == "fp8_group128": + assert ( + self.dtype == torch.float8_e4m3fn + and self.rope_dim == 64 + and self.latent_dim == 576 + and self.quant_group_size == 128 + ), "group128 packed MLA requires latent 576, rope 64 and group size 128" + else: + assert self.kv_encoding in ("fp16", "bf16") + + +@dataclass(frozen=True) +class PackedMLALayerView: + kv: torch.Tensor + indexer_k: torch.Tensor + + +class PackedMLAFullAttentionState(FullAttentionState): + """在同一物理 token 行中存储 MLA latent、量化参数及稀疏 indexer key/scale。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + layout_config: PackedMLAAttentionConfig, + page_size: int = 1, + device: str | torch.device = "cuda", + ) -> None: + """按 packed MLA 配置准备 KV 与 indexer 存储。""" + self.config = layout_config + self.state_name: str | None = None # 由 manager 绑定模型角色时设置。 + # 外层下标是物理层位置；组内多个逻辑层映射到同一份 KV/量化字段存储。 + self._logical_to_storage_layer = { + logical_layer_id: storage_layer_index + for storage_layer_index, logical_layer_group in enumerate(layout_config.logical_layer_groups) + for logical_layer_id in logical_layer_group + } + self.page_size = page_size + self.device = torch.device(device=device) + self.capacity = 0 + self.storage: torch.Tensor | None = None + + def get_state_type(self) -> str: + """返回 packed MLA 编码方案标记，区分 FP8 与原始精度。""" + encoding = ( + f"fp8_group{self.config.quant_group_size}" + if self.config.kv_encoding == "fp8_group128" + else dtype_name(dtype=self.config.dtype) + ) + return "mla_packed_" + encoding + + def allocate_storage(self, token_capacity: int) -> None: + """按统一 token 容量分配本 State 的独立 KV 存储与 HOLD 页。""" + self.capacity = token_capacity + # FP8 行混合保存多种字段，用字节 buffer 承载，字段视图再解释各自的 dtype。 + storage_dtype = torch.uint8 if self.config.kv_encoding == "fp8_group128" else self.config.dtype + self.storage = torch.empty( + size=( + len(self.config.logical_layer_groups), + token_capacity + self.page_size, + 1, + self.config.token_bytes // storage_dtype.itemsize, + ), + dtype=storage_dtype, + device=self.device, + ) + + def get_token_bytes(self) -> int: + """返回全部物理层的单 token 存储字节数，包含编码、indexer 和 padding。""" + return len(self.config.logical_layer_groups) * self.config.token_bytes + + # ======================================================================== + # 二、Full attention：计算视图、KV 写入与完整 token 复制 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_storage_layer_index(self, logical_layer_id: int) -> int: + """按逻辑层号查询存储及配套字段的物理行；共享逻辑层返回同一下标。""" + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int) -> PackedMLALayerView: + """分别返回 KV 与 packed indexer key/scale 的视图，保持原行步长。""" + row = self.storage[self.get_storage_layer_index(logical_layer_id=logical_layer_id)] + raw = row.view(torch.uint8) + kv = ( + raw[..., : self.config.kv_bytes] + if self.config.kv_encoding == "fp8_group128" + else row[..., : self.config.latent_dim] + ) + indexer = raw[..., self.config.indexer_offset : self.config.indexer_offset + self.config.indexer_dim + 4] + return PackedMLALayerView(kv=kv, indexer_k=indexer) + + def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, kv: torch.Tensor) -> None: + """写入 MLA KV 及其量化参数，保留同一 token 行中的 indexer 区域。""" + assert token_indices.is_cuda, "token indices must reside on CUDA" + storage_layer_index = self.get_storage_layer_index(logical_layer_id=logical_layer_id) + if not len(token_indices): + return + assert self.storage.is_cuda, "KV writes require CUDA storage" + if self.config.kv_encoding == "fp8_group128": + from lightllm.common.basemodel.triton_kernel.kv_copy.destindex_copy_kv_flashmla_fp8 import ( + destindex_copy_kv_flashmla_fp8, + ) + + raw = self.storage[storage_layer_index] + nope_dim = self.config.latent_dim - self.config.rope_dim + scale_bytes = nope_dim // self.config.quant_group_size * torch.float32.itemsize + rope_start = nope_dim + scale_bytes + destindex_copy_kv_flashmla_fp8( + KV_nope=kv[..., :nope_dim], + KV_rope=kv[..., nope_dim:], + DestLoc=token_indices, + O_nope=raw[..., :nope_dim].view(torch.float8_e4m3fn), + O_scale=raw[..., nope_dim:rope_start].view(torch.float32), + O_rope=raw[..., rope_start : rope_start + self.config.rope_dim * torch.bfloat16.itemsize].view( + torch.bfloat16 + ), + ) + else: + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + destindex_copy_kv( + K=kv, DestLoc=token_indices, Out=self.storage[storage_layer_index, ..., : self.config.latent_dim] + ) + + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制完整 token 数据；索引使用 Tensor，源与目标位置由调用方保证不重叠。""" + assert src_indices.is_cuda and dst_indices.is_cuda, "token indices must reside on CUDA" + assert ( + src_indices.ndim == dst_indices.ndim == 1 and src_indices.numel() == dst_indices.numel() + ), "token copies require equal one-dimensional index tensors" + assert src_indices.dtype in (torch.int32, torch.int64) and dst_indices.dtype in ( + torch.int32, + torch.int64, + ), "token indices must be integer tensors" + assert self.storage.is_cuda, "token copies require CUDA storage" + if not src_indices.numel(): + return + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 + # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 + # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 + copy_kv_buffer_to_kv_buffer( + src_mem_index=src_indices, + dst_mem_index=dst_indices, + kv_buffer=self.storage, + ) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 CPU 页大小返回字段区间与 TP 切片，作为 CPU 页容量计算依据。 + + 覆盖 KV、动态量化参数及 indexer 有效字段，不包含 GPU 行 padding 或静态标定参数。 + """ + page_tokens = get_env_start_args().cpu_cache_token_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_cpu( + self, + token_indices: torch.Tensor, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + kv_dim=self.config.latent_dim, + indexer_offset=self.config.indexer_offset, + indexer_dim=self.config.indexer_dim, + write=True, + page_readies=cpu_page_ready_flags, + ) + + def read_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + token_indices: torch.Tensor, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + transfer_packed_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=cpu_page_buffers, + page_indices=cpu_page_indices, + layout=layout, + kv_dim=self.config.latent_dim, + indexer_offset=self.config.indexer_offset, + indexer_dim=self.config.indexer_dim, + write=False, + ) + + # ======================================================================== + # 四、DP 迁移 + # ======================================================================== + + def copy_tokens_from_dp( + self, + src_token_indices: torch.Tensor, + dst_token_indices: torch.Tensor, + remote_dp_full_att_state: PackedMLAFullAttentionState, + ) -> None: + """校验 MLA 编码与区域尺寸，再复制完整 KV 和 indexer 数据。""" + assert ( + remote_dp_full_att_state.config.kv_encoding, + remote_dp_full_att_state.config.latent_dim, + remote_dp_full_att_state.config.indexer_dim, + ) == ( + self.config.kv_encoding, + self.config.latent_dim, + self.config.indexer_dim, + ), "incompatible packed MLA format" + assert ( + type(remote_dp_full_att_state) is type(self) and remote_dp_full_att_state.config == self.config + ), "DP copy requires the same encoding and equivalent TP state configurations" + assert src_token_indices.is_cuda and dst_token_indices.is_cuda, "DP token indices must reside on CUDA" + assert src_token_indices.numel() == dst_token_indices.numel(), "source and destination token counts differ" + copy_packed_tokens( + src_storage=remote_dp_full_att_state.storage, + dst_storage=self.storage, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + ) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """按全局 PD 页大小返回字段与切片，交由管理器安排其在联合页中的位置。""" + page_tokens = get_env_start_args().pd_kv_page_size + return self._get_transfer_regions(page_tokens=page_tokens) + + def write_to_pd_page( + self, + token_indices: torch.Tensor, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接写出有效 token 及配套字段，不创建中间 KV 存储。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + kv_dim=self.config.latent_dim, + indexer_offset=self.config.indexer_offset, + indexer_dim=self.config.indexer_dim, + write=True, + ) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + token_indices: torch.Tensor, + pd_page_index: int, + layout: StateTransferLayout, + ) -> None: + """按布局直接恢复有效 token 及配套字段，保留未参与的存储区域。""" + assert pd_page_buffers.is_cuda, "PD pages must reside on CUDA" + transfer_packed_mla_pages( + kv_storage=self.storage, + token_indices=token_indices, + page_buffers=pd_page_buffers, + page_indices=[pd_page_index], + layout=layout, + kv_dim=self.config.latent_dim, + indexer_offset=self.config.indexer_offset, + indexer_dim=self.config.indexer_dim, + write=False, + ) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self, page_tokens: int) -> tuple[TransferRegion, ...]: + """依次排列 KV、动态量化等字段，描述全局页形状和本地 TP 分片。""" + config = self.config + byte_offset = 0 + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, tokens_per_page=%d", + self.state_name, + config.tp_rank, + config.tp_world_size, + page_tokens, + ) + # 阶段一：按本方案固定字段计算形状和 dtype，不分配 meta 存储。 + if config.kv_encoding == "fp8_group128": + field_specs = [ + ("kv_nope", torch.uint8, 512, "dim"), + ("kv_scale", torch.float32, 4, "group"), + ("kv_rope", torch.bfloat16, 64, "dim"), + ] + else: + field_specs = [("kv", config.dtype, config.latent_dim, "dim")] + field_specs.extend( + ( + ("indexer_k", torch.uint8, config.indexer_dim, "dim"), + ("indexer_scale", torch.float32, 1, "scale"), + ) + ) + for field_name, field_dtype, field_width, last_axis in field_specs: + axes = ("layer", "token", "head", last_axis) + local_shape = (len(config.logical_layer_groups), page_tokens, 1, field_width) + global_shape = local_shape + # 全部 TP ranks 保存同一 latent，只有 rank 0 写出共享外部页。 + write_enabled = config.tp_rank == 0 + axis_mappings = tuple( + AxisMapping( + axis=axis, + local_start=0, + external_start=0, + length=dimension, + ) + for axis, dimension in zip(axes, local_shape) + if axis not in ("layer", "token") + ) + + # 阶段二：按字段 dtype 对齐起点，记录本字段区间，再推进下一个字段的字节偏移。 + element_bytes = field_dtype.itemsize + byte_offset = (byte_offset + element_bytes - 1) // element_bytes * element_bytes + field_region = TransferRegion( + field=field_name, + byte_start=byte_offset, + byte_end=byte_offset + prod(global_shape) * element_bytes, + dtype=field_dtype, + axes=axes, + shape=tuple(global_shape), + axis_mappings=axis_mappings, + write_enabled=write_enabled, + ) + regions.append(field_region) + byte_offset = field_region.byte_end + logger.debug( + "%s transfer regions built: fields=%d, total_bytes=%d", + self.state_name, + len(regions), + byte_offset, + ) + return tuple(regions) diff --git a/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py new file mode 100644 index 0000000000..86a0dc4145 --- /dev/null +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_mla_transfer.py @@ -0,0 +1,273 @@ +"""packed MLA 的统一 CPU/PD 页传输：按 token 和页索引直接交换 latent 及配套字段 与外部页。""" + +import torch +import triton +import triton.language as tl + +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_packed_mla_pages( + kv_storage: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + layout: StateTransferLayout, + kv_dim: int, + indexer_offset: int, + indexer_dim: int, + write: bool, + grid_num: int = 24, + page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, +) -> None: + """直接搬运 latent 及配套字段；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 + """ + assert grid_num > 0, "grid_num must be positive" + assert kv_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 latent 及配套字段。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 latent 及配套字段 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + # 只创建原始字节视图；每行保留真实步长，传输仅覆盖已声明字段，不包含 padding。 + raw = kv_storage.view(torch.uint8) + indexer_end = indexer_offset + indexer_dim + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "kv": + local = kv_storage[..., :kv_dim] + axes = ("layer", "token", "head", "dim") + elif region.field == "kv_nope": + # FP8 编码、FP32 scale 和 BF16 RoPE 分别位于固定字节区间。 + local = raw[..., :512] + axes = ("layer", "token", "head", "dim") + elif region.field == "kv_scale": + local = raw[..., 512:528].view(torch.float32) + axes = ("layer", "token", "head", "group") + elif region.field == "kv_rope": + local = raw[..., 528:656].view(torch.bfloat16) + axes = ("layer", "token", "head", "dim") + elif region.field == "indexer_k": + local = raw[..., indexer_offset:indexer_end] + axes = ("layer", "token", "head", "dim") + elif region.field == "indexer_scale": + local = raw[..., indexer_end : indexer_end + 4].view(torch.float32) + axes = ("layer", "token", "head", "scale") + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], + external.shape[3], + external.shape[4], + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个 传输字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + # TODO：当前 CPU/PD 已直接按索引搬运，可进一步将 KV、indexer 及各量化字段 的逐字段 launch 融合， + # 共用 token/page 索引；保留各字段 dtype、stride、TP 切片、write_enabled 和页内 padding。 + for local, external in copy_views: + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_packed_mla_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_packed_mla_pages( + # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 字段末轴参与拷贝的元素数，如 dim、group 或 scale。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + storage_layer_index = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 + local_offset = ( + storage_layer_index * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + storage_layer_index * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py b/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py new file mode 100644 index 0000000000..03d00746bd --- /dev/null +++ b/lightllm/common/attention_state/full_att/packed_mla/packed_transfer.py @@ -0,0 +1,249 @@ +"""单字段 packed 存储的 CPU/PD 页传输，供 MLA 和压缩历史等方案复用。""" + +import torch +import triton +import triton.language as tl + +from ...types import StateTransferLayout +from ...utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +# ======================================================================== +# 一、CPU/PD 共用的页传输入口 +# ======================================================================== + + +def transfer_packed_pages( + kv_storage: torch.Tensor, + token_indices: torch.Tensor, + page_buffers: torch.Tensor, + page_indices: list[int], + layout: StateTransferLayout, + write: bool, + grid_num: int = 24, + page_readies: list[bool] | None = None, + num_warps: int = 8, + num_stages: int = 1, +) -> None: + """直接搬运 packed 存储；CPU 传入页列表，PD 传入 [pd_page_index]，同步由调用方负责。 + + 外部页为 pinned/注册 CPU 内存或 CUDA PD buffer；跨卡指针由调用方提前建立 peer/IPC 映射。 + 单页和多页统一通过页索引 Tensor 定位，共用布局和拷贝 kernel。 + grid_num 控制每个字段发射的 program 数，默认 24；kernel 跳步循环覆盖全部搬运单元。 + num_warps 控制每个 program 的 warp 数，num_stages 控制循环级流水深度。 + """ + assert grid_num > 0, "grid_num must be positive" + assert kv_storage.is_cuda, "page transfers require CUDA State storage" + assert page_buffers.ndim >= 2, "page buffers must include a leading page dimension" + assert isinstance(page_indices, list), "page indices must be a list" + if page_buffers.device.type == "cpu": + assert page_buffers.is_pinned(), "CPU pages must be pinned or registered" + page_tokens = layout.tokens_per_page + assert page_tokens is not None and page_tokens > 0, "KV transfers require a positive token capacity" + assert token_indices.is_cuda, "token indices must reside on CUDA" + token_count = token_indices.numel() + assert len(page_indices) == triton.cdiv(x=token_count, y=page_tokens), "KV token and page counts differ" + if write and page_readies is not None: + assert len(page_readies) == len(page_indices), "page and ready counts differ" + if all(page_readies): + return + if token_count == 0: + return + + # 阶段一：准备索引，不预选或复制 payload。 + # token_indices[i] 是第 i 个 token 的 GPU 物理槽位，直接使用传入的 dtype 和 stride。 + # 跨 GPU 的 KV 由传输进程通过现有 IPC peer 映射导入当前设备，kernel 沿当前流访问。 + # page_ids[j] 是第 j 个逻辑页对应的外部物理页 ID，允许不连续、不递增。 + # CPU/PD、单页/多页统一通过 pinned 索引异步提交到 GPU，临时显存只与页数有关。 + page_ids = torch.tensor(data=page_indices, dtype=torch.int64, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # ready 标记按逻辑页排列，保留完整 token/page 索引，避免跳过中间页后错位。 + # 本 State 的全部字段共享此 Tensor；PD 和读回路径不需要 ready 标记。 + page_readies_tensor = None + if write and page_readies is not None: + page_readies_tensor = torch.tensor(data=page_readies, dtype=torch.bool, pin_memory=True).to( + device=torch.cuda.current_device(), + non_blocking=True, + ) + + # 阶段二：构造全部字段的视图，完成后再统一提交 kernel。 + # page_bytes 是 [物理页数, 每页可见字节数]，第一维仍保留真实页跨度。 + # 一个联合页可能还保存其他 State 的字段，不能用当前 latent 大小代替页跨度。 + page_bytes = page_buffers.view(torch.uint8).view(page_buffers.shape[0], -1) + copy_views = [] + for region in layout.regions: + if write and not region.write_enabled: + continue + if region.field == "kv": + local = kv_storage + axes = ("layer", "token", "head", "dim") + else: + raise AssertionError(f"unknown transfer field {region.field}") + # 从各页中取字段的字节区间，按 dtype 和布局形状解释，不搬运数据。 + external = page_bytes[:, region.byte_start : region.byte_end].view(region.dtype) + external = external.view(page_buffers.shape[0], *region.shape) + # 外部统一为 [page, layer, page_token, head, field_width]；+1 是因为多了 page 轴。 + external = external.permute(0, *(region.axes.index(axis) + 1 for axis in axes)) + local_slices = [slice(None)] * 4 + external_slices = [slice(None)] * 5 + for axis_mapping in region.axis_mappings: + # local_start/external_start 固定表示本地/外部起点，不随读写方向改变。 + # 例：本地 heads [0, 2) 对应外部 heads [4, 6)，切片保留两侧真实 stride。 + axis = axes.index(axis_mapping.axis) + local_slices[axis] = slice(axis_mapping.local_start, axis_mapping.local_start + axis_mapping.length) + external_slices[axis + 1] = slice( + axis_mapping.external_start, axis_mapping.external_start + axis_mapping.length + ) + local = local[tuple(local_slices)] + external = external[tuple(external_slices)] + # token 轴不比较：local 是 GPU 槽位容量，external 是单页容量。 + assert local.dtype == external.dtype + assert (local.shape[0], local.shape[2], local.shape[3]) == ( + external.shape[1], + external.shape[3], + external.shape[4], + ), "local and external copy views must have the same shape" + copy_views.append((local, external)) + + # 阶段三：每个 传输字段一次 launch，覆盖全部层、有效 token 和外部页。 + # elements 是转换后搬运单元数，不是字节数；每个字段固定发射 grid_num 个 program。 + # 超出一轮 grid 覆盖范围的数据由 kernel 跳步处理，没有 gather/scatter 中间存储。 + # stream/event 同步由调用方负责，函数返回不表示 CPU 页已经可读取或复用。 + for local, external in copy_views: + # 先合并连续的 head/width，保留 layer/token/page；不能合并时沿用原视图。 + # head 轴保留为 1，继续使用原来的四维/五维 kernel 接口，再选择最大搬运 dtype。 + local, external = merge_contiguous_tail_dims( + copies=[(local, external)], + start_dim=-2, + keepdim=True, + )[0] + local, external = view_max_copy_dtype(copies=[(local, external)])[0] + elements = local.shape[0] * token_count * local.shape[2] * local.shape[3] + _copy_packed_pages[(grid_num,)]( + LOCAL=local, + EXTERNAL=external, + TOKEN_IDS=token_indices, + TOKEN_STRIDE=token_indices.stride(0), + PAGE_IDS=page_ids, + PAGE_READIES=page_readies_tensor, + HAS_PAGE_READIES=page_readies_tensor is not None, + token_count=token_count, + element_count=elements, + LOCAL_STRIDES=local.stride(), + EXTERNAL_STRIDES=external.stride(), + PAGE_TOKENS=page_tokens, + HEADS=local.shape[2], + WIDTH=local.shape[3], + ELEMENT_BYTES=local.element_size(), + WRITE=write, + BLOCK=1024, + NUM_STAGES=num_stages, + num_warps=num_warps, + ) + + +# ======================================================================== +# 二、CUDA kernel：按索引和实际 stride 双向搬运原始编码 +# ======================================================================== + + +@triton.jit +def _copy_packed_pages( + # 数据指针：一次处理一个字段；调用方已应用 TP 切片，两侧 dtype 相同。 + LOCAL, # GPU 视图：[layer, token_slot, head, dim]。 + EXTERNAL, # CPU/PD 外部页视图：[page, layer, page_token, head, field_width]。 + # 索引指针：均位于 GPU；索引值是物理槽位/页 ID，允许不连续、不递增。 + TOKEN_IDS, # TOKEN_IDS[token] 指向 GPU KV 存储中的 token 槽位。 + PAGE_IDS, # GPU 页索引 Tensor，PAGE_IDS[token // PAGE_TOKENS] 给出外部物理页 ID。 + PAGE_READIES, # 可选的 GPU bool 数组；True 表示该逻辑页已就绪，禁止覆盖。 + # 本次搬运范围：只覆盖有效 token，不包含尾页空余位置。 + token_count, # 本次有效 token 数；调用方保证大于 0。 + element_count, # layer 数 × token_count × HEADS × WIDTH，单位为转换后的搬运单元。 + # 编译期寻址配置：stride 的单位是元素数，不是字节数。 + TOKEN_STRIDE: tl.constexpr, # token 索引步长，直接支持切片视图，不创建连续副本。 + LOCAL_STRIDES: tl.constexpr, # 对应 LOCAL 的四个轴，保留原 GPU 存储步长。 + EXTERNAL_STRIDES: tl.constexpr, # 对应 EXTERNAL 的五个轴，包含真实外部页跨度。 + PAGE_TOKENS: tl.constexpr, # 每个外部页可容纳的 token 数。 + HEADS: tl.constexpr, # 本次 TP 切片参与拷贝的 head 数。 + WIDTH: tl.constexpr, # 字段末轴参与拷贝的元素数，如 dim、group 或 scale。 + # 编译期搬运配置。 + ELEMENT_BYTES: tl.constexpr, # 单个元素字节数，用于选择同宽整数指针。 + HAS_PAGE_READIES: tl.constexpr, # 无 ready 标记时编译掉检查，保持 PD/读回路径。 + WRITE: tl.constexpr, # True：KV→外部页；False：外部页→KV。 + BLOCK: tl.constexpr, # 每个 program 每轮处理的搬运单元数。 + NUM_STAGES: tl.constexpr, # tl.range 的循环流水深度，不是 launch 的同名选项。 +): + # 阶段一：把两侧指针解释为同宽整数指针，逐位保留数据（包括 NaN 位模式）。 + # 这里只改变指针类型，不进行数值转换；元素宽度不变，因此原 stride 仍然有效。 + if ELEMENT_BYTES == 1: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint8)) + elif ELEMENT_BYTES == 2: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint16)) + elif ELEMENT_BYTES == 4: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint32)) + else: + LOCAL = LOCAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + EXTERNAL = EXTERNAL.to(dtype=tl.pointer_type(element_ty=tl.uint64)) + + # 阶段二：确定当前 program 的起点和整张 grid 的步长。 + # 每轮处理 BLOCK 个元素，下一轮前进 grid_num * BLOCK，所有 program 共同覆盖完整范围。 + # 使用 int64 避免大容量寻址溢出；tl.cast 兼容 token_count=1 被特化为常量的情况。 + token_count = tl.cast(input=token_count, dtype=tl.int64) + program_start = tl.program_id(axis=0).to(dtype=tl.int64) * BLOCK + grid_stride = tl.num_programs(axis=0).to(dtype=tl.int64) * BLOCK + block_offsets = tl.arange(start=0, end=BLOCK) + # 起点超出 element_count 的 program 不进入循环；最后一块通过 valid 屏蔽尾部元素。 + for block_start in tl.range(program_start, element_count, grid_stride, num_stages=NUM_STAGES): + offsets = block_start + block_offsets + valid = offsets < element_count + token = offsets // (WIDTH * HEADS) % token_count + + # 阶段三：把逻辑 token 位置映射到两侧物理存储。 + # slot 是 GPU token 槽位；page 是外部物理页；token % PAGE_TOKENS 是页内位置。 + # 例如 PAGE_TOKENS=16 时，逻辑 token 18 使用 TOKEN_IDS[18] 和 PAGE_IDS[1] 的第 2 个位置。 + # 索引读取同样使用 valid，避免尾部无效元素访问 TOKEN_IDS/PAGE_IDS 越界。 + # 写出 ready 页时同时屏蔽源 KV 读取和目标写入，不仅仅跳过 store。 + copy_block = True + if WRITE and HAS_PAGE_READIES: + ready = tl.load(pointer=PAGE_READIES + token // PAGE_TOKENS, mask=valid, other=1) + valid = valid & ~ready + # 一个 block 可能跨页：混合状态保留逐元素 mask；全部无需写出则跳过后续操作。 + copy_block = tl.sum(input=valid.to(dtype=tl.int32), axis=0) > 0 + if copy_block: + dim = offsets % WIDTH + head = offsets // WIDTH % HEADS + storage_layer_index = offsets // (WIDTH * HEADS * token_count) + slot = tl.load(pointer=TOKEN_IDS + token * TOKEN_STRIDE, mask=valid, other=0).to(dtype=tl.int64) + page = tl.load(pointer=PAGE_IDS + token // PAGE_TOKENS, mask=valid, other=0).to(dtype=tl.int64) + + # 阶段四：按各轴真实 stride 计算元素偏移，兼容 TP 切片及轴顺序调整后的视图。 + # EXTERNAL 已指向当前字段在联合页中的起点；这里无需再次加 region.byte_start。 + # EXTERNAL_STRIDES[0] 使用整个外部页的跨度，不能替换为当前 传输字段的大小。 + local_offset = ( + storage_layer_index * LOCAL_STRIDES[0] + + slot * LOCAL_STRIDES[1] + + head * LOCAL_STRIDES[2] + + dim * LOCAL_STRIDES[3] + ) + external_offset = ( + page * EXTERNAL_STRIDES[0] + + storage_layer_index * EXTERNAL_STRIDES[1] + + (token % PAGE_TOKENS) * EXTERNAL_STRIDES[2] + + head * EXTERNAL_STRIDES[3] + + dim * EXTERNAL_STRIDES[4] + ) + + # 阶段五：按方向直接读写两侧存储，不生成 gather/scatter 中间 Tensor。 + # 两个方向使用相同的索引和有效范围；尾页空余、其他 TP 分片及其他字段保持不变。 + if WRITE: + value = tl.load(pointer=LOCAL + local_offset, mask=valid, other=0) + tl.store(pointer=EXTERNAL + external_offset, value=value, mask=valid) + else: + value = tl.load(pointer=EXTERNAL + external_offset, mask=valid, other=0) + tl.store(pointer=LOCAL + local_offset, value=value, mask=valid) diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py new file mode 100644 index 0000000000..6fb18c7ea1 --- /dev/null +++ b/lightllm/common/attention_state/manager.py @@ -0,0 +1,430 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from functools import lru_cache +from typing import List + +import torch + +from lightllm.common.attention_state.types import PageKind + +from lightllm.utils.dist_utils import get_current_device_id +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +from .allocators import StatePageAllocator, FullAttKVAllocator +from .cpu_cache_operator import CpuCacheOperator +from .dp_shared_operator import DPSharedOperator +from .full_att import FullAttentionState +from .bounded_att import BoundedAttentionState +from .pd_operator import PDOperator +from .utils import MODEL_ROLES + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + + +logger = init_logger(__name__) + + +class AttentionStateManager: + """管理模型组合的注意力资源，提供缓存复用及 CPU/DP/PD 传输所需的聚合操作。 + + 通过聚合查询操作各类 State；模型通过两类创建接口决定资源组合，main/draft 共享本管理器。 + 未支持的传输路径必须明确报错。 + + CPU cache、DP shared 和 PD 的实现分别由三个 Operator 持有；管理器只保留 + State 注册、查询和共享容量规划。Operator 只依赖本管理器提供的 State 查询 + 和 allocator，不复制 State 的算法逻辑。 + """ + + # ======================================================================== + # 一、基础管理：共享分配器、资源初始化与聚合查询 + # ======================================================================== + + # 共享分配器：按模型需要启用；调用方直接分配和归还索引。 + full_att_allocator: FullAttKVAllocator | None + bounded_att_big_page_allocator: StatePageAllocator | None + bounded_att_small_page_allocator: StatePageAllocator | None + _state_registration_finished: bool = False + + def __init__(self, request_capacity: int): + """记录共享请求容量；State 和 allocator 由后续初始化流程设置。""" + self.full_att_allocator = None + self.bounded_att_big_page_allocator = None + self.bounded_att_small_page_allocator = None + self._request_capacity = request_capacity + self._cpu_cache_operator = CpuCacheOperator(self) + self._dp_shared_operator = DPSharedOperator(self) + self._pd_operator = PDOperator(self) + + def _reset_operators(self) -> None: + """让浅拷贝后的 Operator 指向新的资源管理器副本。""" + self._cpu_cache_operator = CpuCacheOperator(self) + self._dp_shared_operator = DPSharedOperator(self) + self._pd_operator = PDOperator(self) + + def get_cpu_cache_operator(self) -> CpuCacheOperator: + """返回当前资源管理器唯一的 CPU cache 操作对象。""" + return self._cpu_cache_operator + + def get_dp_shared_operator(self) -> DPSharedOperator: + """返回当前资源管理器唯一的 DP shared 操作对象。""" + return self._dp_shared_operator + + def get_pd_operator(self) -> PDOperator: + """返回当前资源管理器唯一的 PD 操作对象。""" + return self._pd_operator + + def register_state(self, model_role: str, state: FullAttentionState | BoundedAttentionState) -> None: + """按模型角色和具体方案注册 State，保留原对象及其存储。 + + 成员名和 state_name 均为 {model_role}_{state.get_state_type()}_state。 + model_role 为 main 或 draft0..draft35；State 只返回中间的算法/存储方案标记， + 不自行添加角色或 _state。各 State 基类的 get_state_type 注释规定具体命名格式。 + 同角色、同标记不能重复注册；成员名是资源身份，不代表完整布局兼容性。 + """ + assert not self._state_registration_finished, "State registration is already finished" + assert model_role in MODEL_ROLES, f"unknown model role {model_role}" + state_type = state.get_state_type() + assert state_type.isidentifier() and state_type.islower(), f"invalid State type {state_type}" + name = f"{model_role}_{state_type}_state" + assert name not in self.__dict__, f"State {name} is already bound" + assert all( + state is not value for value in self.__dict__.values() + ), "a State cannot belong to multiple model roles" + # config 只描述布局，资源身份仅保存在 State；绑定时保留原配置对象。 + state.state_name = name + setattr(self, name, state) + + @lru_cache(maxsize=None) + def get_state( + self, + model_role: str, + state_type_filter: str, + logical_layer_id: int, + state_class: type[FullAttentionState | BoundedAttentionState] | None = None, + ) -> FullAttentionState | BoundedAttentionState: + """按模型角色、算法/方案字符串和逻辑层选择唯一 State，缓存查询结果。 + + model_role 为 main 或 draft0..draft35，按完整角色前缀匹配，避免 draft1 匹配 draft10。 + state_type_filter 在 get_state_type() 返回的字符串中做包含匹配，如 gqa、gqa_int8、mla_packed。 + 通用推理模板可传空字符串，并通过 state_class 指定资源类别；具体算法传自己的标记。 + 同层存在多类资源时，state_class 可进一步区分；最终必须恰好匹配一个 State。 + 全部模型 State 创建完成后才允许查询；后续原地分配存储不改变缓存对象身份。 + """ + assert self._state_registration_finished, "State registration must finish before querying" + assert model_role in MODEL_ROLES, f"unknown model role {model_role}" + matches = [ + state + for name, state in self.__dict__.items() + if name.startswith(model_role + "_") + and isinstance(state, (FullAttentionState, BoundedAttentionState)) + and state_type_filter in state.get_state_type() + and (state_class is None or isinstance(state, state_class)) + and state.contains_logical_layer(logical_layer_id) + ] + assert len(matches) == 1, ( + f"expected exactly one attention State: {model_role=}, {state_type_filter=}, {logical_layer_id=}, " + f"matches={[state.state_name for state in matches]}" + ) + return matches[0] + + def load_full_att_calibration(self, args: StartArgs | EasyDict) -> None: + """全部 main/draft State 就绪后统一读取静态标定，并交给具体 State 绑定。 + + 文件沿用物理行顺序：main、draft0..draft35；同角色有多个 State 时按最小逻辑层号 + 排序，State 内保留 logical_layer_groups 顺序。非静态 FP8 的 full State 同样占用 + 对应物理行区间，共享逻辑层只占一行；CPU 布局创建不需要调用本流程。 + """ + import json + from .full_att.gqa.static_fp8_state import StaticFP8FullAttentionState + + full_states = self.get_full_attention_states() + calibration_states = [state for state in full_states if isinstance(state, StaticFP8FullAttentionState)] + if not calibration_states: + return + + # 阶段一：只读取一次文件；无文件时由 per_tensor State 绑定单位 scale。 + calibration = None + path = args.kv_quant_calibration_config_path + if path is not None: + with open(path) as file: + calibration = json.load(file) + if any(state.config.scale_mode == "per_tensor" for state in calibration_states): + from lightllm.utils.config_utils import get_model_architectures + + assert calibration["architectures"] == get_model_architectures( + args.model_dir + ), "calibration architecture does not match the model" + + # 阶段二：显式使用模型角色顺序，避免注册先后及 draft10 的字符串排序影响标定行。 + storage_layer_start = 0 + for role in MODEL_ROLES: + states = sorted( + self.get_full_attention_states(role), + key=lambda state: ( + min(layer for group in state.config.logical_layer_groups for layer in group), + state.state_name, + ), + ) + for state in states: + if isinstance(state, StaticFP8FullAttentionState): + state.bind_calibration( + calibration=calibration, + storage_layer_start=storage_layer_start, + ) + storage_layer_start += len(state.config.logical_layer_groups) + logger.info( + "Full-attention calibration bound: states=%d, source=%s", len(calibration_states), path or "unit scales" + ) + + def initialize_storage( + self, + args: StartArgs | EasyDict, + token_capacity: int | None, + mem_fraction: float, + device: str | torch.device | int = "cuda", + ) -> None: + """全部模型完成资源准备后，规划共享容量并初始化缓存存储。 + + 调用前，所有模型的权重和绑定请求容量的运行态已创建，full-attention + 已完成创建标记和所需的静态标定绑定，大小页池尚未分配。调用方在本方法返回后， + 依次执行各模型的运行时初始化，由模型读取 allocator 中的最终容量。 + """ + assert self._state_registration_finished, "State registration must finish before resource initialization" + + # 阶段一：规划共享 token 容量，并向下对齐到完整 KV 页。 + # 未指定容量时，根据剩余显存评估；固定运行态的显存已经在前面占用。 + # PD 模式还需为传输 buffer 预留空间，布局按全部 full-attention State 计算。 + assert self.get_full_attention_states(), "the shared token index domain requires a full-attention State" + pd_reserved_bytes = 0 + if args.run_mode in ("prefill", "decode"): + pd_reserved_bytes = self.get_pd_operator().get_page_bytes() * args.pd_kv_page_num + if token_capacity is None: + # 权重和固定运行态已分配，用当前设备剩余显存扣除保留比例及 PD buffer。 + torch.cuda.empty_cache() + gpu_device = torch.device("cuda", get_current_device_id()) + free_bytes, total_bytes = torch.cuda.mem_get_info(gpu_device) + budget = int(free_bytes - total_bytes * (1 - mem_fraction)) - pd_reserved_bytes + + # 按每页开销直接计算容量，预留一个 HOLD 页；静态量化参数占比较小，不单独扣减。 + assert args.page_size > 0, "capacity planning requires a positive page size" + page_bytes = self.get_full_att_token_bytes() * args.page_size + if page_bytes > budget: + raise MemoryError("attention storage budget cannot fit HOLD pages") + token_capacity = (budget // page_bytes - 1) * args.page_size + + # 各 rank 的显存余量可能不同，统一采用最小容量。 + if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: + capacity = torch.tensor(token_capacity, dtype=torch.int64, device=gpu_device) + torch.distributed.all_reduce(capacity, op=torch.distributed.ReduceOp.MIN) + token_capacity = capacity.item() + logger.info( + f"Attention storage budget: {budget / 1024**3:.3f} GiB, " + f"transfer reserve: {pd_reserved_bytes / 1024**2:.3f} MiB, " + f"profiled token capacity: {token_capacity} (mem_fraction={mem_fraction})" + ) + token_capacity = token_capacity // args.page_size * args.page_size + assert token_capacity > 0, "attention capacity must contain at least one complete page" + + # 阶段二：管理器为已创建的 State 分配实际存储与共享 allocator。 + # State 对象保持不变，查询缓存继续引用相同对象；固定运行态不重复分配。 + big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num + self.allocate_storage( + token_capacity=token_capacity, + page_size=args.page_size, + big_page_tokens=big_page_tokens, + small_page_capacity=args.linear_att_cache_size or 0, + device=device, + ) + + # 创建完成只表示对象集合就绪；实际容量和存储在后续阶段初始化。 + def finish_state_registration(self) -> None: + """确认全部模型的 State 对象已创建并装配完成，允许查询 meta 定义并开始容量规划。""" + assert not self._state_registration_finished, "State registration is already finished" + self._state_registration_finished = True + + def allocate_storage( + self, + token_capacity: int, + page_size: int, + big_page_tokens: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cuda", + ) -> None: + """在 State 对象创建完成后分配 full-attention 存储、共享 allocator 和 bounded 大小页。 + + 请求容量对应的 bounded 运行态（包括压缩窗口）已由各模型创建，这里仅补齐缓存存储。 + """ + assert self._state_registration_finished, "State registration must finish before storage initialization" + assert self.full_att_allocator is None, "resources already initialized" + full_states = self.get_full_attention_states() + assert full_states, "the shared token index domain requires a full-attention State" + for state in full_states: + # 管理器统一设置实际页规格和设备；State 的分配接口只接收 token 容量。 + state.page_size = page_size + state.device = torch.device(device) + state.allocate_storage(token_capacity) + self.full_att_allocator = FullAttKVAllocator(token_capacity, page_size) + + # bounded-attention 大页池末尾保留 CPU 加载/卸载暂存槽位，普通分配不使用它们。 + bounded_states = self.get_bounded_attention_states() + if bounded_states: + assert big_page_tokens > 0, "bounded attention pages require a positive token boundary" + big_page_capacity = (token_capacity + big_page_tokens - 1) // big_page_tokens + 2 + self.bounded_att_big_page_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) + self.bounded_att_small_page_allocator = ( + StatePageAllocator(small_page_capacity) if small_page_capacity else None + ) + for state in bounded_states: + # 无快照区域的 State 在合法缓存边界无需历史，不为它分配大小页 Tensor。 + # 共享槽位 allocator 仍供其他 State 和缓存流程使用，不随单个 State 重复创建。 + if not state.get_cpu_transfer_regions(): + continue + state.initialize_page_pools( + big_page_capacity=big_page_capacity, + small_page_capacity=small_page_capacity, + pin_memory=torch.device(device).type == "cuda", + ) + self.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 + self.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 + + # 存储分配可能改变物理页规格，清除 State 的 CPU/PD 区域和管理器的派生布局缓存。 + for state in (*full_states, *bounded_states): + state.get_cpu_transfer_regions.cache_clear() + state.get_pd_transfer_regions.cache_clear() + # Layout cache 归属于各 Operator；State 存储规格变化后统一失效。 + self.get_cpu_cache_operator().clear_layout_cache() + self.get_pd_operator().clear_layout_cache() + + def get_full_attention_states(self, role: str | None = None) -> tuple[FullAttentionState, ...]: + """返回全部或指定角色的 full-attention State；只缓存对象引用。""" + return self._get_states(FullAttentionState, role) + + def get_bounded_attention_states(self, role: str | None = None) -> tuple[BoundedAttentionState, ...]: + """返回全部或指定角色的 bounded-attention State。""" + return self._get_states(BoundedAttentionState, role) + + def get_full_att_token_bytes(self) -> int: + """累计当前 TP rank 全部 full-attention State 的单 token 存储字节数。""" + total = 0 + for state in self.get_full_attention_states(): + total += state.get_token_bytes() + return total + + # ======================================================================== + # 二、Full attention：token 数据复制 + # ======================================================================== + + # Full attention：在所有参与 State 内复制相同的物理 token 索引，不修改请求映射。 + def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: + """复制 KV 及随 token 保存的量化参数、indexer 字段等完整 token 数据。""" + full_states = self.get_full_attention_states() + assert full_states, "no full-attention storage in this composition" + for state in full_states: + state.copy_full_att_tokens(src_indices, dst_indices) + + # ======================================================================== + # 三、Bounded attention：运行态重置、大小页保存与恢复、MTP 提交 + # ======================================================================== + + # 空前缀初始化：只处理 bounded attention；无此类 State 时可为空操作。 + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """将指定槽位的 bounded 运行态及控制信息恢复到长度 0，保留 KV 和状态页。""" + for state in self.get_bounded_attention_states(): + state.reset_bounded_att_state(req_indices) + + # Bounded attention 状态页：供 radix/CPU cache 复用；未启用对应页池时操作报错。 + # 调用方提供已分配的大小页槽位；发布前完成保存，解除引用并结束使用后直接归还 allocator。 + def save_bounded_att_runtime_to_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """按各请求已执行的绝对 token 长度保存全部 bounded State；三个列表逐项对应。""" + bounded_states = self.get_bounded_attention_states() + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator + assert allocator is not None, f"{page_kind} pages are disabled" + assert ( + len(page_slot_ids) == len(req_indices) == len(state_token_lens) + ), "one length and slot are required per request" + for state in bounded_states: + state.save_bounded_att_runtime_to_pages( + req_indices=req_indices, + state_token_lens=state_token_lens, + page_kind=page_kind, + page_slot_ids=page_slot_ids, + ) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """按各快照对应的绝对 token 长度恢复运行态及 ring 位置等控制信息。""" + bounded_states = self.get_bounded_attention_states() + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + allocator = self.bounded_att_big_page_allocator if page_kind == "big" else self.bounded_att_small_page_allocator + assert allocator is not None, f"{page_kind} pages are disabled" + assert ( + len(page_slot_ids) == len(req_indices) == len(state_token_lens) + ), "one length and slot are required per request" + assert not page_slot_ids or bounded_states, "state restore requires participating bounded states" + for state in bounded_states: + state.restore_bounded_att_runtime_from_pages( + req_indices=req_indices, + state_token_lens=state_token_lens, + page_kind=page_kind, + page_slot_ids=page_slot_ids, + ) + + # MTP 推测解码：只提交指定角色的已接受运行态，由具体 State 管理候选区和游标。 + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + role: str, + ) -> None: + """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" + assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + assert role in MODEL_ROLES, f"unknown model role {role}" + for state in self.get_bounded_attention_states(role): + state.commit_bounded_att_state(req_indices, accepted_steps) + + # ======================================================================== + # 四、CPU/DP/PD 操作均由对应 Operator 提供。 + # ======================================================================== + + # ======================================================================== + # 七、私有辅助实现 + # ======================================================================== + + # 私有实现：缓存创建完成后的成员集合，模型逐层计算直接访问具体 State 成员。 + @lru_cache(maxsize=None) + def _get_states( + self, + state_class: type[FullAttentionState] | type[BoundedAttentionState], + role: str | None, + ): + """按 State 类型和模型角色遍历规范成员，只缓存对象引用。""" + assert self._state_registration_finished, "State registration must finish before querying" + assert role is None or role in MODEL_ROLES, f"unknown model role {role}" + states = ( + state + for state in self.__dict__.values() + if isinstance(state, state_class) and (role is None or state.state_name.startswith(role + "_")) + ) + return tuple(sorted(states, key=lambda state: state.state_name)) diff --git a/lightllm/common/attention_state/pd_operator.py b/lightllm/common/attention_state/pd_operator.py new file mode 100644 index 0000000000..ec1c35c190 --- /dev/null +++ b/lightllm/common/attention_state/pd_operator.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +from collections.abc import Mapping +from copy import copy +from dataclasses import replace +from functools import lru_cache +from typing import TYPE_CHECKING, Literal + +import torch + +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.log_utils import init_logger + +if TYPE_CHECKING: + from lightllm.common.attention_state.manager import AttentionStateManager + + +logger = init_logger(__name__) + + +class PDOperator: + """封装 PD 分离所需的联合布局、页大小和读写操作。""" + + def __init__(self, owner: AttentionStateManager): + self.owner = owner + + @lru_cache(maxsize=None) + def get_pd_full_att_page_layouts(self) -> Mapping[str, StateTransferLayout]: + page_tokens = get_env_start_args().pd_kv_page_size + assert page_tokens > 0, "PD pages require a positive token capacity" + layouts, offset = {}, 0 + for state in self.owner.get_full_attention_states(): + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_pd_transfer_regions() + ) + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=page_tokens, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + @lru_cache(maxsize=None) + def get_pd_bounded_att_page_layouts(self) -> Mapping[str, StateTransferLayout]: + layouts, offset = {}, 0 + for state in self.owner.get_bounded_attention_states(): + regions = tuple( + replace(region, byte_start=offset + region.byte_start, byte_end=offset + region.byte_end) + for region in state.get_pd_transfer_regions() + ) + state_end = (max(region.byte_end for region in regions) + 7) // 8 * 8 + layouts[state.state_name] = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=None, + byte_start=offset, + byte_end=state_end, + regions=regions, + ) + offset = state_end + return layouts + + @lru_cache(maxsize=None) + def get_page_bytes(self) -> int: + token_layouts = self.get_pd_full_att_page_layouts() + runtime_layouts = self.get_pd_bounded_att_page_layouts() + full_bytes = max((layout.byte_end for layout in token_layouts.values()), default=0) + runtime_bytes = max((layout.byte_end for layout in runtime_layouts.values()), default=0) + if token_layouts and full_bytes < runtime_bytes: + logger.warning( + f"PD KV page size ({full_bytes} bytes) is smaller than the complete attention runtime " + f"({runtime_bytes} bytes); using {runtime_bytes} bytes per transfer page" + ) + return max(full_bytes, runtime_bytes) + + def clear_layout_cache(self) -> None: + self.get_pd_full_att_page_layouts.cache_clear() + self.get_pd_bounded_att_page_layouts.cache_clear() + self.get_page_bytes.cache_clear() + + def create_ipc_copy(self) -> AttentionStateManager: + """保留运行态 Tensor，移除 allocator 和 bounded 页池。""" + ipc_copy = copy(self.owner) + ipc_copy.full_att_allocator = None + ipc_copy.bounded_att_big_page_allocator = None + ipc_copy.bounded_att_small_page_allocator = None + for state in self.owner.get_bounded_attention_states(): + state_copy = copy(state) + state_copy.big_pages_pool = None + state_copy.small_pages_pool = None + setattr(ipc_copy, state.state_name, state_copy) + ipc_copy._reset_operators() + return ipc_copy + + def write_page( + self, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + kind: Literal["kv", "att_state"], + layouts: Mapping[str, StateTransferLayout], + token_indices: torch.Tensor | None = None, + req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: + match kind: + case "kv": + assert token_indices is not None, "KV PD pages require token indices" + for state in self.owner.get_full_attention_states(): + state.write_to_pd_page( + token_indices=token_indices, + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + layout=layouts[state.state_name], + ) + case "att_state": + assert req_idx is not None, "state PD pages require a request index" + assert ( + state_token_len is not None and state_token_len >= 0 + ), "PD state requires its prefill prefix length" + for state in self.owner.get_bounded_attention_states(): + state.write_to_pd_page( + req_idx=req_idx, + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + state_token_len=state_token_len, + layout=layouts[state.state_name], + ) + case _: + raise AssertionError(f"unknown PD page kind {kind}") + + def read_page( + self, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + kind: Literal["kv", "att_state"], + layouts: Mapping[str, StateTransferLayout], + token_indices: torch.Tensor | None = None, + req_idx: int | None = None, + state_token_len: int | None = None, + ) -> None: + match kind: + case "kv": + assert token_indices is not None, "KV PD pages require token indices" + for state in self.owner.get_full_attention_states(): + state.read_from_pd_page( + pd_page_buffers=pd_page_buffers, + token_indices=token_indices, + pd_page_index=pd_page_index, + layout=layouts[state.state_name], + ) + case "att_state": + assert req_idx is not None, "state PD pages require a request index" + assert ( + state_token_len is not None and state_token_len >= 0 + ), "PD state requires its prefill prefix length" + for state in self.owner.get_bounded_attention_states(): + state.read_from_pd_page( + pd_page_buffers=pd_page_buffers, + req_idx=req_idx, + pd_page_index=pd_page_index, + state_token_len=state_token_len, + layout=layouts[state.state_name], + ) + case _: + raise AssertionError(f"unknown PD page kind {kind}") diff --git a/lightllm/common/attention_state/types.py b/lightllm/common/attention_state/types.py new file mode 100644 index 0000000000..38c1b95691 --- /dev/null +++ b/lightllm/common/attention_state/types.py @@ -0,0 +1,237 @@ +"""注意力存储字段与传输位置的只读描述，不执行数据搬运。 + +TensorField 描述本地 Tensor；AxisMapping 描述某个轴的本地、外部范围。 +TransferRegion / StateTransferLayout 用于 State 与 CPU/PD 页之间的搬运； +DP 在对应 TP 分片间直接复制，不需要字段映射类型。 + +byte_start、byte_end 使用字节单位，表示左闭右开的区间；shape、轴切片起点和 length 使用元素单位。 +层维度完整、同序对应，token/request 索引由每次搬运的调用参数提供。 +frozen=True 只禁止重新赋值描述属性，不冻结被引用 Tensor 的数据。 +KV 示例统一使用 head_dim=128。 + +PageKind 使用示例，选择 bounded attention 大页池： + >>> page_kind: PageKind = "big" + >>> assert page_kind == "big" +""" + +from dataclasses import dataclass +from typing import Literal + +import torch + +# ======================================================================== +# 一、状态页类别 +# ======================================================================== + +# 选择 bounded attention 的页池及共享分配域，不表示 KV token 页或传输 buffer 的大小。 +# "big" 对应 big_pages_pool 与 bounded_att_big_page_allocator； +# "small" 对应 small_pages_pool 与 bounded_att_small_page_allocator。 +# 两个域中的同一个 slot_id 指向不同状态页。调用方先在对应 allocator 分配槽位， +# 再将类别传给 save_bounded_att_runtime_to_pages / restore_bounded_att_runtime_from_pages 等操作，例如： +# manager.save_bounded_att_runtime_to_pages(req_indices=[0], state_token_lens=[256], page_kind="big", page_slot_ids=[3]) +# 上例假定请求槽位 0 的已接受长度为 256，且大页槽位 3 已由调用方分配。 +PageKind = Literal["big", "small"] + + +# ======================================================================== +# 二、本地存储字段与轴范围 +# ======================================================================== + + +@dataclass(frozen=True) +class TensorField: + """为本地 Tensor 补充轴语义，供 State 内部布局计算及搬运使用。 + + 创建描述只引用已有 Tensor，不分配或复制其底层存储。具体 State 按固定布局组织 + k、v、动态量化 scale、Conv、SSM 等搬运字段，字段名称由所在字典的 key 提供； + packed buffer 的字段可使用切片视图。 + + 示例：本地 K 数据有 2 层、8 个 token、2 个 heads，每个 head 有 128 个元素。 + >>> k = torch.empty((2, 8, 2, 128), dtype=torch.float16) + >>> field = TensorField( + ... tensor=k, axes=("layer", "token", "head", "dim") + ... ) + >>> assert field.tensor is k + """ + + # 实际存储或其视图，提供 dtype、shape、stride、device 等物理信息。 + # 可引用非连续的 packed/分层存储视图，搬运代码沿用其真实 stride。 + tensor: torch.Tensor + + # 按 Tensor 维度顺序列出语义名称，数量等于 tensor.ndim，名称不能重复。 + # 例如 token/head 用于 KV 选取和 TP 切片，request/page 用于运行态或页池描述； + # 选择某个请求或状态页后，搬运视图会移除已选定的 request/page 轴。 + axes: tuple[str, ...] + + +@dataclass(frozen=True) +class AxisMapping: + """某一轴上的逻辑连续范围映射，起点与长度均以元素为单位。 + + TransferRegion 中的方向固定为本地 State 到外部页，读取时反向应用。 + 环形轴按对应 Tensor 的实际轴长度取模，不要求两侧容量相同。 + 环形标记只是布局描述，须由具体 State 的专用 kernel 实现；普通 TP 切片路径不支持取模。 + 当前窗口压缩器直接按 state_token_len 搬运 ring，不依赖这两个标记。 + + 示例：本地 head[0:2] 写到全局传输区域 head[2:4]，读回时反向复制。 + >>> head_mapping = AxisMapping(axis="head", local_start=0, external_start=2, length=2) + >>> assert head_mapping.external_start + head_mapping.length == 4 + + 环形示例：本地 window 容量为 4，从槽位 3 起搬运 4 个位置到线性外部页。 + 本地位置依次为 [3, 0, 1, 2]，外部位置依次为 [0, 1, 2, 3]。 + >>> window_mapping = AxisMapping( + ... axis="window", local_start=3, external_start=0, length=4, local_circular=True, + ... ) + >>> assert window_mapping.local_circular and not window_mapping.external_circular + """ + + # 要映射的轴名，必须同时存在于本地和外部字段中，例如 head、channel、dim。 + # layer 固定完整复制，token 由本次传输数量控制，这两个轴不放入 AxisMapping。 + axis: str + + # 本地 State 轴的起始元素索引，包含该位置且必须非负。 + # 写出时从此处读取，读回时向此处写入。 + local_start: int + + # 外部页轴的起始元素索引，包含该位置且必须非负。 + # 写出时向此处写入，读回时从此处读取。 + # 用于将本地 TP 分片放到全局 head/channel 的对应位置。 + external_start: int + + # 搬运的元素数量，必须为正；普通轴使用半开区间 [start, start + length)。 + # 环形轴对 start + i 取模，length 不得超过该侧 Tensor 的轴容量，避免重复覆盖。 + # 它是该轴上的长度，不是字节数，也不是整个 Tensor 的元素总数。 + length: int + + # 本地 State 的该轴采用环形寻址；容量取自本地 Tensor，而不是外部 region.shape。 + # 例如 tail/SWA 的 window 轴。读回时仍作用于本地侧，不随实际搬运方向交换含义。 + local_circular: bool = False + + # 外部区域的该轴采用环形寻址；容量取自外部 Tensor 的对应轴。 + # 常规 CPU/PD 页按逻辑顺序线性保存，使用默认 False。 + external_circular: bool = False + + +# ======================================================================== +# 三、CPU/PD 外部页布局 +# ======================================================================== + + +@dataclass(frozen=True) +class TransferRegion: + """描述 State 内一个字段与 CPU/PD 外部页区域之间的拷贝映射。 + + 由 StateTransferLayout.regions 组织，同一字段可以对应多个区域映射。 + 区域的 shape 覆盖外部完整数据，axis_mappings 选择当前 rank 的子范围。 + 本地和外部区域的层数必须一致，并按相同存储顺序完整复制。 + + 示例：外部 K 区域有 4 个全局 heads，head_dim=128，当前 rank 持有后 2 个 heads。 + 整个区域占 2 * 8 * 4 * 128 * 2 = 16384 字节（16 KiB），位于页内 [16384, 32768)。 + >>> k_region = TransferRegion( + ... field="k", + ... byte_start=16384, byte_end=32768, + ... dtype=torch.float16, + ... axes=("layer", "token", "head", "dim"), + ... shape=(2, 8, 4, 128), + ... axis_mappings=( + ... AxisMapping(axis="head", local_start=0, external_start=2, length=2), + ... AxisMapping(axis="dim", local_start=0, external_start=0, length=128), + ... ), + ... write_enabled=True, + ... ) + """ + + # 本地搬运字段字典的 key，用来定位 k/v/scale/conv/ssm 等 Tensor。 + # 同一字段可出现在多个 region 中，例如 Conv 的 Q/K/V 使用不同 channel 切片。 + field: str + + # 字段区域起点，包含该字节，必须按 dtype 的元素大小对齐。 + # State 的 get_cpu/pd_transfer_regions 返回相对本 State 起点的区间； + # 管理器将起止位置同时加上 State 的页内起点，最终区间统一相对整页。 + # CPU/PD 都从页集合的第一维选择物理页；字段区间不包含物理页索引。 + byte_start: int + + # 字段区域终点，不包含该字节，与 byte_start 使用相同的坐标基准。 + # byte_end - byte_start 必须等于 prod(shape) * dtype.itemsize， + # 仅包含字段数据；字段之间的对齐空隙不计入该区间。 + byte_end: int + + # 外部区域按此类型解释原始字节，需与本地搬运字段的 dtype 一致。 + # packed 字段可分别使用 uint8、float16 等实际类型；此描述不执行类型转换。 + dtype: torch.dtype + + # 外部区域各维度的语义顺序，与 shape 一一对应且名称唯一。 + # 可与本地 Tensor 的轴顺序不同，搬运时按名称对齐；本地视图需具有相同轴集合。 + axes: tuple[str, ...] + + # 完整外部区域的各轴元素数量，不包含外部页池的 page 轴，也不是本 rank 的局部形状。 + # 例如全局 4 heads、本地 2 heads 时，这里的 head 维仍为 4。 + # 区域连续排列，stride 由 shape 和 dtype 推导；区域末尾必须位于当前页内。 + shape: tuple[int, ...] + + # 描述本地字段到外部区域的轴范围，每个非 layer/token 轴必须恰好出现一次。 + # 无分片的轴也显式填写完整范围，例如 dim: 0 -> 0、length=head_dim。 + # 两个方向共用同一组本地/外部起点，仅反转数据流向；映射应用于全部层和本次有效 token。 + # 环形寻址由支持它的具体 State 搬运实现处理；请求窗口起点根据本次绝对前缀长度计算， + # 不回写缓存布局，也不作为 kernel 的编译期常量。 + axis_mappings: tuple[AxisMapping, ...] + + # 当前 TP rank 是否负责向外部区域写入；False 不影响从该区域读回。 + # KV head 有 TP 副本时，每组只选一个 rank 写入；复制的压缩窗口只由 rank 0 写。 + # 标记属于 region，因为同一 State 内各字段的 TP 分片/复制规则可能不同。 + write_enabled: bool = True + + +@dataclass(frozen=True) +class StateTransferLayout: + """描述一个 State 的 CPU/PD 传输布局，按字段组织 TransferRegion。 + + 每个 layout 通过 state_name 绑定一个 State；regions 描述该 State 各字段 + 与外部页之间的映射，具体轴范围由各 region 的 AxisMapping 描述。 + byte_start / byte_end 记录本 State 在联合页中的完整预留区间，包含自身对齐空间。 + CPU/PD 根据页集合的行选择物理页；PD 通过显式 pd_page_index 指定目标页。 + 实际页容量由 manager 统一计算，可以大于当前 State 的终点。 + 示例：一张 64 KiB 的页中，本 State 的 K/V 各占 16 KiB， + 分别位于 [16 KiB, 32 KiB) 和 [32 KiB, 48 KiB)，前后空间可属于其他资源。 + >>> from dataclasses import replace + >>> k_region = TransferRegion( + ... field="k", byte_start=16384, byte_end=32768, dtype=torch.float16, + ... axes=("layer", "token", "head", "dim"), shape=(2, 8, 4, 128), + ... axis_mappings=( + ... AxisMapping(axis="head", local_start=0, external_start=2, length=2), + ... AxisMapping(axis="dim", local_start=0, external_start=0, length=128), + ... ), + ... ) + >>> layout = StateTransferLayout( + ... state_name="main_gqa_fp16_state", + ... tokens_per_page=8, + ... byte_start=16384, byte_end=49152, + ... regions=(k_region, replace(k_region, field="v", byte_start=32768, byte_end=49152)), + ... ) + >>> assert layout.regions[1].byte_end == layout.byte_end == 49152 + """ + + # 当前 State 的规范身份，与管理器成员名、state.state_name 一致。 + # 用于从联合页布局字典中选择对应资源，区分 main/draft 及不同存储方案。 + state_name: str + + # 整张外部页可容纳的 token 数，不是本次实际传输数量；含 KV 的页使用正整数。 + # CPU/PD 分别使用各自的全局页配置，页容量不随本次传输长度改变。 + # CPU 的 KV + bounded 联合页中，bounded layout 也沿用该 token 容量。 + # 纯运行态页（如 PD 的 att_state 页）为 None，表示按快照而非 token 序列搬运。 + tokens_per_page: int | None + + # 本 State 在外部页中的起始字节，包含该位置，始终相对整页起点。 + # 多个 State 依次排列，后一 State 的起点为前一 State 的终点。 + # regions 已使用页内坐标，搬运时不再重复加上此起点。 + byte_start: int + + # 本 State 在外部页中的结束字节，不包含该位置；涵盖全部字段及本 State 的对齐、预留空间。 + # byte_end - byte_start 是本 State 的占用，不是整张物理页的跨度。 + # 同类布局的最大 byte_end 表示整组 State 的占用；CPU 物理页再按 16 字节对齐， + # PD 物理页由 PDOperator.get_page_bytes() 取 KV 与完整运行态布局占用的最大值。 + byte_end: int + + # 当前 State 参与传输的字段区域；其他 State 的区域放在它们各自的 layout 中。 + # 包含所需的 KV、动态量化参数或运行态字段，所有偏移均相对同一张完整外部页。 + regions: tuple[TransferRegion, ...] diff --git a/lightllm/common/attention_state/utils.py b/lightllm/common/attention_state/utils.py new file mode 100644 index 0000000000..5de4267819 --- /dev/null +++ b/lightllm/common/attention_state/utils.py @@ -0,0 +1,166 @@ +"""注意力 State 的命名和传输视图工具。""" + +from __future__ import annotations + +from collections.abc import Callable +from copy import deepcopy +from functools import WRAPPER_ASSIGNMENTS, wraps +from typing import ParamSpec, TypeVar + +import torch + +_P = ParamSpec("_P") +_R = TypeVar("_R") + +MAX_DRAFT_MODELS = 36 +MODEL_ROLES = ("main", *(f"draft{i}" for i in range(MAX_DRAFT_MODELS))) + + +# ======================================================================== +# 一、State 命名与返回值副本 +# ======================================================================== + + +def dtype_name(dtype: torch.dtype): + return {torch.bfloat16: "bf16", torch.float16: "fp16", torch.float32: "fp32"}[dtype] + + +def deepcopy_result(func: Callable[_P, _R]) -> Callable[_P, _R]: + """每次返回结果的独立深拷贝，放在 lru_cache 外层以保护缓存原对象。 + + 保留被装饰函数的缓存管理接口，调用方仍可使用 cache_clear/cache_info/cache_parameters。 + """ + + @wraps(func, assigned=WRAPPER_ASSIGNMENTS + ("cache_clear", "cache_info", "cache_parameters")) + def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R: + return deepcopy(func(*args, **kwargs)) + + return wrapped + + +# ======================================================================== +# 二、传输视图：合并两端共同支持的尾部维度 +# ======================================================================== + + +def merge_contiguous_tail_dims( + copies: list[tuple[torch.Tensor, torch.Tensor]], + start_dim: int = 0, + keepdim: bool = False, +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """合并批次两侧共同连续的尾部维度，只创建 view，不复制数据或改变 dtype。 + + 调用方保证批次及 Tensor 非空，同一侧的 shape/stride 相同，两侧待合并的尾部 shape 相同； + 前缀维度的数量和容量可以不同，例如 GPU token 存储与外部页集合。 + 根据首对 Tensor 确定共同的合并范围，再为整批分别创建两侧形状的视图。 + 从末轴向前尝试，遇到任一侧无法合并就停止，不单独整理前面的维度。 + + start_dim 是允许参与合并的最前轴，之前的轴保持独立。例如 + [layer, window, head, dim] 传入 start_dim=2，只合并 head/dim，保留环形 window。 + 支持负数，例如 start_dim=-2 对齐两侧最后两维，不要求两侧前缀维数一致。 + keepdim=True 将被合并的前面各轴保留为 1，元素数集中到末轴，供固定维数的 + token/page kernel 使用；默认 False 直接减少维数，供 bounded 聚合 kernel 使用。 + 直接按 shape/stride 判断连续性，长度为 1 的轴不约束 stride。 + 只依赖公开 Tensor 接口，确定形状后一次性创建 view，不使用可能复制数据的 reshape。 + """ + source, destination = copies[0] + source_shape, destination_shape = source.shape, destination.shape + source_strides, destination_strides = source.stride(), destination.stride() + ndim = len(source_shape) + assert -ndim <= start_dim < ndim, "invalid tail merge boundary" + if start_dim < 0: + start_dim += ndim + tail_dims = ndim - start_dim + assert ( + len(destination_shape) >= tail_dims and source_shape[start_dim:] == destination_shape[-tail_dims:] + ), "copy tail shapes must match" + + # 阶段一：从末轴向前检查；两侧的前缀容量和轴数独立，按距末轴的距离对齐。 + first_tail_dim = ndim - 1 + tail_elements = source_shape[-1] + source_tail_stride, destination_tail_stride = source_strides[-1], destination_strides[-1] + prefix_dim_difference = len(destination_shape) - ndim + for source_axis in range(ndim - 2, start_dim - 1, -1): + destination_axis = source_axis + prefix_dim_difference + axis_size = source_shape[source_axis] + if axis_size != 1: + if tail_elements == 1: + # 内侧只有一个元素，其 stride 不影响寻址；以当前非单元素轴为准。 + source_tail_stride = source_strides[source_axis] + destination_tail_stride = destination_strides[destination_axis] + elif ( + source_strides[source_axis] != tail_elements * source_tail_stride + or destination_strides[destination_axis] != tail_elements * destination_tail_stride + ): + break + # 当前轴长度为 1 时忽略其 stride，但保留已累计的连续区间和元素步长。 + tail_elements *= axis_size + first_tail_dim = source_axis + + if first_tail_dim == ndim - 1: + return copies + + # 阶段二：整批复用普通 tuple 形状，减少 view 的参数解析开销。 + # keepdim 保留固定维数 kernel 所需的单元素轴。 + merge_dims = ndim - first_tail_dim + tail_shape = (1,) * (merge_dims - 1) + (tail_elements,) if keepdim else (tail_elements,) + source_transfer_shape = tuple(source_shape[:-merge_dims]) + tail_shape + destination_transfer_shape = tuple(destination_shape[:-merge_dims]) + tail_shape + if source_transfer_shape == source_shape and destination_transfer_shape == destination_shape: + return copies + return [(src.view(source_transfer_shape), dst.view(destination_transfer_shape)) for src, dst in copies] + + +# ======================================================================== +# 三、传输视图：为整批选择最宽的搬运 dtype +# ======================================================================== + + +def view_max_copy_dtype( + copies: list[tuple[torch.Tensor, torch.Tensor]], +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """将整批视图按最大可用整数宽度搬运，不转换数值、不复制数据。 + + 调用方保证批次非空、dtype 相同，同一侧的 shape/stride 相同，两侧的末轴长度相同。 + 两侧的外层 shape/stride 可以不同。依次尝试 8/4/2 字节宽度，整批统一选择一种 dtype； + 不能扩大时返回原列表。只改变末轴的元素数，保留其余轴及存储间隙。 + 本函数不合并维度；有需要时先调用 merge_contiguous_tail_dims。 + 先用末轴字节数、首对 Tensor 的外层 stride 和 Storage 偏移过滤候选； + 同一侧的 stride 由整批共用，其余切片的 Storage 偏移继续由 view 校验。 + 只捕获创建视图的 RuntimeError，任一视图失败就整批退回较窄的 dtype。 + """ + source, destination = copies[0] + element_bytes = source.element_size() + if element_bytes >= 8 or source.stride(-1) != 1 or destination.stride(-1) != 1: + return copies + row_bytes = source.shape[-1] * element_bytes + # 按行宽选择候选类型，例如 12 字节只尝试 INT32、INT16。 + if row_bytes % 8 == 0: + copy_dtypes = ((8, torch.int64), (4, torch.int32), (2, torch.int16)) + elif row_bytes % 4 == 0: + copy_dtypes = ((4, torch.int32), (2, torch.int16)) + elif row_bytes % 2 == 0: + copy_dtypes = ((2, torch.int16),) + else: + return copies + + # 首对布局即可排除必然失败的候选，避免 MTP Conv 等带间隙切片反复抛异常。 + # 这些值以原 dtype 的元素数计量，扩宽后必须仍能表示为整数 stride/offset。 + strides_and_offsets = ( + *source.stride()[:-1], + *destination.stride()[:-1], + source.storage_offset(), + destination.storage_offset(), + ) + for transfer_bytes, transfer_dtype in copy_dtypes: + if transfer_bytes <= element_bytes: + break + elements_per_unit = transfer_bytes // element_bytes + if any(value % elements_per_unit for value in strides_and_offsets): + continue + # 批次中其他切片的起点可能不同；保留 view 的整批校验和较窄类型回退。 + try: + return [(src.view(transfer_dtype), dst.view(transfer_dtype)) for src, dst in copies] + except RuntimeError: + continue + return copies diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index ee8d06ed9a..1122df3551 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -4,7 +4,7 @@ import torch from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Optional, TYPE_CHECKING, Tuple, Union, Dict +from typing import Optional, TYPE_CHECKING, Tuple, Dict from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_env_start_args @@ -75,16 +75,6 @@ def uses_causal_attention(self) -> bool: is_parallel_block_draft = self.model.is_mtp_draft_model and args.mtp_mode in ("dspark", "dflash", "dflash2") return not is_parallel_block_draft - def _find_layer_index( - self, k: torch.Tensor, v: torch.Tensor, att_state: Union["BasePrefillAttState", "BaseDecodeAttState"] - ) -> int: - kv_buffer = att_state.infer_state.mem_manager.kv_buffer - layer_count = len(kv_buffer) - find_dict = {kv_buffer[i].data_ptr(): i for i in range(layer_count)} - key = min(k.data_ptr(), v.data_ptr()) - assert key in find_dict - return find_dict[key] - @dataclass class AttControl: @@ -113,6 +103,8 @@ class AttControl: linear_att_prefill_dict: Dict = None linear_att_decode: bool = False linear_att_decode_dict: Dict = None + # 静态 FP8 标定按此逻辑层号查询 State 的物理行；其他后端可不使用。 + logical_layer_id: Optional[int] = None @dataclass diff --git a/lightllm/common/basemodel/attention/fa3/fp8.py b/lightllm/common/basemodel/attention/fa3/fp8.py index 01633e294b..4e33d187a0 100644 --- a/lightllm/common/basemodel/attention/fa3/fp8.py +++ b/lightllm/common/basemodel/attention/fa3/fp8.py @@ -3,7 +3,10 @@ import torch from ..base_att import AttControl from lightllm.utils.sgl_utils import flash_attn_with_kvcache -from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import q_per_head_static_fp8_quant +from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import ( + q_per_head_fp8_quant, + q_per_head_static_fp8_quant, +) from lightllm.utils.log_utils import init_logger from .fp import Fa3AttBackend, Fa3PrefillAttState, Fa3DecodeAttState @@ -39,10 +42,13 @@ class Fp8Fa3PrefillAttState(Fa3PrefillAttState): def init_state(self): super().init_state() batch_size = self.infer_state.batch_size - mem_manager = self.backend.model.mem_manager + state = getattr( + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_head_state", + ) - offline_scales: torch.Tensor = mem_manager.scales - head_num = mem_manager.head_num + offline_scales: torch.Tensor = state.scales + head_num = state.config.head_num # 为了减少推理计算量，在推理外部初始化k_descale和v_descale self.k_descale = ( offline_scales[:, :head_num].view(-1, 1, head_num).expand(offline_scales.shape[0], batch_size, head_num) @@ -69,10 +75,16 @@ def prefill_att( k=k, v=v, alloc_func=alloc_func, + logical_layer_id=att_control.logical_layer_id, ) def _fp8_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, ) -> torch.Tensor: self.backend: Fp8Fa3AttBackend = self.backend # for typing @@ -80,12 +92,29 @@ def _fp8_prefill_att( q_head_dim = q.shape[2] k_head_num = k.shape[1] k_head_dim = k.shape[2] - cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) - cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) - layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, att_state=self) - static_q_scales = self.backend.model.mem_manager.q_scales[layer_index] - q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales) - q_scale = static_q_scales.view(1, k_head_num).expand(self.infer_state.b_seq_len.shape[0], k_head_num) + cache_k = k.view( + size=(-1, 1, k_head_num, k_head_dim), + ).view(torch.float8_e4m3fn) + cache_v = v.view( + size=(-1, 1, k_head_num, k_head_dim), + ).view(torch.float8_e4m3fn) + state = getattr( + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_head_state", + ) + storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) + q = q.reshape(q.shape[0], k_head_num, -1) + if state.q_scales is None: + # 旧标定文件仅包含 K/V；按本次请求的 Q 长度动态计算每个 KV head 的 scale。 + q, q_scale = q_per_head_fp8_quant( + q=q, + seq_lens=self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1], + b1_start_loc=self.cu_seqlens_q, + ) + else: + static_q_scales = state.q_scales[storage_layer_index] + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(self.infer_state.b_seq_len.shape[0], k_head_num) o = flash_attn_with_kvcache( q=q.reshape(-1, q_head_num, q_head_dim), k_cache=cache_k, @@ -99,8 +128,8 @@ def _fp8_prefill_att( window_size=(-1, -1), softcap=0.0, q_descale=q_scale, - k_descale=self.k_descale[layer_index], - v_descale=self.v_descale[layer_index], + k_descale=self.k_descale[storage_layer_index], + v_descale=self.v_descale[storage_layer_index], return_softmax_lse=False, ) return o @@ -114,12 +143,15 @@ class Fp8Fa3DecodeAttState(Fa3DecodeAttState): def init_state(self): self.backend: Fp8Fa3AttBackend = self.backend - mem_manager = self.backend.model.mem_manager + state = getattr( + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_head_state", + ) super().init_state() att_batch_size = self.b_att_seq_len.shape[0] - offline_scales: torch.Tensor = mem_manager.scales - head_num = mem_manager.head_num + offline_scales: torch.Tensor = state.scales + head_num = state.config.head_num # 为了减少推理计算量，在推理外部初始化k_descale和v_descale self.k_descale = ( @@ -152,6 +184,7 @@ def decode_att( k=k, v=v, alloc_func=alloc_func, + logical_layer_id=att_control.logical_layer_id, ) def _fp8_decode_att( @@ -159,21 +192,39 @@ def _fp8_decode_att( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + logical_layer_id: int, alloc_func=torch.empty, ): k_head_num = k.shape[1] k_head_dim = k.shape[2] - cache_k = k.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) - cache_v = v.view(-1, 1, k_head_num, k_head_dim).view(torch.float8_e4m3fn) + cache_k = k.view( + size=(-1, 1, k_head_num, k_head_dim), + ).view(torch.float8_e4m3fn) + cache_v = v.view( + size=(-1, 1, k_head_num, k_head_dim), + ).view(torch.float8_e4m3fn) - layer_index = self.backend._find_layer_index(k=cache_k, v=cache_v, att_state=self) + state = getattr( + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_head_state", + ) + storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) q_head_num = q.shape[1] att_batch_size = self.b_att_seq_len.shape[0] - static_q_scales = self.backend.model.mem_manager.q_scales - q = q_per_head_static_fp8_quant(q.reshape(q.shape[0], k_head_num, -1), static_q_scales[layer_index]) - q_scale = static_q_scales[layer_index].view(1, k_head_num).expand(att_batch_size, k_head_num) + q = q.reshape(q.shape[0], k_head_num, -1) + if state.q_scales is None: + # 使用 FA3 的实际 Q 分组，兼容普通 decode、MTP 多 token 及零长度填充组。 + q, q_scale = q_per_head_fp8_quant( + q=q, + seq_lens=self.cu_seqlens_q[1:] - self.cu_seqlens_q[:-1], + b1_start_loc=self.cu_seqlens_q, + ) + else: + static_q_scales = state.q_scales[storage_layer_index] + q = q_per_head_static_fp8_quant(q=q, scales=static_q_scales) + q_scale = static_q_scales.view(1, k_head_num).expand(att_batch_size, k_head_num) o = flash_attn_with_kvcache( q=q.reshape(-1, q_head_num, k_head_dim), k_cache=cache_k, @@ -187,8 +238,8 @@ def _fp8_decode_att( window_size=(-1, -1), softcap=0.0, q_descale=q_scale.view(att_batch_size, k_head_num), - k_descale=self.k_descale[layer_index], - v_descale=self.v_descale[layer_index], + k_descale=self.k_descale[storage_layer_index], + v_descale=self.v_descale[storage_layer_index], return_softmax_lse=False, ) return o diff --git a/lightllm/common/basemodel/attention/flashinfer/fp8.py b/lightllm/common/basemodel/attention/flashinfer/fp8.py index 189f8daddb..260333e191 100644 --- a/lightllm/common/basemodel/attention/flashinfer/fp8.py +++ b/lightllm/common/basemodel/attention/flashinfer/fp8.py @@ -58,20 +58,30 @@ def prefill_att( k=k, v=v, alloc_func=alloc_func, + logical_layer_id=att_control.logical_layer_id, ) def _fp8_prefill_att( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, alloc_func=torch.empty + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + logical_layer_id: int, + alloc_func=torch.empty, ) -> torch.Tensor: o_tensor = alloc_func(q.shape, q.dtype, device="cuda") k = k.unsqueeze(1).view(torch.float8_e4m3fn) v = v.unsqueeze(1).view(torch.float8_e4m3fn) - layer_index = self.backend._find_layer_index(k=k, v=v, att_state=self) - k_descale = self.infer_state.mem_manager.cpu_scales[layer_index][0] - v_descale = self.infer_state.mem_manager.cpu_scales[layer_index][1] + state = getattr( + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_tensor_state", + ) + storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) + k_descale = state.cpu_scales[storage_layer_index][0] + v_descale = state.cpu_scales[storage_layer_index][1] self.prefill_wrapper.run( - q, - (k, v), + q=q, + paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor, @@ -105,6 +115,7 @@ def decode_att( k=k, v=v, alloc_func=alloc_func, + logical_layer_id=att_control.logical_layer_id, ) def _fp8_decode_att( @@ -112,19 +123,24 @@ def _fp8_decode_att( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + logical_layer_id: int, alloc_func=torch.empty, ): o_tensor = alloc_func(q.shape, q.dtype, device="cuda") k = k.unsqueeze(1).view(torch.float8_e4m3fn) v = v.unsqueeze(1).view(torch.float8_e4m3fn) - layer_index = self.backend._find_layer_index(k=k, v=v, att_state=self) + state = getattr( + self.infer_state.att_state_manager, + f"{self.infer_state.att_state_role}_gqa_fp8_static_per_tensor_state", + ) + storage_layer_index = state.get_storage_layer_index(logical_layer_id=logical_layer_id) - k_descale = self.infer_state.mem_manager.cpu_scales[layer_index][0] - v_descale = self.infer_state.mem_manager.cpu_scales[layer_index][1] + k_descale = state.cpu_scales[storage_layer_index][0] + v_descale = state.cpu_scales[storage_layer_index][1] self.decode_wrapper.run( - q, - (k, v), + q=q, + paged_kv_cache=(k, v), k_scale=k_descale, v_scale=v_descale, out=o_tensor, diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 2cde561013..986214568f 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.bounded_att import GDNState import dataclasses import weakref import torch @@ -139,7 +140,13 @@ 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) + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=GDNState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: @@ -210,16 +217,27 @@ def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step == 0: self._init_normal_decode_state() - elif self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(draft_step + 1) + return + + # MTP 元数据按模型初始化，沿用每个角色只有一个 GDN State 的约束。 + # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 + (state,) = ( + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states( + self.infer_state.att_state_role + ) + if isinstance(state, GDNState) + ) + if self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state(draft_step + 1, req_to_mtp_state_index=state.mtp_state_index) else: - self._init_fixed_mtp_decode_state(draft_step) + self._init_fixed_mtp_decode_state(draft_step, req_to_mtp_state_index=state.mtp_state_index) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx self.b_ssm_buffer_idx = self.infer_state.b_req_idx - def _init_dynamic_mtp_decode_state(self, mtp_size: int): + def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: ( self.b1_mtp_cu_q_seq_len, self.b_conv_buffer_idx, @@ -227,12 +245,12 @@ 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=req_to_mtp_state_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) - def _init_fixed_mtp_decode_state(self, draft_step: int): + def _init_fixed_mtp_decode_state(self, draft_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: mtp_size = draft_step + 1 batch_size = self.infer_state.batch_size assert batch_size % mtp_size == 0, ( @@ -249,7 +267,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 = req_to_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 +299,13 @@ 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) + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=GDNState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=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..cd2f8afb39 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -10,6 +10,7 @@ import torch +from lightllm.common.attention_state.bounded_att import KDAState from lightllm.common.basemodel.attention.base_att import ( AttControl, BaseAttBackend, @@ -83,7 +84,13 @@ 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) + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, + state_type_filter="kda", + logical_layer_id=layer_num, + state_class=KDAState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) conv_states = conv_states[..., : backend.conv_kernel_size - 1] mixed_qkv = causal_conv1d_fn( mixed_qkv.transpose(0, 1), @@ -134,16 +141,27 @@ def init_state(self): mtp_step = self.backend.mtp_step if mtp_step == 0: self._init_normal_decode_state() - elif self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(mtp_step + 1) + return + + # MTP 元数据按模型初始化，沿用每个角色只有一个 KDA State 的约束。 + # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 + (state,) = ( + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states( + self.infer_state.att_state_role + ) + if isinstance(state, KDAState) + ) + if self.backend.uses_dynamic_spec_verify_layout(): + self._init_dynamic_mtp_decode_state(mtp_step + 1, req_to_mtp_state_index=state.mtp_state_index) else: - self._init_fixed_mtp_decode_state(mtp_step) + self._init_fixed_mtp_decode_state(mtp_step, req_to_mtp_state_index=state.mtp_state_index) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx self.b_ssm_buffer_idx = self.infer_state.b_req_idx - def _init_dynamic_mtp_decode_state(self, mtp_size: int): + def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: ( self.b1_mtp_cu_q_seq_len, self.b_conv_buffer_idx, @@ -151,12 +169,12 @@ 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=req_to_mtp_state_index, hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, ) self._init_mtp_ssm_buffer_idx(mtp_size) - def _init_fixed_mtp_decode_state(self, mtp_step: int): + def _init_fixed_mtp_decode_state(self, mtp_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: mtp_size = mtp_step + 1 batch_size = self.infer_state.batch_size assert batch_size % mtp_size == 0, ( @@ -173,7 +191,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 = req_to_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 +222,13 @@ 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) + state = self.infer_state.att_state_manager.get_state( + model_role=self.infer_state.att_state_role, + state_type_filter="kda", + logical_layer_id=layer_num, + state_class=KDAState, + ) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=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/attention/nsa/dsv4_fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py index 536db3d4fe..93560de07a 100644 --- a/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/dsv4_fp8_flashmla_sparse.py @@ -18,42 +18,72 @@ class DeepseekV4FlashMlaFp8SparseAttBackend(BaseAttBackend): def __init__(self, model): super().__init__(model=model) - self.cache_config = model.req_manager.cache_config + from lightllm.models.deepseek_v4.config import create_attention_config + + config = create_attention_config( + model_config=model.config, + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + ) self.tp_q_head_num = model.config["num_attention_heads"] // model.tp_world_size_ # Sparse FlashMLA uses 64-head tiles; TP can leave fewer local Q heads. self.padded_q_head_num = ((self.tp_q_head_num + 63) // 64) * 64 - self.compress_ratios = tuple(dict.fromkeys(model.config["compress_ratios"])) + self.compress_ratios = tuple(dict.fromkeys(config.compress_ratios)) + self.layer_compress_ratios = dict(zip(config.logical_layer_ids, config.compress_ratios)) + self.swa_config = config.create_swa_config() + self.compressed_configs = { + ratio: config.create_compressed_config(compress_ratio=ratio) for ratio in self.compress_ratios if ratio + } @classmethod def get_swa_prefill_workspace(cls, model): """由 backend 复用计算临时区；模型在 KV 容量评估前调用，以计入实际显存占用。""" - if model.args.run_mode == "decode": - return None, 0 - config = model.req_manager.cache_config - page_token_num = config.swa_page_token_num + from lightllm.models.deepseek_v4.config import create_attention_config + + config = create_attention_config( + model_config=model.config, + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + ).create_swa_config() + page_token_num = config.page_token_size max_forward_tokens = max(model.batch_max_tokens, model.max_req_num) + 2 * model.tp_world_size_ capacity = ( triton.cdiv( - max_forward_tokens + model.max_req_num * (config.sliding_window + 2 * page_token_num), page_token_num + max_forward_tokens + model.max_req_num * (config.window_size + 2 * page_token_num), page_token_num ) * page_token_num ) + if model.args.run_mode == "decode": + # D 节点初始化仍会对 HOLD 请求执行一次 prefill；只保留末尾的 HOLD 页， + # 不为普通请求预留长 prefill 临时区，实际请求的前缀由 P 节点传入。 + capacity = 0 regions = 1 + int(model.args.enable_prefill_microbatch_overlap) # 每个 microbatch 使用独立临时区；最后一页留给 HOLD 请求。 token_num = regions * capacity + page_token_num + # 物理行宽由 SWA 原子配置给出，backend 仅决定 workspace 的 token 容量。 + slot_bytes = config.get_packed_kv_bytes(page_token_num) // page_token_num + layer_count = len(config.logical_layer_ids) buffer = cls.get_gpu_workspace_buffer( - key_name="deepseek_v4_swa_prefill", - workspace_size=config.layer_num * token_num * config.swa_slot_bytes, + key_name=f"deepseek_v4_swa_prefill_{model.att_state_role}", + workspace_size=layer_count * token_num * slot_bytes, dtype=torch.uint8, ) - return buffer.view(config.layer_num, token_num, config.swa_slot_bytes), capacity - - def get_att_input_params(self, infer_state, layer_index): - from lightllm.models.deepseek_v4.triton_kernel.cache_io import view_swa_kv - - swa, compressed_kv = infer_state.req_manager.get_att_input_params(layer_index) + return buffer.view(layer_count, token_num, slot_bytes), capacity + + def get_att_input_params(self, infer_state, logical_layer_id: int): + state = infer_state.swa_state + swa_buffer = state.get_layer_view(logical_layer_id=logical_layer_id) + config = state.config + swa = swa_buffer.view(-1, config.get_packed_kv_bytes(config.page_token_size)) + ratio = self.layer_compress_ratios[logical_layer_id] + compressed_kv = infer_state.get_compressed_kv_view(logical_layer_id=logical_layer_id).kv if ratio else None if infer_state.is_prefill: - swa = view_swa_kv(infer_state.prefill_att_state.swa_buffer[layer_index], self.cache_config) + swa = infer_state.prefill_att_state.swa_buffer[config.logical_layer_ids.index(logical_layer_id)].view( + -1, + config.get_packed_kv_bytes(config.page_token_size), + ) return swa, compressed_kv def _flashmla_att( @@ -66,20 +96,12 @@ def _flashmla_att( ) -> torch.Tensor: """Attend to SWA plus optional compressed KV and return all padded Q heads.""" swa_packed, compressed_packed = packed_kv - config = self.cache_config - bytes_per_token = config.kv_bytes - # Expose FlashMLA's [page, token, KV head, byte] axes while retaining padded page strides. + # 原子配置提供计算视图；backend 不再计算编码字节数或解释压缩块布局。 compressed_cache = None if compressed_packed is not None: - page_size = config.compressed_kv_token_alignment // nsa_dict["compress_ratio"] - bytes_per_page = page_size * bytes_per_token - compressed_cache = compressed_packed[:, :bytes_per_page].view( - compressed_packed.shape[0], page_size, 1, bytes_per_token - ) - swa_bytes_per_page = config.swa_page_token_num * bytes_per_token - swa_cache = swa_packed[:, :swa_bytes_per_page].view( - swa_packed.shape[0], config.swa_page_token_num, 1, bytes_per_token - ) + config = self.compressed_configs[nsa_dict["compress_ratio"]] + compressed_cache = config.get_attention_view(packed_kv=compressed_packed) + swa_cache = self.swa_config.get_attention_view(packed_kv=swa_packed) padded_out, _ = flashmla.flash_mla_with_kvcache( q=q.unsqueeze(1), @@ -116,22 +138,22 @@ class _PrefillAttState(BasePrefillAttState): def prepare_swa_workspace(self, infer_state): from lightllm.models.deepseek_v4.triton_kernel.swa_cache import prepare_swa_prefill_workspace - requests = infer_state.req_manager + swa_state = infer_state.swa_state self.swa_buffer, capacity = self.backend.get_swa_prefill_workspace(self.backend.model) - self.HOLD_SWA_SLOT = self.swa_buffer.shape[1] - self.backend.cache_config.swa_page_token_num + self.HOLD_SWA_SLOT = self.swa_buffer.shape[1] - swa_state.config.page_token_size self.prefill_token_ranges = torch.zeros( - (requests.max_request_num + 1, 2), dtype=torch.int32, device=self.swa_buffer.device + (swa_state.request_capacity, 2), dtype=torch.int32, device=self.swa_buffer.device ) self.swa_prefill_workspace_offsets = torch.full( - (requests.max_request_num + 1,), -1, dtype=torch.int32, device=self.swa_buffer.device + (swa_state.request_capacity,), -1, dtype=torch.int32, device=self.swa_buffer.device ) prepare_swa_prefill_workspace( - requests, - self, - infer_state.b_req_idx, - infer_state.b_ready_cache_len, - infer_state.b_seq_len, - infer_state.microbatch_index * capacity, + swa_state=swa_state, + prefill_state=self, + b_req_idx=infer_state.b_req_idx, + b_ready_cache_len=infer_state.b_ready_cache_len, + b_seq_len=infer_state.b_seq_len, + workspace_base=infer_state.microbatch_index * capacity, ) def init_state(self): diff --git a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py index 29c1d0fb56..10ad62d6b8 100644 --- a/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py +++ b/lightllm/common/basemodel/attention/nsa/fp8_flashmla_sparse.py @@ -84,11 +84,16 @@ def _nsa_prefill_att( # 当前推理生成的token kv部分从 prefill_cache_kv 中获取，历史 # 部分kv 从 packed_kv 中获取, 并进行反量化，这样可以避免 prefill_cache_kv # 部分的数据进行重复的反量化操作，提升整体的性能。 - kv, topk_indices = self.infer_state.mem_manager.get_prefill_kv_cache_and_remap_indices( + from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( + get_prefill_kv_cache_and_remap_indices_triton, + ) + + kv, topk_indices = get_prefill_kv_cache_and_remap_indices_triton( packed_kv=packed_kv, - topk_indices=topk_mem_indices, + topk_mem_indices=topk_mem_indices, prefill_mem_index=self.infer_state.mem_index, prefill_cache_kv=prefill_cache_kv, + prefill_dtype=prefill_cache_kv.dtype, ) else: kv = prefill_cache_kv diff --git a/lightllm/common/basemodel/attention/triton/int4kv.py b/lightllm/common/basemodel/attention/triton/int4kv.py index e4dd2cee2f..22e723de3b 100644 --- a/lightllm/common/basemodel/attention/triton/int4kv.py +++ b/lightllm/common/basemodel/attention/triton/int4kv.py @@ -70,9 +70,6 @@ def _groupsize_quant_prefill_att( # o_tensor = alloc_func(q.shape, q.dtype, device=q.device) # batch_size = self.infer_state.b_seq_len.shape[0] - assert k.untyped_storage().data_ptr() == v.untyped_storage().data_ptr() - assert k_scale.untyped_storage().data_ptr() == v_scale.untyped_storage().data_ptr() - total_token_num = self.infer_state.total_token_num head_dim = k.shape[2] * 2 # 2个4bit存储为一个int8, 所以维度需要翻倍，才是解量化后的精度 k_dequant = alloc_func((total_token_num, k.shape[1], head_dim), dtype=q.dtype, device=q.device) diff --git a/lightllm/common/basemodel/attention/triton/int8kv.py b/lightllm/common/basemodel/attention/triton/int8kv.py index d863ac4014..9da88626c8 100644 --- a/lightllm/common/basemodel/attention/triton/int8kv.py +++ b/lightllm/common/basemodel/attention/triton/int8kv.py @@ -69,12 +69,7 @@ def _groupsize_quant_prefill_att( v_scale: torch.Tensor, alloc_func=torch.empty, ) -> torch.Tensor: - # o_tensor = alloc_func(q.shape, q.dtype, device=q.device) - # batch_size = self.infer_state.b_seq_len.shape[0] - - assert k.untyped_storage().data_ptr() == v.untyped_storage().data_ptr() - assert k_scale.untyped_storage().data_ptr() == v_scale.untyped_storage().data_ptr() - + # 反量化内核分别使用 K/V 及 scale 的指针和 stride，兼容共享或独立存储。 total_token_num = self.infer_state.total_token_num k_dequant = alloc_func((total_token_num, k.shape[1], k.shape[2]), dtype=q.dtype, device=q.device) v_dequant = alloc_func((total_token_num, v.shape[1], v.shape[2]), dtype=q.dtype, device=q.device) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 34a1c9c71c..78af939391 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -8,15 +8,13 @@ import torch import torch.nn.functional as F import triton -from typing import final, List +from typing import final, List, TYPE_CHECKING from tqdm import tqdm from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.kv_cache_mem_manager import MemoryManager -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class from lightllm.common.req_manager import ReqManager -from lightllm.common.build_utils import repair_config +from lightllm.common.attention_state import AttentionStateManager from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( select_kv_index_from_req, select_kv_index_from_req_prefill, @@ -27,12 +25,10 @@ from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed from lightllm.utils.log_utils import init_logger -from lightllm.utils.dist_utils import get_dp_world_size -from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory +from lightllm.utils.dist_utils import get_dp_world_size, get_current_rank_in_dp from lightllm.utils.envs_utils import ( get_env_start_args, get_llm_data_type, - get_added_mtp_kv_layer_num, ) from lightllm.distributed.communication_op import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput, PostLayerOutput @@ -41,7 +37,6 @@ ) from lightllm.common.basemodel.mtp_manager import MtpManager from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch -from lightllm.utils.envs_utils import set_model_init_status from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType from lightllm.utils.infer_utils import post_empty_cache from lightllm.utils.torch_memory_saver_utils import ( @@ -51,6 +46,12 @@ from .attention import get_prefill_att_backend_class, get_decode_att_backend_class from .attention import BaseAttBackend +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + logger = init_logger(__name__) torch.backends.cudnn.enabled = True @@ -70,6 +71,7 @@ class TpPartBaseModel: # infer state class infer_state_class = InferStateInfo + is_hybrid_att_model = False def __init__(self, kvargs): self.args = get_env_start_args() @@ -97,6 +99,7 @@ def __init__(self, kvargs): self.expert_dtype = kvargs.get("expert_dtype", None) self.mem_fraction = kvargs.get("mem_fraction", 0.9) self.tp_world_size_ = get_dp_world_size() + self.tp_rank_ = get_current_rank_in_dp() self.enable_tpsp_mix_mode = get_env_start_args().enable_tpsp_mix_mode self.graph_max_batch_size = self._align_decode_batch_size(self.graph_max_batch_size) @@ -109,49 +112,111 @@ def __init__(self, kvargs): self._init_quant() enable_weight_cpu_backup = self.args.enable_weight_cpu_backup - with profile_mtp_weight_memory(self): - with self.torch_memory_saver.region(tag=MemoryTag.WEIGHT, enable_cpu_backup=enable_weight_cpu_backup): - self._init_weights() - with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - self._init_req_manager() - 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._check_mem_size() + with self.torch_memory_saver.region(tag=MemoryTag.WEIGHT, enable_cpu_backup=enable_weight_cpu_backup): + self._init_weights() self._init_infer_layer() self._init_some_value() self._init_custom() self.load_weights(self.weight_dict) + # The backend calls each model's runtime stages after all weights are loaded. + self.graph = None + return + + @property + def att_state_role(self) -> str: + """当前模型的资源前缀；每个 draft 独立拥有 State。""" + from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS + + if not self.is_mtp_draft_model: + return "main" + draft_index = len(self.mtp_previous_draft_models) + assert draft_index < MAX_DRAFT_MODELS, "at most 36 draft models are supported" + return f"draft{draft_index}" + + def get_req_manager_max_seq_len(self) -> int: + """返回请求到 token 索引映射所需的长度容量，不表示模型支持的最大上下文长度。""" + return max(0, self.batch_max_tokens or 0, self.max_seq_length or 0) + + def bind_req_manager(self, req_manager: ReqManager) -> None: + """绑定启动流程创建的共享请求管理器；main/draft 共用请求槽位和 token 索引映射。""" + self.req_manager = req_manager + + def prepare_attention_states(self) -> None: + """main 创建共享管理器，draft 引用同一实例，再准备当前模型自己的 State。""" + assert getattr(self, "att_state_manager", None) is None, "resources already initialized" + if self.is_mtp_draft_model: + self.att_state_manager = self.main_model.att_state_manager + assert not self.att_state_manager._state_registration_finished, "State registration is already finished" + assert ( + self.max_req_num + 1 == self.att_state_manager._request_capacity + ), "models must share the request domain" + with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + if not self.is_mtp_draft_model: + self.att_state_manager = AttentionStateManager(request_capacity=self.max_req_num + 1) + # 无 attention 的 draft 只共享管理器，不创建自身 State。 + if self.is_mtp_draft_model and self.args.mtp_mode in ("vanilla_no_att", "eagle_no_att"): + return + + # 当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列；各工厂自行选择参与层及共享分组。 + logical_layer_ids = tuple(layer.layer_num_ for layer in self.layers_infer) + full_states = self.create_full_att_states( + model_config=self.config, + args=self.args, + model_logical_layer_ids=logical_layer_ids, + tp_rank=self.tp_rank_, + tp_world_size=self.tp_world_size_, + ) + bounded_states = self.create_bounded_att_states( + model_config=self.config, + args=self.args, + model_logical_layer_ids=logical_layer_ids, + request_capacity=self.max_req_num + 1, + tp_rank=self.tp_rank_, + tp_world_size=self.tp_world_size_, + device="cuda", + ) + for state in (*full_states, *bounded_states): + self.att_state_manager.register_state(self.att_state_role, state) + def init_runtime(self) -> None: + """Initialize this model's backends, warmup and graphs after shared allocation.""" + self.req_manager.bind_hold_page(self.att_state_manager.full_att_allocator.hold_token_indices) + self._check_mem_size() self._init_att_backend() self._init_att_backend1() - - logger.info(f"use prefill att backend: {self.prefill_att_backend.__class__.__name__}") - logger.info(f"use decode att backend: {self.decode_att_backend.__class__.__name__}") + self._init_hidden_collector() + role = self.att_state_role + logger.info( + f"{role} prefill/decode attention: {type(self.prefill_att_backend).__name__} / " + f"{type(self.decode_att_backend).__name__}" + ) if self.prefill_att_backend1 is not None: - logger.info(f"use prefill att backend1: {self.prefill_att_backend1.__class__.__name__}") - logger.info(f"use decode att backend1: {self.decode_att_backend1.__class__.__name__}") + logger.info( + f"{role} prefill/decode attention1: {type(self.prefill_att_backend1).__name__} / " + f"{type(self.decode_att_backend1).__name__}" + ) - self._init_hidden_collector() self._autotune_warmup() self._init_padded_req() self._init_cudagraph() self._init_prefill_cuda_graph() - self._check_max_len_infer() - torch.cuda.empty_cache() - set_model_init_status(True) - return + if not self.is_mtp_draft_model: + self._check_max_len_infer() + + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """返回当前模型的有效配置，供模型初始化和 CPU cache 共用。 + + 不依赖模型实例、权重或 CUDA；不修改输入配置，重复调用应保持结果不变。 + 子类负责模型专有的配置覆盖，通用工具仅展开 text_config 并同步字段别名。 + """ + from lightllm.common.basemodel.basemodel_utils import normalize_model_config + + return normalize_model_config(model_config) def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: - self.config = json.load(json_file) - # rename keys - repair_config(self.config, same_names=["num_attention_heads", "n_head"]) - repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) + self.config = self.normalize_model_config(json.load(json_file)) if self.finetune_config: self.config["vocab_size"] = self.finetune_config.vocab_size return @@ -196,20 +261,43 @@ def load_weights(self, weight_dict: dict): [weight.verify_load() for weight in self.trans_layers_weight] return - 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.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_attention_heads"] // self.tp_world_size_, - head_dim=self.config["n_embed"] // self.config["num_attention_heads"], - layer_num=self.config["n_layer"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """返回本模型的 full-attention 定义列表；默认没有此类资源，不分配 KV 容量。 + + model_logical_layer_ids 是当前模型拥有的全局逻辑层 ID，按模型内部层顺序排列； + 具体模型自行选择参与层并定义 logical_layer_groups，不要求使用全部层。 + """ + return [] + + @classmethod + def create_bounded_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", + ) -> list["BoundedAttentionState"]: + """返回本模型的 bounded 运行态列表；meta 用于布局查询，大小页后续分配。 + + model_logical_layer_ids 描述当前模型拥有的全部全局逻辑层，具体模型自行选择参与层。 + """ + return [] def _check_mem_size(self): - self.max_total_token_num = self.mem_manager.size + self.max_total_token_num = self.att_state_manager.full_att_allocator.size assert ( self.max_total_token_num > self.batch_max_tokens @@ -229,17 +317,6 @@ def _check_mem_size(self): return - 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 = ReqManager(self.max_req_num, create_max_seq_len, None) - return - def _init_infer_layer(self, start_layer_index=0): self.pre_infer = self.pre_layer_infer_class(network_config=self.config) self.post_infer = self.post_layer_infer_class(network_config=self.config) @@ -260,6 +337,11 @@ def _init_some_value(self): return def _init_att_backend(self): + role = self.att_state_role + if not self.att_state_manager.get_full_attention_states(role): + # FFN-only and linear-only models do not need a full-attention backend. + self.prefill_att_backend = self.decode_att_backend = None + return self.prefill_att_backend: BaseAttBackend = get_prefill_att_backend_class(index=0)(model=self) self.decode_att_backend: BaseAttBackend = get_decode_att_backend_class(index=0)(model=self) return @@ -352,6 +434,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.is_draft_model = self.is_mtp_draft_model + infer_state.att_state_role = self.att_state_role infer_state.return_all_prompt_logics = self.return_all_prompt_logics infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num @@ -374,8 +457,8 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.multimodal_params = model_input.multimodal_params - infer_state.mem_manager = self.mem_manager infer_state.req_manager = self.req_manager + infer_state.att_state_manager = self.att_state_manager infer_state.mem_index = self._select_mem_indexes(model_input) infer_state.microbatch_index = microbatch_index @@ -385,13 +468,21 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens if infer_state.is_prefill: - infer_state.prefill_att_state = self.prefill_att_backend.create_att_prefill_state(infer_state=infer_state) + infer_state.prefill_att_state = ( + self.prefill_att_backend.create_att_prefill_state(infer_state=infer_state) + if self.prefill_att_backend is not None + else None + ) if self.prefill_att_backend1 is not None: infer_state.prefill_att_state1 = self.prefill_att_backend1.create_att_prefill_state( infer_state=infer_state ) else: - infer_state.decode_att_state = self.decode_att_backend.create_att_decode_state(infer_state=infer_state) + infer_state.decode_att_state = ( + self.decode_att_backend.create_att_decode_state(infer_state=infer_state) + if self.decode_att_backend is not None + else None + ) if self.decode_att_backend1 is not None: infer_state.decode_att_state1 = self.decode_att_backend1.create_att_decode_state( infer_state=infer_state @@ -988,9 +1079,9 @@ def _check_max_len_infer(self): logger.info("begin check max_len infer") dummy_input_ids = torch.ones(self.batch_max_tokens, dtype=torch.int64, device="cuda") b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda") - page_size = self.mem_manager.page_size + page_size = self.att_state_manager.full_att_allocator.page_size alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size - mem_indexes = self.mem_manager.alloc(alloc_token_num).cuda() + mem_indexes = self.att_state_manager.full_att_allocator.alloc(alloc_token_num).cuda() self.req_manager.req_to_token_indexs[b_req_idx[0], : len(mem_indexes)] = mem_indexes b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = self.batch_max_tokens @@ -1024,7 +1115,7 @@ def _check_max_len_infer(self): torch.argmax(prob_out, dim=1, keepdim=True) prob_out = None self.req_manager.free_all() - self.mem_manager.free_all() + self.att_state_manager.full_att_allocator.free_all() logger.info(f"check max_len {self.batch_max_tokens} infer ok") except (RuntimeError, torch.OutOfMemoryError) as e: logger.exception(str(e)) @@ -1069,9 +1160,9 @@ def _autotune_warmup(self): 0, 10000, (input_len,), dtype=torch.int64, device="cuda", generator=rand_gen ) b_req_idx = torch.tensor([self.req_manager.alloc()], dtype=torch.int32, device="cuda") - page_size = self.mem_manager.page_size + page_size = self.att_state_manager.full_att_allocator.page_size alloc_token_num = triton.cdiv(len(dummy_input_ids), page_size) * page_size - mem_indexes = self.mem_manager.alloc(alloc_token_num).cuda() + mem_indexes = self.att_state_manager.full_att_allocator.alloc(alloc_token_num).cuda() self.req_manager.req_to_token_indexs[b_req_idx[0], : len(mem_indexes)] = mem_indexes b_seq_len = torch.ones(1, dtype=torch.int32, device="cuda") b_seq_len[:] = input_len @@ -1105,14 +1196,14 @@ def _autotune_warmup(self): ) del model_output self.req_manager.free_all() - self.mem_manager.free_all() + self.att_state_manager.full_att_allocator.free_all() gc.collect() torch.cuda.empty_cache() except Exception as e: logger.warning(f"autotune warmup for length {input_len} failed: {str(e)}") logger.exception(str(e)) self.req_manager.free_all() - self.mem_manager.free_all() + self.att_state_manager.full_att_allocator.free_all() gc.collect() torch.cuda.empty_cache() self.layers_num = layer_num_bak diff --git a/lightllm/common/basemodel/basemodel_utils.py b/lightllm/common/basemodel/basemodel_utils.py new file mode 100644 index 0000000000..2aa5d02d37 --- /dev/null +++ b/lightllm/common/basemodel/basemodel_utils.py @@ -0,0 +1,12 @@ +"""模型启动与 CPU cache 共用的配置展开及字段别名处理。""" + +from lightllm.common.build_utils import repair_config + + +def normalize_model_config(model_config: dict[str, object]) -> dict[str, object]: + """复制文本配置并补齐通用别名；模型专有覆盖由模型类在调用前完成。""" + config = model_config.get("text_config", model_config).copy() + repair_config(config, same_names=["num_attention_heads", "n_head"]) + repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(config, same_names=["num_hidden_layers", "n_layer"]) + return config diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index ee888c1a00..ab0e0c0721 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -292,7 +292,7 @@ def warmup(self, model): del b_req_idx del b_seq_len - model.mem_manager.free_all() + model.att_state_manager.full_att_allocator.free_all() model.req_manager.free_all() # release local tensors for var_name, var_value in list(locals().items()): @@ -354,7 +354,7 @@ def warmup_overlap(self, model): _, _ = model.microbatch_overlap_decode(decode_batches[0], decode_batches[1]) - model.mem_manager.free_all() + model.att_state_manager.full_att_allocator.free_all() model.req_manager.free_all() del decode_batches diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index bf25149450..83b312f991 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -1,7 +1,7 @@ import torch import triton import collections -from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.attention_state.manager import AttentionStateManager from lightllm.common.req_manager import ReqManager from lightllm.distributed import CustomProcessGroup from typing import Tuple, Any, Optional, List @@ -47,7 +47,8 @@ def __init__(self): self.is_prefill: bool = None self.is_draft_model: bool = False - self.mem_manager: MemoryManager = None + self.att_state_manager: AttentionStateManager = None + self.att_state_role: str = "main" # main 或 draft0 .. draft35，区分共享管理器中的 State。 self.req_manager: ReqManager = None self.mem_index: torch.Tensor = None @@ -134,11 +135,13 @@ def init_some_extra_state(self, model): def init_att_state(self): if self.is_prefill: - self.prefill_att_state.init_state() + if self.prefill_att_state is not None: + self.prefill_att_state.init_state() if self.prefill_att_state1 is not None: self.prefill_att_state1.init_state() else: - self.decode_att_state.init_state() + if self.decode_att_state is not None: + self.decode_att_state.init_state() if self.decode_att_state1 is not None: self.decode_att_state1.init_state() self.check_att_states() @@ -178,7 +181,8 @@ def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"): if attr_ is not None and attr_.data_ptr() != attr_value.data_ptr(): attr_.copy_(attr_value, non_blocking=True) - self.decode_att_state.copy_for_decode_cuda_graph(new_infer_state.decode_att_state) + if self.decode_att_state is not None: + self.decode_att_state.copy_for_decode_cuda_graph(new_infer_state.decode_att_state) if self.decode_att_state1 is not None: self.decode_att_state1.copy_for_decode_cuda_graph(new_infer_state.decode_att_state1) return diff --git a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py index f0cc129c09..b65cab258a 100755 --- a/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py +++ b/lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.full_att import FullAttentionState import os import torch import torch.distributed as dist @@ -33,10 +34,15 @@ def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight) -> Tuple[to raise Exception("need to impl") def _post_cache_kv(self, cache_kv, infer_state: InferStateInfo, layer_weight): - mem_manager = infer_state.mem_manager - mem_manager.operator.copy_kv_to_mem_manager( - layer_index=self.layer_num_, - mem_index=infer_state.mem_index, + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + state.write_full_att_kv( + logical_layer_id=self.layer_num_, + token_indices=infer_state.mem_index, kv=cache_kv, ) return diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index 28d44fefbf..4f21d90cfc 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -226,7 +226,7 @@ def warmup(self, model): del b_req_idx del b_seq_len - model.mem_manager.free_all() + model.att_state_manager.full_att_allocator.free_all() model.req_manager.free_all() # release local tensors for var_name, var_value in list(locals().items()): @@ -289,7 +289,7 @@ def warmup_overlap(self, model): _, _ = model.microbatch_overlap_prefill(prefill_batches[0], prefill_batches[1]) - model.mem_manager.free_all() + model.att_state_manager.full_att_allocator.free_all() model.req_manager.free_all() del prefill_batches diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py index be5eb6d8a9..b94126fcaa 100644 --- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py +++ b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv_fp8.py @@ -9,6 +9,7 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( KV, Dest_loc, Out, + Out_v, scale, stride_kv_bs, stride_kv_h, @@ -22,6 +23,7 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( FP8_MIN: tl.constexpr, FP8_MAX: tl.constexpr, IS_PER_TENSOR_QUANT: tl.constexpr, + SPLIT_KV: tl.constexpr, ): cur_index = tl.program_id(0) offs_h = tl.arange(0, BLOCK_HEAD) @@ -30,7 +32,18 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( dest_index = tl.load(Dest_loc + cur_index).to(tl.int64) kv_ptrs = KV + cur_index * stride_kv_bs + stride_kv_h * offs_h[:, None] + stride_kv_d * offs_d[None, :] - o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] + if SPLIT_KV: + # 两块独立分配的 K/V 具有相同形状和 stride，输入仍按 K heads、V heads 排列。 + k_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] + v_ptrs = ( + Out_v + + dest_index * stride_o_bs + + stride_o_h * (offs_h[:, None] - head_num // 2) + + stride_o_d * offs_d[None, :] + ) + o_ptrs = tl.where(offs_h[:, None] < head_num // 2, k_ptrs, v_ptrs) + else: + o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :] # to fp8 if not IS_PER_TENSOR_QUANT: @@ -50,15 +63,20 @@ def _fwd_kernel_destindex_copy_kv_per_head_fp8( @torch.no_grad() -def destindex_copy_kv_fp8(KV, DestLoc, scales, Out, is_per_tensor_quant=False): +def destindex_copy_kv_fp8(KV, DestLoc, scales, Out, is_per_tensor_quant=False, Out_v=None): """ - 当 is_per_tensor_quant 为 False 时，为 为 per_head 量化。 + 当 is_per_tensor_quant 为 False 时，为 per_head 量化。 + Out_v 非空时，Out/Out_v 分别为 K/V 存储；否则沿用合并输出布局。 """ seq_len = DestLoc.shape[0] head_num = KV.shape[1] assert head_num % 2 == 0 head_dim = KV.shape[2] - assert KV.shape[1] == Out.shape[1] and KV.shape[2] == Out.shape[2] + if Out_v is None: + assert KV.shape[1] == Out.shape[1] and KV.shape[2] == Out.shape[2] + else: + assert KV.shape[1] == 2 * Out.shape[1] and KV.shape[2] == Out.shape[2] + assert Out_v.shape == Out.shape and Out_v.stride() == Out.stride() BLOCK_HEAD = triton.next_power_of_2(head_num) grid = (seq_len,) num_warps = 1 @@ -67,6 +85,7 @@ def destindex_copy_kv_fp8(KV, DestLoc, scales, Out, is_per_tensor_quant=False): KV, DestLoc, Out, + Out if Out_v is None else Out_v, scales, KV.stride(0), KV.stride(1), @@ -80,6 +99,7 @@ def destindex_copy_kv_fp8(KV, DestLoc, scales, Out, is_per_tensor_quant=False): FP8_MIN=torch.finfo(torch.float8_e4m3fn).min, FP8_MAX=torch.finfo(torch.float8_e4m3fn).max, IS_PER_TENSOR_QUANT=is_per_tensor_quant, + SPLIT_KV=Out_v is not None, num_warps=num_warps, num_stages=1, ) 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 82% 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 index d5859c6867..c7e9da2bdf 100644 --- 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 @@ -5,33 +5,33 @@ @triton.jit -def _quant_scale(max_nope, fp8_max): +def _quant_scale(max_nope: tl.tensor, fp8_max: float): return tl.exp2(tl.ceil(tl.log2(tl.maximum(max_nope / fp8_max, 1e-4)))) @triton.jit def _fwd_kernel_destindex_copy_kv_flashmla_fp8( - KV_nope, - KV_rope, - Dest_loc, - O_nope, - O_scale, - O_rope, - stride_kv_nope_bs, - stride_kv_nope_h, - stride_kv_nope_d, - stride_kv_rope_bs, - stride_kv_rope_h, - stride_kv_rope_d, - stride_o_nope_bs, - stride_o_nope_h, - stride_o_nope_d, - stride_o_scale_bs, - stride_o_scale_h, - stride_o_scale_d, - stride_o_rope_bs, - stride_o_rope_h, - stride_o_rope_d, + KV_nope: tl.tensor, + KV_rope: tl.tensor, + Dest_loc: tl.tensor, + O_nope: tl.tensor, + O_scale: tl.tensor, + O_rope: tl.tensor, + stride_kv_nope_bs: tl.tensor, + stride_kv_nope_h: tl.tensor, + stride_kv_nope_d: tl.tensor, + stride_kv_rope_bs: tl.tensor, + stride_kv_rope_h: tl.tensor, + stride_kv_rope_d: tl.tensor, + stride_o_nope_bs: tl.tensor, + stride_o_nope_h: tl.tensor, + stride_o_nope_d: tl.tensor, + stride_o_scale_bs: tl.tensor, + stride_o_scale_h: tl.tensor, + stride_o_scale_d: tl.tensor, + stride_o_rope_bs: tl.tensor, + stride_o_rope_h: tl.tensor, + stride_o_rope_d: tl.tensor, FP8_MIN: tl.constexpr, FP8_MAX: tl.constexpr, BLOCK_DMODEL_NOPE: tl.constexpr, 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 87% 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 index 0d0447fec8..7d5995db40 100644 --- 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 @@ -1,3 +1,5 @@ +"""Compact packed FP8 MLA history and newly computed prefill KV for attention.""" + import torch import triton import triton.language as tl @@ -5,9 +7,9 @@ @triton.jit def _build_prefill_row_table_kernel( - prefill_mem_index_ptr, - row_table_ptr, - prefill_token_num, + prefill_mem_index_ptr: tl.tensor, + row_table_ptr: tl.tensor, + prefill_token_num: tl.tensor, ): pid = tl.program_id(0) if pid < prefill_token_num: @@ -17,24 +19,24 @@ def _build_prefill_row_table_kernel( @triton.jit def _fill_compact_kv_kernel( - packed_nope_ptr, - packed_scale_ptr, - packed_rope_ptr, - unique_mem_index_ptr, - prefill_row_table_ptr, - prefill_kv_ptr, - compact_kv_ptr, - packed_nope_stride_s, - packed_nope_stride_d, - packed_scale_stride_s, - packed_scale_stride_d, - packed_rope_stride_s, - packed_rope_stride_d, - prefill_kv_stride_s, - prefill_kv_stride_d, - compact_kv_stride_s, - compact_kv_stride_d, - unique_num, + packed_nope_ptr: tl.tensor, + packed_scale_ptr: tl.tensor, + packed_rope_ptr: tl.tensor, + unique_mem_index_ptr: tl.tensor, + prefill_row_table_ptr: tl.tensor, + prefill_kv_ptr: tl.tensor, + compact_kv_ptr: tl.tensor, + packed_nope_stride_s: tl.tensor, + packed_nope_stride_d: tl.tensor, + packed_scale_stride_s: tl.tensor, + packed_scale_stride_d: tl.tensor, + packed_rope_stride_s: tl.tensor, + packed_rope_stride_d: tl.tensor, + prefill_kv_stride_s: tl.tensor, + prefill_kv_stride_d: tl.tensor, + compact_kv_stride_s: tl.tensor, + compact_kv_stride_d: tl.tensor, + unique_num: tl.tensor, KV_NOPE_DIM: tl.constexpr, KV_ROPE_DIM: tl.constexpr, GROUP_SIZE: tl.constexpr, diff --git a/lightllm/common/basemodel/triton_kernel/kv_move.py b/lightllm/common/basemodel/triton_kernel/kv_move.py index d2ce0764b7..ce99a007bd 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_move.py +++ b/lightllm/common/basemodel/triton_kernel/kv_move.py @@ -8,6 +8,8 @@ def _copy_kv_buffer_to_kv_buffer( mem_num, src_mem_index, dst_mem_index, + src_index_stride: tl.constexpr, + dst_index_stride: tl.constexpr, kv_buffer, kv_buffer_stride_l, kv_buffer_stride_s, @@ -24,8 +26,8 @@ def _copy_kv_buffer_to_kv_buffer( kv_buffer_stride_d = tl.cast(kv_buffer_stride_d, dtype=tl.int64) for i in range(start_index, mem_num, grid_num): - src_mem = tl.load(src_mem_index + i) - dst_mem = tl.load(dst_mem_index + i) + src_mem = tl.load(src_mem_index + i * src_index_stride) + dst_mem = tl.load(dst_mem_index + i * dst_index_stride) for j in range(tl.cdiv(kv_buffer_tail_dim, BLOCK)): offs = j * BLOCK + tl.arange(0, BLOCK) mask = offs < kv_buffer_tail_dim @@ -58,6 +60,8 @@ def copy_kv_buffer_to_kv_buffer( mem_num=len(src_mem_index), src_mem_index=src_mem_index, dst_mem_index=dst_mem_index, + src_index_stride=src_mem_index.stride(0), + dst_index_stride=dst_mem_index.stride(0), kv_buffer=kv_buffer, kv_buffer_stride_l=kv_buffer.stride(0), kv_buffer_stride_s=kv_buffer.stride(1), diff --git a/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py b/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py deleted file mode 100644 index b0c9a60649..0000000000 --- a/lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py +++ /dev/null @@ -1,531 +0,0 @@ -import torch -import triton -import triton.language as tl -from lightllm.common.state_cache_manager import LinearAttCacheConfig - - -@triton.jit -def _copy_kv_buffer_to_cpu_cache( - page_num, - mem_indexes_ptr, # [move_token_num] - page_indexes_ptr, # [page_num], - page_readies_ptr, # [page_num], - big_page_buffer_ids, # [page_num] - cpu_cache_full_att, # [all_page_num, head, xdim] - cpu_cache_full_att_stride_p, - cpu_cache_full_att_stride_h, - cpu_cache_full_att_stride_d, - cpu_cache_conv, # [all_page_num, tp_world_size, xxdim] - cpu_cache_conv_stride_p, - cpu_cache_conv_stride_t, - cpu_cache_conv_stride_d, - cpu_cache_ssm, # [all_page_num, tp_world_size, xxxdim] - cpu_cache_ssm_stride_p, - cpu_cache_ssm_stride_t, - cpu_cache_ssm_stride_d, - gpu_kv_full_att_state, # [token_size, full_att_layer_num, xdim] - gpu_kv_full_att_stride_s, - gpu_kv_full_att_stride_l, - gpu_kv_full_att_stride_d, - cpu_kv_conv_state, # [buffer_count, xxxxxdim] - cpu_kv_conv_stride_s, - cpu_kv_conv_stride_d, - cpu_kv_ssm_state, # [buffer_count, xxxxxxxdim] - cpu_kv_ssm_stride_s, - cpu_kv_ssm_stride_d, - gpu_full_att_tail_dim, - cpu_kv_conv_tail_dim, - cpu_kv_ssm_tail_dim, - tp_rank, - full_att_layer_num, - big_page_token_num, - head_scale_size, - BLOCK: tl.constexpr, -): - split_index_start = tl.program_id(0) - grid_num = tl.num_programs(0) - # 将 所有stride 切成 tl.int64 - cpu_cache_full_att_stride_p = tl.cast(cpu_cache_full_att_stride_p, tl.int64) - cpu_cache_full_att_stride_h = tl.cast(cpu_cache_full_att_stride_h, tl.int64) - cpu_cache_full_att_stride_d = tl.cast(cpu_cache_full_att_stride_d, tl.int64) - cpu_cache_conv_stride_p = tl.cast(cpu_cache_conv_stride_p, tl.int64) - cpu_cache_conv_stride_t = tl.cast(cpu_cache_conv_stride_t, tl.int64) - cpu_cache_conv_stride_d = tl.cast(cpu_cache_conv_stride_d, tl.int64) - cpu_cache_ssm_stride_p = tl.cast(cpu_cache_ssm_stride_p, tl.int64) - cpu_cache_ssm_stride_t = tl.cast(cpu_cache_ssm_stride_t, tl.int64) - cpu_cache_ssm_stride_d = tl.cast(cpu_cache_ssm_stride_d, tl.int64) - gpu_kv_full_att_stride_s = tl.cast(gpu_kv_full_att_stride_s, tl.int64) - gpu_kv_full_att_stride_l = tl.cast(gpu_kv_full_att_stride_l, tl.int64) - gpu_kv_full_att_stride_d = tl.cast(gpu_kv_full_att_stride_d, tl.int64) - cpu_kv_conv_stride_s = tl.cast(cpu_kv_conv_stride_s, tl.int64) - cpu_kv_conv_stride_d = tl.cast(cpu_kv_conv_stride_d, tl.int64) - cpu_kv_ssm_stride_s = tl.cast(cpu_kv_ssm_stride_s, tl.int64) - cpu_kv_ssm_stride_d = tl.cast(cpu_kv_ssm_stride_d, tl.int64) - - for block_index in range(page_num): - cpu_page_index = tl.load(page_indexes_ptr + block_index).to(tl.int64) - run_flag = 1 - if cpu_page_index == -1: - run_flag = 0 - ready_state = tl.load(page_readies_ptr + block_index) - if ready_state: - run_flag = 0 - if tp_rank % head_scale_size == 0: - head_flag = 1 - else: - head_flag = 0 - - mem_start_ptr = mem_indexes_ptr + big_page_token_num * block_index - for i in range(split_index_start, tl.cdiv(gpu_full_att_tail_dim, BLOCK) * run_flag * head_flag, grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < gpu_full_att_tail_dim - per_token_size = gpu_full_att_tail_dim // big_page_token_num - per_layer_size = per_token_size // full_att_layer_num - mem_offs = gpu_start_i // (per_token_size) - mem_index = tl.load(mem_start_ptr + mem_offs, mask=mask, other=-1) - layer_index = (gpu_start_i // (per_layer_size)) % full_att_layer_num - dim_index = gpu_start_i % per_layer_size - gpu_full_att_data = tl.load( - gpu_kv_full_att_state - + mem_index * gpu_kv_full_att_stride_s - + layer_index * gpu_kv_full_att_stride_l - + dim_index * gpu_kv_full_att_stride_d, - mask=mask & (mem_index != -1), - other=0, - ) - dest_cpu_cache_full_att_ptr = ( - cpu_cache_full_att - + cpu_page_index * cpu_cache_full_att_stride_p - + (tp_rank // head_scale_size) * cpu_cache_full_att_stride_h - + gpu_start_i - ) - tl.store(dest_cpu_cache_full_att_ptr, gpu_full_att_data, mask=mask & (mem_index != -1)) - - big_page_idx = tl.load(big_page_buffer_ids + block_index) - - for i in range(split_index_start, tl.cdiv(cpu_kv_conv_tail_dim, BLOCK) * run_flag, grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < cpu_kv_conv_tail_dim - cpu_kv_conv_data = tl.load( - cpu_kv_conv_state + big_page_idx * cpu_kv_conv_stride_s + gpu_start_i, - mask=mask, - other=0, - ) - dest_cpu_cache_conv_ptr = ( - cpu_cache_conv - + cpu_page_index * cpu_cache_conv_stride_p - + tp_rank * cpu_cache_conv_stride_t - + gpu_start_i - ) - tl.store(dest_cpu_cache_conv_ptr, cpu_kv_conv_data, mask=mask) - - for i in range(split_index_start, tl.cdiv(cpu_kv_ssm_tail_dim, BLOCK) * run_flag, grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < cpu_kv_ssm_tail_dim - - cpu_kv_ssm_data = tl.load( - cpu_kv_ssm_state + big_page_idx * cpu_kv_ssm_stride_s + gpu_start_i, - mask=mask, - other=0, - ) - dest_cpu_cache_ssm_ptr = ( - cpu_cache_ssm + cpu_page_index * cpu_cache_ssm_stride_p + tp_rank * cpu_cache_ssm_stride_t + gpu_start_i - ) - tl.store(dest_cpu_cache_ssm_ptr, cpu_kv_ssm_data, mask=mask) - - return - - -def copy_kv_buffer_to_cpu_cache( - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - big_page_buffer_ids: torch.Tensor, - gpu_kv_full_att_state: torch.Tensor, # [full_att_layer_num, s, head_num, head_dim] - cpu_kv_conv_state: torch.Tensor, # [s, linear_layer_num, dim] - cpu_kv_ssm_state: torch.Tensor, # [s, linear_layer_num, xdim] - cpu_cache_tensor: torch.Tensor, # [page_num, 1, 1, 1, xxdim] - tp_rank: int, - tp_world_size: int, - big_page_token_num: int, - linear_config: LinearAttCacheConfig, - grid_num: int = 12, -): - assert len(page_indexes) == len(page_readies) == len(big_page_buffer_ids) - assert len(mem_indexes) % len(page_indexes) == 0 - - BLOCK = 4096 - if linear_config.full_att_all_num_kv_heads % tp_world_size == 0: - # tp world size 不比 kv 的 head 多时 - head_scale_size = 1 - else: - head_scale_size = tp_world_size // linear_config.full_att_all_num_kv_heads - - cpu_page_num = cpu_cache_tensor.shape[0] - cpu_cache_tensor = cpu_cache_tensor.view(cpu_page_num, -1).view(dtype=torch.uint8) - a = linear_config.get_cpu_cache_full_att_bytes() - b = linear_config.get_cpu_cache_conv_bytes() - c = linear_config.get_cpu_cache_ssm_bytes() - - if head_scale_size == 1: - cpu_cache_full_att = cpu_cache_tensor[:, 0:a].view(cpu_page_num, tp_world_size, -1) - else: - cpu_cache_full_att = cpu_cache_tensor[:, 0:a].view(cpu_page_num, linear_config.full_att_all_num_kv_heads, -1) - - cpu_cache_full_att = cpu_cache_full_att.view(dtype=torch.uint64) - - cpu_cache_conv = cpu_cache_tensor[:, a : (a + b)].view(cpu_page_num, tp_world_size, -1).view(dtype=torch.uint64) - cpu_cache_ssm = ( - cpu_cache_tensor[:, (a + b) : (a + b + c)].view(cpu_page_num, tp_world_size, -1).view(dtype=torch.uint64) - ) - - gpu_kv_full_att_state = gpu_kv_full_att_state.view( - gpu_kv_full_att_state.shape[0], gpu_kv_full_att_state.shape[1], -1 - ).view(dtype=torch.uint64) - - gpu_kv_full_att_state = gpu_kv_full_att_state.permute(1, 0, 2) # [s, layer_num, xxdim] - - cpu_kv_conv_state = cpu_kv_conv_state.view(cpu_kv_conv_state.shape[0], -1).view(dtype=torch.uint64) - cpu_kv_ssm_state = cpu_kv_ssm_state.view(cpu_kv_ssm_state.shape[0], -1).view(dtype=torch.uint64) - - gpu_full_att_tail_dim = gpu_kv_full_att_state.shape[-1] * gpu_kv_full_att_state.shape[-2] * big_page_token_num - cpu_kv_conv_tail_dim = cpu_kv_conv_state.shape[-1] - cpu_kv_ssm_tail_dim = cpu_kv_ssm_state.shape[-1] - full_att_layer_num = gpu_kv_full_att_state.shape[-2] - - assert full_att_layer_num == linear_config.get_full_att_kv_layer_num_with_draft_model() - assert gpu_full_att_tail_dim == cpu_cache_full_att.shape[-1] - assert cpu_cache_conv.shape[-1] == cpu_kv_conv_state.shape[-1] - assert cpu_cache_ssm.shape[-1] == cpu_kv_ssm_state.shape[-1] - assert gpu_kv_full_att_state.stride(2) == 1 - assert ( - gpu_full_att_tail_dim % big_page_token_num == 0 - and (gpu_full_att_tail_dim // big_page_token_num) % full_att_layer_num == 0 - ) - assert (tp_rank // head_scale_size) < linear_config.full_att_all_num_kv_heads - - grid = (grid_num,) - _copy_kv_buffer_to_cpu_cache[grid]( - page_num=len(page_indexes), - mem_indexes_ptr=mem_indexes, - page_indexes_ptr=page_indexes, - page_readies_ptr=page_readies, - big_page_buffer_ids=big_page_buffer_ids, - cpu_cache_full_att=cpu_cache_full_att, - cpu_cache_full_att_stride_p=cpu_cache_full_att.stride(0), - cpu_cache_full_att_stride_h=cpu_cache_full_att.stride(1), - cpu_cache_full_att_stride_d=cpu_cache_full_att.stride(2), - cpu_cache_conv=cpu_cache_conv, - cpu_cache_conv_stride_p=cpu_cache_conv.stride(0), - cpu_cache_conv_stride_t=cpu_cache_conv.stride(1), - cpu_cache_conv_stride_d=cpu_cache_conv.stride(2), - cpu_cache_ssm=cpu_cache_ssm, - cpu_cache_ssm_stride_p=cpu_cache_ssm.stride(0), - cpu_cache_ssm_stride_t=cpu_cache_ssm.stride(1), - cpu_cache_ssm_stride_d=cpu_cache_ssm.stride(2), - gpu_kv_full_att_state=gpu_kv_full_att_state, - gpu_kv_full_att_stride_s=gpu_kv_full_att_state.stride(0), - gpu_kv_full_att_stride_l=gpu_kv_full_att_state.stride(1), - gpu_kv_full_att_stride_d=gpu_kv_full_att_state.stride(2), - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_conv_stride_s=cpu_kv_conv_state.stride(0), - cpu_kv_conv_stride_d=cpu_kv_conv_state.stride(1), - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_kv_ssm_stride_s=cpu_kv_ssm_state.stride(0), - cpu_kv_ssm_stride_d=cpu_kv_ssm_state.stride(1), - gpu_full_att_tail_dim=gpu_full_att_tail_dim, - cpu_kv_conv_tail_dim=cpu_kv_conv_tail_dim, - cpu_kv_ssm_tail_dim=cpu_kv_ssm_tail_dim, - tp_rank=tp_rank, - full_att_layer_num=full_att_layer_num, - big_page_token_num=big_page_token_num, - head_scale_size=head_scale_size, - BLOCK=BLOCK, - ) - - -@triton.jit -def _copy_cpu_cache_to_kv_buffer( - page_num, - mem_indexes_ptr, # [move_token_num] - page_indexes_ptr, # [page_num], - big_page_buffer_ids, # [page_num] - cpu_cache_full_att, # [all_page_num, head, xdim] - cpu_cache_full_att_stride_p, - cpu_cache_full_att_stride_h, - cpu_cache_full_att_stride_d, - cpu_cache_conv, # [all_page_num, tp_world_size, xxdim] - cpu_cache_conv_stride_p, - cpu_cache_conv_stride_t, - cpu_cache_conv_stride_d, - cpu_cache_ssm, # [all_page_num, tp_world_size, xxxdim] - cpu_cache_ssm_stride_p, - cpu_cache_ssm_stride_t, - cpu_cache_ssm_stride_d, - gpu_kv_full_att_state, # [token_size, full_att_layer_num, xdim] - gpu_kv_full_att_stride_s, - gpu_kv_full_att_stride_l, - gpu_kv_full_att_stride_d, - cpu_kv_conv_state, # [buffer_count, xxxxxdim] - cpu_kv_conv_stride_s, - cpu_kv_conv_stride_d, - cpu_kv_ssm_state, # [buffer_count, xxxxxxxdim] - cpu_kv_ssm_stride_s, - cpu_kv_ssm_stride_d, - gpu_full_att_tail_dim, - cpu_kv_conv_tail_dim, - cpu_kv_ssm_tail_dim, - tp_rank, - full_att_layer_num, - big_page_token_num, - head_scale_size, - BLOCK: tl.constexpr, -): - split_index_start = tl.program_id(0) - grid_num = tl.num_programs(0) - # 将 所有stride 切成 tl.int64 - cpu_cache_full_att_stride_p = tl.cast(cpu_cache_full_att_stride_p, tl.int64) - cpu_cache_full_att_stride_h = tl.cast(cpu_cache_full_att_stride_h, tl.int64) - cpu_cache_full_att_stride_d = tl.cast(cpu_cache_full_att_stride_d, tl.int64) - cpu_cache_conv_stride_p = tl.cast(cpu_cache_conv_stride_p, tl.int64) - cpu_cache_conv_stride_t = tl.cast(cpu_cache_conv_stride_t, tl.int64) - cpu_cache_conv_stride_d = tl.cast(cpu_cache_conv_stride_d, tl.int64) - cpu_cache_ssm_stride_p = tl.cast(cpu_cache_ssm_stride_p, tl.int64) - cpu_cache_ssm_stride_t = tl.cast(cpu_cache_ssm_stride_t, tl.int64) - cpu_cache_ssm_stride_d = tl.cast(cpu_cache_ssm_stride_d, tl.int64) - gpu_kv_full_att_stride_s = tl.cast(gpu_kv_full_att_stride_s, tl.int64) - gpu_kv_full_att_stride_l = tl.cast(gpu_kv_full_att_stride_l, tl.int64) - gpu_kv_full_att_stride_d = tl.cast(gpu_kv_full_att_stride_d, tl.int64) - cpu_kv_conv_stride_s = tl.cast(cpu_kv_conv_stride_s, tl.int64) - cpu_kv_conv_stride_d = tl.cast(cpu_kv_conv_stride_d, tl.int64) - cpu_kv_ssm_stride_s = tl.cast(cpu_kv_ssm_stride_s, tl.int64) - cpu_kv_ssm_stride_d = tl.cast(cpu_kv_ssm_stride_d, tl.int64) - - for block_index in range(page_num): - cpu_page_index = tl.load(page_indexes_ptr + block_index).to(tl.int64) - - mem_start_ptr = mem_indexes_ptr + big_page_token_num * block_index - for i in range(split_index_start, tl.cdiv(gpu_full_att_tail_dim, BLOCK), grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < gpu_full_att_tail_dim - per_token_size = gpu_full_att_tail_dim // big_page_token_num - per_layer_size = per_token_size // full_att_layer_num - mem_offs = gpu_start_i // (per_token_size) - mem_index = tl.load(mem_start_ptr + mem_offs, mask=mask, other=-1) - layer_index = (gpu_start_i // (per_layer_size)) % full_att_layer_num - dim_index = gpu_start_i % per_layer_size - - src_cpu_cache_full_att_ptr = ( - cpu_cache_full_att - + cpu_page_index * cpu_cache_full_att_stride_p - + (tp_rank // head_scale_size) * cpu_cache_full_att_stride_h - + gpu_start_i - ) - cpu_full_att_data = tl.load(src_cpu_cache_full_att_ptr, mask=mask & (mem_index != -1), other=0) - - tl.store( - gpu_kv_full_att_state - + mem_index * gpu_kv_full_att_stride_s - + layer_index * gpu_kv_full_att_stride_l - + dim_index * gpu_kv_full_att_stride_d, - cpu_full_att_data, - mask=mask & (mem_index != -1), - ) - - big_page_idx = tl.load(big_page_buffer_ids + block_index) - - for i in range(split_index_start, tl.cdiv(cpu_kv_conv_tail_dim, BLOCK), grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < cpu_kv_conv_tail_dim - - src_cpu_cache_conv_ptr = ( - cpu_cache_conv - + cpu_page_index * cpu_cache_conv_stride_p - + tp_rank * cpu_cache_conv_stride_t - + gpu_start_i - ) - cpu_kv_conv_data = tl.load(src_cpu_cache_conv_ptr, mask=mask, other=0) - - tl.store( - cpu_kv_conv_state + big_page_idx * cpu_kv_conv_stride_s + gpu_start_i, - cpu_kv_conv_data, - mask=mask, - ) - - for i in range(split_index_start, tl.cdiv(cpu_kv_ssm_tail_dim, BLOCK), grid_num): - gpu_start_i = i * BLOCK + tl.arange(0, BLOCK) - mask = gpu_start_i < cpu_kv_ssm_tail_dim - - src_cpu_cache_ssm_ptr = ( - cpu_cache_ssm + cpu_page_index * cpu_cache_ssm_stride_p + tp_rank * cpu_cache_ssm_stride_t + gpu_start_i - ) - cpu_kv_ssm_data = tl.load(src_cpu_cache_ssm_ptr, mask=mask, other=0) - - tl.store( - cpu_kv_ssm_state + big_page_idx * cpu_kv_ssm_stride_s + gpu_start_i, - cpu_kv_ssm_data, - mask=mask, - ) - - return - - -def copy_cpu_cache_to_kv_buffer( - mem_indexes: torch.Tensor, - big_page_buffer_ids: torch.Tensor, - page_indexes: torch.Tensor, - gpu_full_att_kv_state: torch.Tensor, # [layer_num, s, head_num, head_dim] - cpu_kv_conv_state: torch.Tensor, # [layer_num, s, dim] - cpu_kv_ssm_state: torch.Tensor, # [layer_num, s, xdim] - cpu_cache_tensor: torch.Tensor, # [page_num, 1, 1, tp_world_size, xxdim] - tp_rank: int, - tp_world_size: int, - big_page_token_num: int, - linear_config: LinearAttCacheConfig, - grid_num: int = 12, -): - assert len(mem_indexes) % len(page_indexes) == 0 - - BLOCK = 4096 - if linear_config.full_att_all_num_kv_heads % tp_world_size == 0: - head_scale_size = 1 - else: - head_scale_size = tp_world_size // linear_config.full_att_all_num_kv_heads - - cpu_page_num = cpu_cache_tensor.shape[0] - cpu_cache_tensor = cpu_cache_tensor.view(cpu_page_num, -1).view(dtype=torch.uint8) - a = linear_config.get_cpu_cache_full_att_bytes() - b = linear_config.get_cpu_cache_conv_bytes() - c = linear_config.get_cpu_cache_ssm_bytes() - - if head_scale_size == 1: - cpu_cache_full_att = cpu_cache_tensor[:, 0:a].view(cpu_page_num, tp_world_size, -1) - else: - cpu_cache_full_att = cpu_cache_tensor[:, 0:a].view(cpu_page_num, linear_config.full_att_all_num_kv_heads, -1) - - cpu_cache_full_att = cpu_cache_full_att.view(dtype=torch.uint64) - - cpu_cache_conv = cpu_cache_tensor[:, a : (a + b)].view(cpu_page_num, tp_world_size, -1).view(dtype=torch.uint64) - cpu_cache_ssm = ( - cpu_cache_tensor[:, (a + b) : (a + b + c)].view(cpu_page_num, tp_world_size, -1).view(dtype=torch.uint64) - ) - - gpu_full_att_kv_state = gpu_full_att_kv_state.view( - gpu_full_att_kv_state.shape[0], gpu_full_att_kv_state.shape[1], -1 - ).view(dtype=torch.uint64) - gpu_full_att_kv_state = gpu_full_att_kv_state.permute(1, 0, 2) # [s, layer_num, xxdim] - - cpu_kv_conv_state = cpu_kv_conv_state.view(cpu_kv_conv_state.shape[0], -1).view(dtype=torch.uint64) - cpu_kv_ssm_state = cpu_kv_ssm_state.view(cpu_kv_ssm_state.shape[0], -1).view(dtype=torch.uint64) - - gpu_full_att_tail_dim = gpu_full_att_kv_state.shape[-1] * gpu_full_att_kv_state.shape[-2] * big_page_token_num - cpu_kv_conv_tail_dim = cpu_kv_conv_state.shape[-1] - cpu_kv_ssm_tail_dim = cpu_kv_ssm_state.shape[-1] - full_att_layer_num = gpu_full_att_kv_state.shape[-2] - - assert gpu_full_att_tail_dim == cpu_cache_full_att.shape[-1] - assert cpu_cache_conv.shape[-1] == cpu_kv_conv_state.shape[-1] - assert cpu_cache_ssm.shape[-1] == cpu_kv_ssm_state.shape[-1] - assert gpu_full_att_kv_state.stride(2) == 1 - - assert (tp_rank // head_scale_size) < linear_config.full_att_all_num_kv_heads - - grid = (grid_num,) - _copy_cpu_cache_to_kv_buffer[grid]( - page_num=len(page_indexes), - mem_indexes_ptr=mem_indexes, - page_indexes_ptr=page_indexes, - big_page_buffer_ids=big_page_buffer_ids, - cpu_cache_full_att=cpu_cache_full_att, - cpu_cache_full_att_stride_p=cpu_cache_full_att.stride(0), - cpu_cache_full_att_stride_h=cpu_cache_full_att.stride(1), - cpu_cache_full_att_stride_d=cpu_cache_full_att.stride(2), - cpu_cache_conv=cpu_cache_conv, - cpu_cache_conv_stride_p=cpu_cache_conv.stride(0), - cpu_cache_conv_stride_t=cpu_cache_conv.stride(1), - cpu_cache_conv_stride_d=cpu_cache_conv.stride(2), - cpu_cache_ssm=cpu_cache_ssm, - cpu_cache_ssm_stride_p=cpu_cache_ssm.stride(0), - cpu_cache_ssm_stride_t=cpu_cache_ssm.stride(1), - cpu_cache_ssm_stride_d=cpu_cache_ssm.stride(2), - gpu_kv_full_att_state=gpu_full_att_kv_state, - gpu_kv_full_att_stride_s=gpu_full_att_kv_state.stride(0), - gpu_kv_full_att_stride_l=gpu_full_att_kv_state.stride(1), - gpu_kv_full_att_stride_d=gpu_full_att_kv_state.stride(2), - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_conv_stride_s=cpu_kv_conv_state.stride(0), - cpu_kv_conv_stride_d=cpu_kv_conv_state.stride(1), - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_kv_ssm_stride_s=cpu_kv_ssm_state.stride(0), - cpu_kv_ssm_stride_d=cpu_kv_ssm_state.stride(1), - gpu_full_att_tail_dim=gpu_full_att_tail_dim, - cpu_kv_conv_tail_dim=cpu_kv_conv_tail_dim, - cpu_kv_ssm_tail_dim=cpu_kv_ssm_tail_dim, - tp_rank=tp_rank, - full_att_layer_num=full_att_layer_num, - big_page_token_num=big_page_token_num, - head_scale_size=head_scale_size, - BLOCK=BLOCK, - ) - - -@triton.jit -def _copy_linear_att_state_to_linear_att_state( - src_conv_state, - dst_conv_state, - src_ssm_state, - dst_ssm_state, - conv_size, - ssm_size, - BLOCK: tl.constexpr, -): - pid = tl.program_id(0) - grid_num = tl.num_programs(0) - - # copy conv state - num_conv_blocks = tl.cdiv(conv_size, BLOCK) - for i in range(pid, num_conv_blocks, grid_num): - start = i * BLOCK + tl.arange(0, BLOCK) - mask = start < conv_size - data = tl.load(src_conv_state + start, mask=mask, other=0) - tl.store(dst_conv_state + start, data, mask=mask) - - # copy ssm state - num_ssm_blocks = tl.cdiv(ssm_size, BLOCK) - for i in range(pid, num_ssm_blocks, grid_num): - start = i * BLOCK + tl.arange(0, BLOCK) - mask = start < ssm_size - data = tl.load(src_ssm_state + start, mask=mask, other=0) - tl.store(dst_ssm_state + start, data, mask=mask) - - -def copy_linear_att_state_to_linear_att_state( - src_conv_state: torch.Tensor, - src_ssm_state: torch.Tensor, - dst_conv_state: torch.Tensor, - dst_ssm_state: torch.Tensor, - grid_num: int = 16, -): - assert src_conv_state.shape == dst_conv_state.shape - assert src_ssm_state.shape == dst_ssm_state.shape - - BLOCK = 4096 - - src_conv_flat = src_conv_state.view(-1).view(dtype=torch.uint8) - dst_conv_flat = dst_conv_state.view(-1).view(dtype=torch.uint8) - src_ssm_flat = src_ssm_state.view(-1).view(dtype=torch.uint8) - dst_ssm_flat = dst_ssm_state.view(-1).view(dtype=torch.uint8) - - conv_size = src_conv_flat.shape[0] - ssm_size = src_ssm_flat.shape[0] - - grid = (grid_num,) - _copy_linear_att_state_to_linear_att_state[grid]( - src_conv_state=src_conv_flat, - dst_conv_state=dst_conv_flat, - src_ssm_state=src_ssm_flat, - dst_ssm_state=dst_ssm_flat, - conv_size=conv_size, - ssm_size=ssm_size, - BLOCK=BLOCK, - ) diff --git a/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py b/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py index a31292b4fc..08e548d779 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/q_per_head_fp8_quant.py @@ -82,7 +82,8 @@ def q_per_head_fp8_quant(q, seq_lens, b1_start_loc, scale_out=None, token_batch_ if scale_out is None: scale_out = torch.empty((B, H), dtype=torch.float32, device=q.device) if token_batch_ids is None: - token_batch_ids = torch.repeat_interleave(torch.arange(B, device=q.device), seq_lens) + # 输出 token 数已知，避免读取 GPU 上的长度总和，允许在 CUDA Graph 中调用。 + token_batch_ids = torch.repeat_interleave(torch.arange(B, device=q.device), seq_lens, output_size=T) _per_head_max_reduce_kernel[(B, H)]( q, diff --git a/lightllm/common/kv_cache_mem_manager/__init__.py b/lightllm/common/kv_cache_mem_manager/__init__.py deleted file mode 100644 index 3cd6a211d7..0000000000 --- a/lightllm/common/kv_cache_mem_manager/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -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 .deepseek4_mem_manager import DeepseekV4MemoryManager -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 .glm5_next_mem_manager import Glm5NextMemManager -from .fp8_glm5_next_mem_manager import FP8Glm5NextMemManager - -__all__ = [ - "KvCacheAllocator", - "MemoryManager", - "ReadOnlyStaticsMemoryManager", - "PPLINT4KVMemoryManager", - "PPLINT8KVMemoryManager", - "Deepseek2MemoryManager", - "Deepseek3_2MemoryManager", - "DeepseekV4MemoryManager", - "FP8PerTokenGroupQuantDeepseek3_2MemoryManager", - "FP8StaticPerHeadQuantMemManager", - "FP8StaticPerTensorQuantMemManager", - "Qwen3NextMemManager", - "Glm5NextMemManager", - "FP8Glm5NextMemManager", -] diff --git a/lightllm/common/kv_cache_mem_manager/allocator.py b/lightllm/common/kv_cache_mem_manager/allocator.py deleted file mode 100644 index 1331943b1a..0000000000 --- a/lightllm/common/kv_cache_mem_manager/allocator.py +++ /dev/null @@ -1,103 +0,0 @@ -import torch -from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt -from lightllm.utils.dist_utils import get_current_rank_in_node -from lightllm.utils.log_utils import init_logger -from typing import Union, List - -logger = init_logger(__name__) - - -class KvCacheAllocator: - def __init__(self, size: int) -> None: - self.size = size - self.mem_state = torch.arange( - 0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self.mark_start = 0 - self.mark_end = self.size - - self._mem_state_return = torch.arange( - 0, self.size * 3, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self._return_start = 0 - - self.can_use_mem_size = self.size - - rank_in_node = get_current_rank_in_node() - # 用共享内存进行共享，router 模块读取进行精确的调度估计；基础层会统一添加服务前缀以防止实例冲突。 - self.shared_can_use_token_num = SharedInt(f"mem_manger_can_use_token_num_{rank_in_node}") - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - return - - def alloc(self, need_size) -> torch.Tensor: - if need_size > self.mark_end - self.mark_start: - logger.error(f"warn no enough cache need_size {need_size} left_size {self.can_use_mem_size}") - assert False, "error alloc state" - - start = self.mark_start - end = self.mark_start + need_size - self.mark_start += need_size - - self.can_use_mem_size -= need_size - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - - # 利用缓冲区返回，避免异步情况下的内存竞争 - if self._return_start + need_size > self._mem_state_return.shape[0]: - self._return_start = 0 - ans = self._mem_state_return[self._return_start : self._return_start + need_size] - ans.copy_(self.mem_state[start:end]) - self._return_start += need_size - return ans - - def free(self, free_index: Union[torch.Tensor, List[int]]): - """_summary_ - - Args: - free_index (torch.Tensor): _description_ - """ - - end = self.mark_start - start = self.mark_start - len(free_index) - assert start >= 0, f"error free state start: {self.mark_start} free len {len(free_index)}" - - if isinstance(free_index, list): - self.mem_state.numpy()[start:end] = free_index - else: - # 从 gpu 到 cpu 的拷贝操作是流内阻塞操作 - self.mem_state[start:end] = free_index - - self.mark_start -= len(free_index) - - self.can_use_mem_size += len(free_index) - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - - if self.can_use_mem_size == len(self.mem_state): - logger.debug(f"freed all gpu mem size {self.can_use_mem_size}") - return - - def free_all(self): - self.mem_state.numpy()[:] = list(range(0, len(self.mem_state))) - self.mark_start = 0 - self.mark_end = len(self.mem_state) - self.can_use_mem_size = len(self.mem_state) - self.shared_can_use_token_num.set_value(self.can_use_mem_size) - return - - def resize(self, new_size: int) -> None: - """ - just for test code - """ - self.size = new_size - self.mem_state = torch.arange( - 0, self.size, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self.mark_start = 0 - self.mark_end = self.size - - self._mem_state_return = torch.arange( - 0, self.size * 3, dtype=torch.int32, device="cpu", requires_grad=False, pin_memory=True - ) - self._return_start = 0 - - self.can_use_mem_size = self.size - self.shared_can_use_token_num.set_value(self.can_use_mem_size) 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/deepseek4_mem_manager.py b/lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py deleted file mode 100644 index 6794446508..0000000000 --- a/lightllm/common/kv_cache_mem_manager/deepseek4_mem_manager.py +++ /dev/null @@ -1,175 +0,0 @@ -import math - -import torch -import triton - -from .mem_manager import MemoryManager -from .operator import DeepseekV4MemOperator -from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4CacheConfig, DeepseekV4StateCacheManager -from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_compressed_kv_pages -from lightllm.utils.envs_utils import get_env_start_args - - -class DeepseekV4MemoryManager(MemoryManager): - operator_class = DeepseekV4MemOperator - - def __init__( - self, - size, - dtype, - head_num, - head_dim, - layer_num, - cache_config: DeepseekV4CacheConfig, - always_copy=False, - mem_fraction=0.9, - ): - self.cache_config = cache_config - super().__init__( - size, - torch.uint8, - head_num, - max(cache_config.compressed_kv_slot_bytes), - len(cache_config.compressed_kv_slot_bytes), - always_copy, - mem_fraction, - ) - - def get_cell_size(self): - return sum(self.cache_config.compressed_kv_slot_bytes) - - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - self.kv_buffer = [ - torch.zeros((size + self.page_size, head_num, slot_bytes), dtype=dtype, device="cuda") - for slot_bytes in self.cache_config.compressed_kv_slot_bytes - ] - args = get_env_start_args() - big_page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num - # Reserve the same CPU-cache load/offload staging slots as the other hybrid models. - self.big_page_buffers = DeepseekV4StateCacheManager( - triton.cdiv(size, big_page_tokens) + 2, self.cache_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 - - def _free_buffers(self): - self.kv_buffer = self.big_page_buffers = None - - def free_all(self): - super().free_all() - self.big_page_buffers.clear_to_init_state() - - def get_att_input_params(self, layer_index): - index = self.cache_config.layer_to_compressed_kv_idx.get(layer_index) - return self.kv_buffer[index] if index is not None else None - - def get_index_kv_buffer(self, index): - return {f"layer_{layer}": buffer[index].unsqueeze(0) for layer, buffer in enumerate(self.kv_buffer)} - - def load_index_kv_buffer(self, index, load_tensor_dict): - for layer, buffer in enumerate(self.kv_buffer): - buffer[index] = load_tensor_dict[f"layer_{layer}"][0].to(buffer.device) - - def __getstate__(self): - state = self.__dict__.copy() - state["big_page_buffers"] = None # Keep process-local pinned checkpoints out of CUDA IPC. - return state - - def write_to_shm(self, req_manager): - # Transport receives GPU views, just as Qwen's mem manager exports linear state. - for name in ( - "req_to_swa", - "req_to_c4_state", - "req_to_c4_indexer_state", - "req_to_c128_state", - ): - setattr(self, name, getattr(req_manager, name)) - return super().write_to_shm(req_manager) - - def get_paged_kv_move_buffer_shape(self, page_num, page_size): - compressed_kv_bytes = page_size * self.get_cell_size() - swa_shape, *compressor_shapes = self.cache_config.get_runtime_state_shapes(get_env_start_args().mtp_step) - state_bytes = ( - math.prod(swa_shape) + sum(math.prod(shape) for shape in compressor_shapes) * torch.float32.itemsize - ) - return (page_num, 1, 1, 1, max(compressed_kv_bytes, state_bytes)) - - def alloc_paged_kv_move_buffer(self, page_num, page_size): - self.pd_kv_page_size = page_size - kv_move_buffer = super().alloc_paged_kv_move_buffer(page_num, page_size) - self.att_state_page_helper = DeepseekV4AttStatePageHelper(self) - return kv_move_buffer - - 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": - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - dp_mems = mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] - return self.att_state_page_helper.write_req_to_page(page_index, req_idx, dp_mems) - pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] - pin_indexes.numpy()[:] = mem_indexes - page = self.kv_move_buffer[page_index] - copy_compressed_kv_pages( - mem_manager=mem_managers[dp_index * dp_world_size], - mem_indexes=pin_indexes.cuda(non_blocking=True), - page_indexes=torch.zeros(1, dtype=torch.int32, device=page.device), - kv_move_buffer=page.view(1, -1), - to_pages=True, - page_size=self.pd_kv_page_size, - round_tail=True, - ) - - 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": - assert page_kind == "att_state", f"unknown page_kind={page_kind}" - dp_mems = mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] - return self.att_state_page_helper.read_page_to_req(page_index, req_idx, dp_mems) - pin_indexes = self._buffer_mem_indexes_tensors[page_index][: len(mem_indexes)] - pin_indexes.numpy()[:] = mem_indexes - indexes = pin_indexes.cuda(non_blocking=True) - page = self.kv_move_buffer[page_index] - page_ids = torch.zeros(1, dtype=torch.int32, device=page.device) - for mem in mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size]: - copy_compressed_kv_pages( - mem_manager=mem, - mem_indexes=indexes, - page_indexes=page_ids, - kv_move_buffer=page.view(1, -1), - to_pages=False, - page_size=self.pd_kv_page_size, - round_tail=True, - ) - - -class DeepseekV4AttStatePageHelper: - """整块迁移请求的 SWA 和 compressor ring，保持物理槽位顺序。""" - - def __init__(self, mem_manager): - self.mem_manager = mem_manager - - def write_req_to_page(self, page_index, req_idx, dp_mems): - page = self.mem_manager.kv_move_buffer[page_index] - self._copy_req_state(page, dp_mems[0], req_idx, write=True) - - def read_page_to_req(self, page_index, req_idx, dp_mems): - page = self.mem_manager.kv_move_buffer[page_index] - for mem in dp_mems: - self._copy_req_state(page, mem, req_idx, write=False) - - def _copy_req_state(self, page, mem, req_idx, *, write): - data = page.reshape(-1) - offset = 0 - for buffer in (mem.req_to_swa, mem.req_to_c4_state, mem.req_to_c4_indexer_state, mem.req_to_c128_state): - if buffer is None: - continue - state = buffer[:, req_idx] - nbytes = state.numel() * state.element_size() - page_state = data[offset : offset + nbytes].view(state.dtype).view(state.shape) - if write: - page_state.copy_(state, non_blocking=True) - else: - state.copy_(page_state, non_blocking=True) - offset += nbytes diff --git a/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py deleted file mode 100644 index f28f400331..0000000000 --- a/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py +++ /dev/null @@ -1,19 +0,0 @@ -from .glm5_next_mem_manager import Glm5NextMemManager -from .operator import LinearAttMemOperator - - -class FP8Glm5NextMemOperator(LinearAttMemOperator): - def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): - from lightllm.models.glm5_next.triton_kernel.destindex_copy_kv_flashmla_fp8 import ( - destindex_copy_kv_flashmla_fp8, - ) - - packed_kv = self.mem_manager.get_att_input_params(layer_index) - destindex_copy_kv_flashmla_fp8(kv, mem_index, packed_kv) - - -class FP8Glm5NextMemManager(Glm5NextMemManager): - operator_class = FP8Glm5NextMemOperator - - def get_att_input_params(self, layer_index): - return self._layer_buffer(layer_index)[:, :, : self.linear_config.FP8_MLA_KV_BYTES] 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_head_quant_mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py deleted file mode 100755 index 2d5e8a1e5e..0000000000 --- a/lightllm/common/kv_cache_mem_manager/fp8_static_per_head_quant_mem_manager.py +++ /dev/null @@ -1,98 +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_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 - -logger = init_logger(__name__) - - -class FP8StaticPerHeadQuantMemManager(MemoryManager): - operator_class = FP8StaticPerHeadQuantMemOperator - - 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 - - cfg, all_scales = self._load_and_check_config() - q_cfg = cfg["q_calibration"] - - all_head_num = cfg["num_head"] - # A joint target+draft config uses its leading target scale rows on a target-only server. - all_scales = all_scales[: self.layer_num].to(device="cuda") - - factor = (get_dp_world_size() * head_num) // all_head_num - all_scales = torch.repeat_interleave(input=all_scales, repeats=factor, dim=-1) - rank_in_dp = get_current_rank_in_dp() - - v_offset = all_scales.shape[1] // 2 - start_head = rank_in_dp * head_num - end_head = start_head + head_num - k_scales = all_scales[:, start_head:end_head].contiguous() - v_scales = all_scales[:, v_offset + start_head : v_offset + end_head].contiguous() - self.scales = torch.cat((k_scales, v_scales), dim=-1) - - all_q_scales = torch.tensor(q_cfg["scales"], dtype=torch.float32, device="cuda").view(q_cfg["scales_shape"])[ - : self.layer_num - ] - q_factor = (get_dp_world_size() * head_num) // q_cfg["num_head"] - all_q_scales = torch.repeat_interleave(input=all_q_scales, repeats=q_factor, dim=-1) - self.q_scales = all_q_scales[:, start_head:end_head].contiguous() - return - - def _load_and_check_config(self): - config_path = get_env_start_args().kv_quant_calibration_config_path - if config_path is None: - raise ValueError("fp8kv_sph requires kv_quant_calibration_config_path with q_calibration") - - logger.info(f"kv_quant_calibration_config_path {config_path} is set, will load kv quant calibration config") - if not os.path.exists(config_path): - raise FileNotFoundError(f"kv_quant_calibration_config {config_path} not found") - - with open(config_path, "r") as f: - cfg = json.load(f) - - num_layers, num_heads = cfg["num_layers"], cfg["num_head"] - scales = torch.tensor(cfg["scales"], dtype=torch.float32) - expected_shape = [num_layers, 2 * num_heads] - if cfg["scales_shape"] != expected_shape or list(scales.shape) != expected_shape: - raise ValueError(f"calibration scales and scales_shape must match {expected_shape}") - - runtime_draft_layers = get_added_mtp_kv_layer_num() - runtime_target_layers = self.layer_num - runtime_draft_layers - if "num_target_layers" in cfg or "num_draft_layers" in cfg: - target_layers, draft_layers = cfg["num_target_layers"], cfg["num_draft_layers"] - if ( - draft_layers < 0 - or target_layers + draft_layers != num_layers - or target_layers != runtime_target_layers - or (runtime_draft_layers > 0 and draft_layers != runtime_draft_layers) - ): - raise ValueError( - f"invalid calibration layer layout: target={target_layers}, draft={draft_layers}, " - f"runtime target={runtime_target_layers}, draft={runtime_draft_layers}" - ) - elif num_layers != self.layer_num: - raise ValueError(f"legacy calibration num_layers {num_layers} do not match runtime {self.layer_num}") - - q_cfg = cfg.get("q_calibration") - if not q_cfg: - raise ValueError("fp8kv_sph calibration config requires q_calibration") - assert (get_dp_world_size() * self.head_num) % num_heads == 0 - 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 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 deleted file mode 100644 index 74f0f26810..0000000000 --- a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py +++ /dev/null @@ -1,122 +0,0 @@ -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 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): - 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.""" - - 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): - assert head_num == 1 - self.kv_buffer = torch.empty((layer_num, size + self.page_size, 1, head_dim), dtype=dtype, device="cuda") - self._init_linear_att_buffers() - - def _layer_buffer(self, layer_index): - return self.kv_buffer[self.linear_config.get_full_att_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/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py deleted file mode 100755 index 7108f39d79..0000000000 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ /dev/null @@ -1,308 +0,0 @@ -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.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 - - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): - self.size = size - self.head_num = head_num - self.head_dim = head_dim - self.layer_num = layer_num - self.always_copy = always_copy - self.dtype = dtype - # profile the max total token num if the size is None - self.profile_size(mem_fraction) - - # A physical KV page must never straddle the allocator boundary. The - # unused remainder is intentionally dropped so every managed page has - # a stable ``mem_index // page_size`` id. - self.page_size = get_env_start_args().page_size - self.size = self.size // self.page_size * self.page_size - - self.allocator = KvCacheAllocator(self.size) - - self._init_buffers( - self.size, - dtype, - head_num, - head_dim, - layer_num, - ) - self.HOLD_TOKEN_MEMINDEXES = tuple(range(self.size, self.size + self.page_size)) - - # 构建对外的操作类接口 - self.operator: BaseMemManagerOperator = self.operator_class(self) - - 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 - - def get_cell_size(self): - return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype) - - 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 ( - page_num, - page_size, - self.layer_num, - 2 * num_kv_head, - self.head_dim, - ) - - def get_pd_kv_move_buffer_size(self): - args = get_env_start_args() - if args.run_mode not in ["prefill", "decode"]: - return 0 - shape = self.get_paged_kv_move_buffer_shape( - page_num=args.pd_kv_page_num, - page_size=args.pd_kv_page_size, - ) - return math.prod(shape) * torch._utils._element_size(self.dtype) - - def profile_size(self, mem_fraction): - if self.size is not None: - return - - torch.cuda.empty_cache() - world_size = dist.get_world_size() - 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 - 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()}") - dist.all_reduce(tensor, op=dist.ReduceOp.MIN) - self.size = tensor.item() - logger.info( - f"{str(available_memory)} GB space is available after load the model weight\n" - f"{str(pd_kv_move_buffer_size / 1024 ** 2)} MB is reserved for PD KV transfer buffer\n" - f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n" - f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n" - ) - return - - 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}" - # 在初始化 kv_buffer 时，每层额外保留一整页；这部分空间不会被 allocator 对外分配。 - # 这一页用于多 DP、overlap microbatch 等模式下的 padding 请求， - # 页内索引存储在 HOLD_TOKEN_MEMINDEXES 中，与 req_manager 的 HOLD_REQUEST_ID 作用类似。 - self.kv_buffer = torch.empty( - (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 - - def alloc(self, need_size) -> torch.Tensor: - assert ( - need_size % self.page_size == 0 - ), f"KV cache allocation size {need_size} must be a multiple of page_size {self.page_size}" - return self.allocator.alloc(need_size) - - def free(self, free_index: Union[torch.Tensor, List[int]]) -> None: - self.allocator.free(free_index) - - def free_all(self): - self.allocator.free_all() - - def resize_mem(self, new_size): - """ - just for test code - """ - size = new_size // self.page_size * self.page_size - dtype = self.dtype - head_num = self.head_num - head_dim = self.head_dim - layer_num = self.layer_num - - self.size = size - self.allocator.resize(size) - self._free_buffers() - self._init_buffers(size, dtype, head_num, head_dim, layer_num) - self.HOLD_TOKEN_MEMINDEXES = tuple(range(self.size, self.size + self.page_size)) - return - - def get_index_kv_buffer(self, index): - return {"kv_buffer": self.kv_buffer[:, index]} - - 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)] - # Each consumer needs its own IPC ref-counter; its pickle size can vary. - obj_size = max(len(obj_bytes) for obj_bytes in obj_bytes_array) - shm = create_or_link_shm( - 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) -> "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: - """ - 读取一些统计信息 - """ - - def __init__(self) -> None: - args = get_env_start_args() - self.global_world_size = args.tp - self.node_world_size = args.tp // args.nnodes - self.dp_world_size = self.global_world_size // args.dp - # 兼容多机 dp size=1 纯 tp 模式的情况 - self.is_multinode_tp = args.dp == 1 and args.nnodes > 1 - self.shared_tp_infos = [ - SharedInt(f"mem_manger_can_use_token_num_{rank_in_node}") - for rank_in_node in range(0, self.node_world_size, self.dp_world_size) - ] - - def get_unrefed_token_num(self, dp_rank_in_node: int): - if self.is_multinode_tp: - return self.shared_tp_infos[0].get_value() - return self.shared_tp_infos[dp_rank_in_node].get_value() diff --git a/lightllm/common/kv_cache_mem_manager/mem_utils.py b/lightllm/common/kv_cache_mem_manager/mem_utils.py deleted file mode 100644 index 398f407e36..0000000000 --- a/lightllm/common/kv_cache_mem_manager/mem_utils.py +++ /dev/null @@ -1,76 +0,0 @@ -from . import ( - 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__) - - -@lru_cache(maxsize=None) -def select_mem_manager_class(): - # case 1 - # 先判断是否是 deepseek 系列的模型 - model_class = get_llm_model_class() - from lightllm.models import DeepseekV4TpPartModel - - if issubclass(model_class, DeepseekV4TpPartModel): - from . import DeepseekV4MemoryManager - - mem_class = DeepseekV4MemoryManager - logger.info(f"Model kv cache using default, mem_manager class: {mem_class}") - return mem_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 - - 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/operator/__init__.py b/lightllm/common/kv_cache_mem_manager/operator/__init__.py deleted file mode 100644 index 442c2e300e..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -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, - DeepseekV4MemOperator, - FP8PerTokenGroupQuantDeepseek3_2MemOperator, -) -from .fp8_quant import ( - FP8StaticPerHeadQuantMemOperator, - FP8StaticPerTensorQuantMemOperator, -) diff --git a/lightllm/common/kv_cache_mem_manager/operator/base.py b/lightllm/common/kv_cache_mem_manager/operator/base.py deleted file mode 100644 index 682b8a5d60..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/base.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -from abc import ABC, abstractmethod -from typing import List, TYPE_CHECKING - -if TYPE_CHECKING: - from ..mem_manager import MemoryManager - from lightllm.server.router.model_infer.infer_batch import InferReq - -# 定义一个抽象基类 -class BaseMemManagerOperator(ABC): - def __init__(self, mem_manager: "MemoryManager") -> None: - super().__init__() - self.mem_manager = mem_manager - - @abstractmethod - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - pass - - def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): - raise NotImplementedError() - - # cpu cache 的相关操作接口 - def load_cpu_cache_to_gpu( - self, mem_indexes: torch.Tensor, page_indexes: torch.Tensor, cpu_cache_client, req: "InferReq" - ): - raise NotImplementedError() - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client, - req: "InferReq", - ): - raise NotImplementedError() - - # dp 间共享 kv 的操作接口, 提升dp 模式下的kv 共享效率，降低调度的难度 - def copy_kv_from_other_dp_ranks( - self, - mem_managers: List["MemoryManager"], - move_token_indexes: torch.Tensor, - token_dp_indexes: torch.Tensor, - mem_indexes: torch.Tensor, - dp_size_in_node: int, - rank_in_dp: int, - ): - raise NotImplementedError() 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 6be23b5cfe..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/deepseek.py +++ /dev/null @@ -1,207 +0,0 @@ -from typing import TYPE_CHECKING - -import torch -from .normal import NormalMemOperator -from .base import BaseMemManagerOperator -from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args - -if TYPE_CHECKING: - from ..deepseek4_mem_manager import DeepseekV4MemoryManager - from lightllm.common.req_manager.deepseek4 import DeepseekV4ReqManager - 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 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 - - -class DeepseekV4MemOperator(BaseMemManagerOperator): - def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): - """仅复制压缩 KV 和 indexer K；SWA、compressor 状态由 ReqManager 恢复。""" - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer - - src = src_mem_index.to(device="cuda", dtype=torch.int64, non_blocking=True) - dst = dst_mem_index.to(device="cuda", dtype=torch.int64, non_blocking=True) - for buffer in self.mem_manager.kv_buffer: - copy_kv_buffer_to_kv_buffer(src, dst, buffer.view(1, buffer.shape[0], -1)) - - def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): - # 模型推理层通过融合算子写入 SWA，compressor/indexer 直接写入压缩 KV。 - raise NotImplementedError("DeepSeek-V4 writes packed KV using request-owned SWA slots") - - 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", - ) -> None: - from lightllm.utils.dist_utils import get_current_rank_in_dp - from lightllm.server.router.model_infer.infer_batch import g_infer_context - from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_kv_buffer_to_cpu_cache, copy_state_cache - - # MLA KV and compression states are replicated across TP ranks. - if get_current_rank_in_dp() != 0: - return - - mem_manager: "DeepseekV4MemoryManager" = self.mem_manager - req_manager: "DeepseekV4ReqManager" = g_infer_context.req_manager - page_token_num = get_env_start_args().cpu_cache_token_page_size - - 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) // page_token_num) * page_token_num - start_kv_len = (len(big_page_buffer_ids_cpu) + 1) * page_token_num - for seq_len in range(start_kv_len, max_kv_len + 1, page_token_num): - big_page_buffer_ids_cpu.append(req.hybrid_len_to_big_page_id[seq_len]) - - # 碎页：将尾部小页状态复制到临时大页，KV 由拷贝算子按实际 token 数处理。 - has_tail_page = len(mem_indexes) % page_token_num != 0 - if has_tail_page: - assert req.tail_small_page_buffer_id is not None - temp_page_id = mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID - copy_state_cache( - src_state=req_manager.small_page_buffers.buffer[req.tail_small_page_buffer_id], - dst_state=mem_manager.big_page_buffers.buffer[temp_page_id], - ) - big_page_buffer_ids_cpu.append(temp_page_id) - - big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( - non_blocking=True - ) - copy_kv_buffer_to_cpu_cache( - mem_manager=mem_manager, - mem_indexes=mem_indexes, - page_indexes=page_indexes, - page_readies=page_readies, - big_page_buffer_ids=big_page_buffer_ids_gpu, - cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, - ) - - def load_cpu_cache_to_gpu( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ) -> None: - from lightllm.server.router.model_infer.infer_batch import g_infer_context - from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_cpu_cache_to_kv_buffer, copy_state_cache - - mem_manager: "DeepseekV4MemoryManager" = self.mem_manager - req_manager: "DeepseekV4ReqManager" = g_infer_context.req_manager - page_token_num = get_env_start_args().cpu_cache_token_page_size - - big_page_num = len(mem_indexes) // page_token_num - max_kv_len = (req.cur_kv_len // page_token_num) * page_token_num - big_page_buffer_ids_cpu = [] - for _ 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 -= page_token_num - big_page_buffer_ids_cpu.reverse() - - # 碎页：先将状态读入临时大页，再恢复请求运行态。 - has_tail_page = len(mem_indexes) % page_token_num != 0 - if has_tail_page: - 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 - ) - copy_cpu_cache_to_kv_buffer( - mem_manager=mem_manager, - mem_indexes=mem_indexes, - big_page_buffer_ids=big_page_buffer_ids_gpu, - page_indexes=page_indexes, - cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor, - ) - - if has_tail_page and req.cur_kv_len == req.hybrid_cache_len: - g_infer_context.radix_cache.free_one_small_page_buffer() - small_page_buffer_id = req_manager.small_page_buffers.alloc_one_state_cache() - if small_page_buffer_id is not None: - req.tail_small_page_buffer_id = small_page_buffer_id - # 临时大页会被复用，先将碎页状态保留到小页池。 - copy_state_cache( - src_state=mem_manager.big_page_buffers.buffer[big_page_buffer_ids_cpu[-1]], - dst_state=req_manager.small_page_buffers.buffer[small_page_buffer_id], - ) - - req_manager.restore_big_page_state(big_page_buffer_idx=big_page_buffer_ids_cpu[-1], req=req) diff --git a/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py b/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py deleted file mode 100644 index 4b1f0ac207..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/fp8_quant.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -from typing import TYPE_CHECKING -from .base import BaseMemManagerOperator -from lightllm.utils.log_utils import init_logger - -if TYPE_CHECKING: - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient -logger = init_logger(__name__) - - -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 ( - FP8StaticPerHeadQuantMemManager, - ) - - mem_manager: FP8StaticPerHeadQuantMemManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import ( - destindex_copy_kv_fp8, - ) - - scales = mem_manager.scales - destindex_copy_kv_fp8( - kv, - mem_index, - scales[layer_index], - mem_manager.kv_buffer[layer_index].view(torch.float8_e4m3fn), - ) - return - - -class FP8StaticPerTensorQuantMemOperator(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.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import ( - destindex_copy_kv_fp8, - ) - - destindex_copy_kv_fp8( - kv, - mem_index, - mem_manager.scales[layer_index], - mem_manager.kv_buffer[layer_index].view(torch.float8_e4m3fn), - is_per_tensor_quant=True, - ) - return 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/normal.py b/lightllm/common/kv_cache_mem_manager/operator/normal.py deleted file mode 100644 index 3c53ace079..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/normal.py +++ /dev/null @@ -1,115 +0,0 @@ -import torch -from typing import TYPE_CHECKING, List -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 NormalMemOperator(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.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 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 len(mem_indexes) % args.cpu_cache_token_page_size == 0 - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu - - load_cpu_kv_to_gpu( - gpu_mem_indexes=mem_indexes, - gpu_kv_cache=mem_manager.kv_buffer, - gpu_kv_cache_scale=None, - cpu_kv_cache=cpu_cache_client.cpu_kv_cache_tensor, - cpu_kv_cache_scale=None, - page_indexes=page_indexes, - tp_index=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - grid_num=16, - ) - 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", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda - args = get_env_start_args() - assert len(mem_indexes) % args.cpu_cache_token_page_size == 0 - assert len(mem_indexes) // args.cpu_cache_token_page_size == len(page_indexes) - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu - - offload_gpu_kv_to_cpu( - token_indexes=mem_indexes, - gpu_kv_cache=mem_manager.kv_buffer, - gpu_kv_cache_scale=None, - cpu_kv_cache=cpu_cache_client.cpu_kv_cache_tensor, - cpu_kv_cache_scale=None, - page_indexes=page_indexes, - page_readies=page_readies, - tp_index=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - grid_num=16, - ) - return - - def copy_kv_from_other_dp_ranks( - self, - mem_managers: List, - move_token_indexes: torch.Tensor, - token_dp_indexes: torch.Tensor, - mem_indexes: torch.Tensor, - dp_size_in_node: int, - rank_in_dp: int, - ): - 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()) - 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 - - # 一次性传输所有层 - kv_trans_for_dp( - input_mems=self.mem_ptrs_tensor.cuda(non_blocking=True), - input_idx=move_token_indexes, - input_dp_idx=token_dp_indexes, - output=self.mem_manager.kv_buffer, - output_idx=mem_indexes, - dp_size_in_node=dp_size_in_node, - rank_in_dp=rank_in_dp, - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/operator/quant.py b/lightllm/common/kv_cache_mem_manager/operator/quant.py deleted file mode 100644 index a3a1c1d01c..0000000000 --- a/lightllm/common/kv_cache_mem_manager/operator/quant.py +++ /dev/null @@ -1,129 +0,0 @@ -import torch -from typing import TYPE_CHECKING, List -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 QuantScaleMemOperator(BaseMemManagerOperator): - """ - 对于kv cache中包含独立的对应每个token的scale变量的memManager使用。 - """ - - 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 len(mem_indexes) % args.cpu_cache_token_page_size == 0 - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - - cpu_cache_meta = cpu_cache_client.kv_cache_tensor_meta - cpu_kv_cache = cpu_cache_client.cpu_kv_cache_tensor[:, :, :, :, 0 : cpu_cache_meta.head_dim] - cpu_kv_cache_scale = cpu_cache_client.cpu_kv_cache_tensor[:, :, :, :, cpu_cache_meta.head_dim :].view( - mem_manager.scale_buffer.dtype - ) - - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu - - load_cpu_kv_to_gpu( - gpu_mem_indexes=mem_indexes, - gpu_kv_cache=mem_manager.kv_buffer, - gpu_kv_cache_scale=mem_manager.scale_buffer, - cpu_kv_cache=cpu_kv_cache, - cpu_kv_cache_scale=cpu_kv_cache_scale, - page_indexes=page_indexes, - tp_index=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - grid_num=16, - ) - 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", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda - args = get_env_start_args() - assert len(mem_indexes) % args.cpu_cache_token_page_size == 0 - assert len(mem_indexes) // args.cpu_cache_token_page_size == len(page_indexes) - from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager - - mem_manager: MemoryManager = self.mem_manager - - cpu_cache_meta = cpu_cache_client.kv_cache_tensor_meta - cpu_kv_cache = cpu_cache_client.cpu_kv_cache_tensor[:, :, :, :, 0 : cpu_cache_meta.head_dim] - cpu_kv_cache_scale = cpu_cache_client.cpu_kv_cache_tensor[:, :, :, :, cpu_cache_meta.head_dim :].view( - mem_manager.scale_buffer.dtype - ) - - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu - - offload_gpu_kv_to_cpu( - token_indexes=mem_indexes, - gpu_kv_cache=mem_manager.kv_buffer, - gpu_kv_cache_scale=mem_manager.scale_buffer, - cpu_kv_cache=cpu_kv_cache, - cpu_kv_cache_scale=cpu_kv_cache_scale, - page_indexes=page_indexes, - page_readies=page_readies, - tp_index=get_current_rank_in_dp(), - tp_world_size=get_dp_world_size(), - grid_num=16, - ) - return - - -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 - - mem_manager: PPLINT4KVMemoryManager = self.mem_manager - from ...basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import ( - destindex_copy_int4kv, - ) - - destindex_copy_int4kv( - kv, - mem_index, - mem_manager.kv_buffer[layer_index], - mem_manager.scale_buffer[layer_index], - quant_group_size=mem_manager.group_quant_size, - ) - return - - -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 - - mem_manager: PPLINT8KVMemoryManager = self.mem_manager - from ...basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import ( - destindex_copy_quantize_kv, - ) - - destindex_copy_quantize_kv( - kv, - mem_index, - mem_manager.kv_buffer[layer_index], - mem_manager.scale_buffer[layer_index], - quant_group_dim=mem_manager.group_quant_size, - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py b/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py deleted file mode 100755 index 6b02cbed63..0000000000 --- a/lightllm/common/kv_cache_mem_manager/ppl_int4kv_mem_manager.py +++ /dev/null @@ -1,50 +0,0 @@ -import torch -from typing import Tuple, Any -from .mem_manager import MemoryManager -from .operator import PPLInt4KVMemOperator - - -class PPLINT4KVMemoryManager(MemoryManager): - - operator_class = PPLInt4KVMemOperator - - 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 // 2 * 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 // 2), 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/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/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index 7f2e6ac9b5..711915bd04 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -1,15 +1,7 @@ from .base import ReqManager -from .deepseek4 import DeepseekV4ReqManager -from .linear_att import ReqManagerForMamba -from .glm5_next import Glm5NextReqManager -from .hybrid_base import HybridAttentionReqManager from .req_sampling_params import ReqSamplingParamsManager __all__ = [ "ReqManager", - "DeepseekV4ReqManager", - "HybridAttentionReqManager", - "ReqManagerForMamba", - "Glm5NextReqManager", "ReqSamplingParamsManager", ] diff --git a/lightllm/common/req_manager/base.py b/lightllm/common/req_manager/base.py index cfe5d716ee..0088c6b82b 100644 --- a/lightllm/common/req_manager/base.py +++ b/lightllm/common/req_manager/base.py @@ -2,7 +2,6 @@ import torch -from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger @@ -52,7 +51,7 @@ def is_all_free(self): class ReqManager: - def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryManager): + def __init__(self, max_request_num, max_sequence_length): from .req_sampling_params import ReqSamplingParamsManager # 这里对最大请求数量的管理在默认上多申请了一个，主要是 index 为 max_request_num 代表 @@ -60,7 +59,7 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana # 的那个batch size 进行运行，所有 padding 的请求都会使用预留的这个请求管理 id 进行处理 # 这样让 DP 的实现更为简化一些。 self.req_list = _ReqLinkedList(max_request_num) - page_size = get_env_start_args().page_size + self.page_size = page_size = get_env_start_args().page_size max_sequence_length = (max_sequence_length + page_size - 1) // page_size * page_size self.req_to_token_indexs = torch.zeros( (max_request_num + 1, max_sequence_length), dtype=torch.int32, device="cuda" @@ -68,43 +67,31 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager: MemoryMana self.req_sampling_params_manager = ReqSamplingParamsManager(max_request_num) self.max_request_num = max_request_num self.HOLD_REQUEST_ID = max_request_num - self.mem_manager = None - if mem_manager is not None: - self.bind_mem_manager(mem_manager) - def bind_mem_manager(self, mem_manager: MemoryManager): - self.mem_manager = mem_manager - - self.init_hold_request_indexs() - return - - def init_hold_request_indexs(self): + def bind_hold_page(self, token_indices): + """Map padding requests to the resource owner's reserved physical page.""" assert ( self.req_list.is_all_free() ), "hold request indexes can only be initialized when all requests are released" - - # HOLD_REQUEST_ID 对应的请求行供 DP padding、overlap microbatch 等占位请求使用。将该行 - # 按 page_size 划分后，每一页都映射到 mem_manager 额外保留的同一个物理页；这样占位请求 - # 无论访问哪一个逻辑位置，都会落到合法且不会参与正常分配的 KV cache 地址上。 hold_row = self.req_to_token_indexs[self.HOLD_REQUEST_ID] - hold_page = torch.tensor( - self.mem_manager.HOLD_TOKEN_MEMINDEXES, + hold_page = torch.as_tensor( + token_indices, dtype=hold_row.dtype, device=hold_row.device, ) - hold_row.view(-1, self.mem_manager.page_size).copy_(hold_page) - return + if hold_page.ndim != 1 or hold_page.numel() != self.page_size: + raise ValueError("HOLD indices must describe exactly one physical page") + hold_row.view(-1, self.page_size).copy_(hold_page) def alloc(self): return self.req_list.alloc() - def free(self, free_req_indexes: List[int], free_token_index): + def free(self, free_req_indexes: List[int]): for req_index in free_req_indexes: self.req_list.free(req_index) if self.req_list.is_all_free(): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") - self.mem_manager.free(free_token_index) def free_req(self, free_req_index: int): self.req_list.free(free_req_index) @@ -112,10 +99,6 @@ def free_req(self, free_req_index: int): logger.debug(f"freed all request size {self.req_list.can_alloc_size}") return - def free_token(self, free_token_index): - self.mem_manager.free(free_token_index) - return - def free_all(self): self.req_list = _ReqLinkedList(self.max_request_num) return diff --git a/lightllm/common/req_manager/deepseek4.py b/lightllm/common/req_manager/deepseek4.py deleted file mode 100644 index f90d9fc4b9..0000000000 --- a/lightllm/common/req_manager/deepseek4.py +++ /dev/null @@ -1,88 +0,0 @@ -import torch - -from .hybrid_base import HybridAttentionReqManager -from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4StateCacheManager - - -class DeepseekV4ReqManager(HybridAttentionReqManager): - """持有跨 forward 的请求映射、SWA ring 和 compressor ring，并负责状态保存和恢复。 - - 当前 forward 的派生索引、页表和元数据由 infer_state 构建并持有， - 模型 workspace 提供投影与量化的计算 scratch。 - """ - - def __init__( - self, - max_request_num, - max_sequence_length, - mem_manager, - cache_config, - *, - mtp_step=0, - ): - super().__init__(max_request_num, max_sequence_length, mem_manager) - self.cache_config = cache_config - self.sliding_window = cache_config.sliding_window - swa_shape, c4_shape, indexer_shape, c128_shape = cache_config.get_runtime_state_shapes(mtp_step) - self.swa_capacity_tokens = swa_shape[1] - self.HOLD_SWA_SLOT = self.HOLD_REQUEST_ID * self.swa_capacity_tokens - self.c4_state_tokens = c4_shape[1] - self.c128_state_tokens = c128_shape[1] - self.req_to_swa = torch.zeros( - (cache_config.layer_num, max_request_num + 1, self.swa_capacity_tokens, cache_config.swa_slot_bytes), - dtype=torch.uint8, - device="cuda", - ) - self.req_to_c4_state = self._alloc_compressor_state(*c4_shape) - self.req_to_c4_indexer_state = self._alloc_compressor_state(*indexer_shape) - self.req_to_c128_state = self._alloc_compressor_state(*c128_shape) - - def _alloc_compressor_state(self, layer_count, token_capacity, state_dim): - buffer = torch.zeros( - (layer_count, self.max_request_num + 1, token_capacity, state_dim), dtype=torch.float32, device="cuda" - ) - buffer[:, self.HOLD_REQUEST_ID, :, state_dim // 2 :].fill_(float("-inf")) - return buffer - - def get_att_input_params(self, layer_index): - from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_compressed_kv_buffer, view_swa_kv - - return view_swa_kv(self.req_to_swa[layer_index], self.cache_config), get_compressed_kv_buffer( - self.mem_manager, layer_index - ) - - def create_small_page_cache_manager(self, size): - self.small_page_buffers = DeepseekV4StateCacheManager(size, self.cache_config) - return self.small_page_buffers - - def save_state(self, req, buffer_idx, state_cache_manager): - from lightllm.models.deepseek_v4.triton_kernel.swa_cache import gather_swa_state - - config = self.cache_config - swa, c4, indexer = state_cache_manager.get_state_cache(buffer_idx) - # 保存时 cur_kv_len 还是 chunk 起点；运行态对应本轮 prefill 终点。 - req_idx = req.req_idx - checkpoint_len = req.get_chuncked_input_token_len() - positions = torch.arange(checkpoint_len - config.sliding_window, checkpoint_len, device="cuda") - swa.copy_(gather_swa_state(self, req_idx, positions), non_blocking=True) - if not config.n_c4: - return - positions = torch.arange(checkpoint_len - config.short_compress_ratio, checkpoint_len, device="cuda") - state_slots = positions % self.c4_state_tokens - # compressor 已将 chunk 尾部写回 ring；大小页只保存当前终点。 - for target, runtime in ((c4, self.req_to_c4_state), (indexer, self.req_to_c4_indexer_state)): - target.copy_(runtime[:, req_idx].index_select(1, state_slots), non_blocking=True) - - def restore_state(self, req, state_cache_manager, buffer_idx): - from lightllm.models.deepseek_v4.triton_kernel.swa_cache import scatter_swa_state - - config = self.cache_config - swa, c4, indexer = state_cache_manager.get_state_cache(buffer_idx) - positions = torch.arange(req.cur_kv_len - config.sliding_window, req.cur_kv_len, device="cuda") - scatter_swa_state(self, req.req_idx, positions, swa.cuda(non_blocking=True)) - if config.n_c4: - positions = torch.arange(req.cur_kv_len - config.short_compress_ratio, req.cur_kv_len, device="cuda") - slots = positions % self.c4_state_tokens - self.req_to_c4_state[:, req.req_idx].index_copy_(1, slots, c4.cuda(non_blocking=True)) - self.req_to_c4_indexer_state[:, req.req_idx].index_copy_(1, slots, indexer.cuda(non_blocking=True)) - # Aligned checkpoints close C128 groups; the next group overwrites its tail. diff --git a/lightllm/common/req_manager/glm5_next.py b/lightllm/common/req_manager/glm5_next.py deleted file mode 100644 index 492dfeb73c..0000000000 --- a/lightllm/common/req_manager/glm5_next.py +++ /dev/null @@ -1,35 +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, - # Raw keys/gates are computation state, even when packed KV is FP8. - dtype=linear_config.conv_state_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 1f78003317..0000000000 --- a/lightllm/common/req_manager/hybrid_base.py +++ /dev/null @@ -1,69 +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 数。""" - - 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, reqs: List["InferReq"], buffer_indexes: List[int]): - """批量保存请求 GPU 运行态到已分配的大页槽位，buffer_indexes 中的 -1 表示跳过。 - - b_req_idx 与 reqs 分别为同一批请求的 GPU 索引张量和 CPU 请求对象列表。 - 默认逐请求保存，模型可覆盖为批量拷贝算子。 - """ - for req, buffer_idx in zip(reqs, buffer_indexes): - if buffer_idx != -1: - self.save_state(req, buffer_idx, self.big_page_buffers) - - @abstractmethod - def save_state(self, req: "InferReq", buffer_idx: int, state_cache_manager): - """保存本轮 prefill 终点的运行态；在更新 req.cur_kv_len 前调用，大小页共用。""" - - 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 5ae9f4966c..0000000000 --- a/lightllm/common/req_manager/linear_att.py +++ /dev/null @@ -1,119 +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, reqs: List["InferReq"], buffer_indexes: List[int]): - from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer - - buffer_indexes = 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: "InferReq", buffer_idx: int, state_cache_manager: LinearAttCacheManager): - # checkpoint 只保存标准 conv 窗口和请求的基准 SSM 状态，不包含 MTP 扩展运行态。 - req_idx = req.req_idx - 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): - layer_idx_in_linear = self.linear_config.get_linear_state_layer_index(layer_idx_in_all) - 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/__init__.py b/lightllm/common/state_cache_manager/__init__.py deleted file mode 100644 index 758c6a88d4..0000000000 --- a/lightllm/common/state_cache_manager/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -from .base import StateCacheManager -from .deepseek4 import DeepseekV4CacheConfig, DeepseekV4StateCacheManager -from .layer_cache import LayerCache -from .linear_att import LinearAttCacheConfig, LinearAttCacheManager -from .glm5_next import Glm5NextCacheConfig - - -def get_hybrid_cache_config(): - """Return the model-specific layout used by hybrid CPU/disk cache pages.""" - from transformers.configuration_utils import PretrainedConfig - from lightllm.utils.envs_utils import get_env_start_args - - args = get_env_start_args() - model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir) - model_type = model_cfg["model_type"] - if model_type == "deepseek_v4": - return DeepseekV4CacheConfig.from_model_config(model_cfg, args) - if model_type in ("glm5_next", "glm5_next_text"): - return Glm5NextCacheConfig.from_model_config(model_cfg, args) - if model_type in ("qwen3_5", "sensenova_loop", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): - return LinearAttCacheConfig.load_from_args() - raise ValueError("No hybrid state-cache layout registered for this model") diff --git a/lightllm/common/state_cache_manager/base.py b/lightllm/common/state_cache_manager/base.py deleted file mode 100644 index aa131fe784..0000000000 --- a/lightllm/common/state_cache_manager/base.py +++ /dev/null @@ -1,57 +0,0 @@ -import collections -from abc import ABC, abstractmethod -from typing import List, Optional - -from lightllm.utils.log_utils import init_logger - - -logger = init_logger(__name__) - - -class StateCacheManager(ABC): - """CPU checkpoint 槽位池；大小页分别实例化，具体状态布局由子类负责。 - - 尾部 keep_num 个槽位保留给 CPU cache 碎页传输，不参与普通分配。 - 本类不管理 GPU 运行态，也不判断页面边界、前缀匹配或淘汰策略。 - """ - - def __init__(self, size: int, keep_num: int = 0): - self.size = size - self.keep_num = keep_num - assert 0 <= keep_num <= size, f"invalid keep_num {keep_num} for size {size}" - self.free_list = collections.deque(range(size - keep_num)) - - @abstractmethod - def get_state_cache(self, buffer_idx: int): - """返回指定槽位的状态视图；可以是单个 Tensor 或多个 Tensor。""" - - def alloc_one_state_cache(self) -> Optional[int]: - return None if not self.free_list else self.free_list.popleft() - - def alloc_state_cache(self, need_size: int) -> Optional[List[int]]: - if need_size > len(self.free_list): - logger.error(f"warn no enough cache need_size {need_size} free_size {len(self.free_list)}") - return None - return [self.free_list.popleft() for _ in range(need_size)] - - def free_state_cache(self, free_indexes: List[int]): - alloc_upper_bound = self.size - self.keep_num - for idx in free_indexes: - assert 0 <= idx < alloc_upper_bound, ( - f"free index {idx} out of alloc range [0, {alloc_upper_bound}), " f"reserved tail num {self.keep_num}" - ) - self.free_list.extend(free_indexes) - assert ( - len(self.free_list) <= alloc_upper_bound - ), f"free cache num {len(self.free_list)} should not be larger than alloc size {alloc_upper_bound}" - - def get_free_cache_num(self): - return len(self.free_list) - - def get_used_cache_num(self): - # Preserve the existing accounting: reserved slots count as used. - return self.size - len(self.free_list) - - def clear_to_init_state(self): - """重置空闲槽位；子类同时清零自身的 checkpoint buffer。""" - self.free_list = collections.deque(range(self.size - self.keep_num)) diff --git a/lightllm/common/state_cache_manager/deepseek4.py b/lightllm/common/state_cache_manager/deepseek4.py deleted file mode 100644 index 368bd8e222..0000000000 --- a/lightllm/common/state_cache_manager/deepseek4.py +++ /dev/null @@ -1,146 +0,0 @@ -import math -from dataclasses import dataclass - -import torch -import triton - -from lightllm.utils.envs_utils import get_env_start_args - -from .base import StateCacheManager - - -DSV4_COMPRESSED_KV_TOKEN_ALIGNMENT = 256 - - -@dataclass -class DeepseekV4CacheConfig: - """V4 KV 布局和大小页状态缓存配置。""" - - compress_rates: tuple[int, ...] - head_dim: int - rope_head_dim: int - indexer_head_dim: int - sliding_window: int - - # Packed FP8 KV format shared by Flash and Pro. - quant_group_dim = 64 - scale_alignment_bytes = 8 - # 每个 SWA 物理页的 token 数，用于 KV/scale 打包和 ring 容量对齐。 - swa_page_token_num = 128 - compressed_kv_token_alignment = DSV4_COMPRESSED_KV_TOKEN_ALIGNMENT - short_compress_ratio = 4 - long_compress_ratio = 128 - - def __post_init__(self): - self.compress_rates = tuple(self.compress_rates) - self.layer_num = len(self.compress_rates) - c4_layers = [layer for layer, rate in enumerate(self.compress_rates) if rate == self.short_compress_ratio] - c128_layers = [layer for layer, rate in enumerate(self.compress_rates) if rate == self.long_compress_ratio] - self.layer_to_c4_idx = {layer: slot for slot, layer in enumerate(c4_layers)} - self.layer_to_c128_idx = {layer: slot for slot, layer in enumerate(c128_layers)} - self.n_c4, self.n_c128 = len(self.layer_to_c4_idx), len(self.layer_to_c128_idx) - compressed_kv_layers = [layer for layer, rate in enumerate(self.compress_rates) if rate] - self.layer_to_compressed_kv_idx = {layer: slot for slot, layer in enumerate(compressed_kv_layers)} - self.kv_data_bytes = self.head_dim + self.rope_head_dim - self.kv_scale_bytes = ( - triton.cdiv((self.head_dim - self.rope_head_dim) // self.quant_group_dim, self.scale_alignment_bytes) - * self.scale_alignment_bytes - ) - self.kv_bytes = self.kv_data_bytes + self.kv_scale_bytes - self.indexer_bytes = self.indexer_head_dim + torch.float32.itemsize - self.swa_slot_bytes = self.get_packed_kv_bytes(self.swa_page_token_num) // self.swa_page_token_num - self.compressed_kv_slot_bytes = self._get_compressed_kv_slot_bytes() - - @classmethod - def from_model_config(cls, config, args): - from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num - - layer_num = config["num_hidden_layers"] - compress_rates = list(config["compress_ratios"]) - if args.mtp_mode is not None: - layer_num += get_added_mtp_kv_layer_num() - # Runtime-added draft layers use sliding-window attention. - compress_rates.extend([0] * (layer_num - len(compress_rates))) - return cls( - compress_rates=tuple(compress_rates[:layer_num]), - head_dim=config["head_dim"], - rope_head_dim=config["qk_rope_head_dim"], - indexer_head_dim=config["index_head_dim"], - sliding_window=config["sliding_window"], - ) - - def get_packed_kv_bytes(self, token_count): - return triton.cdiv(token_count * self.kv_bytes, self.kv_data_bytes) * self.kv_data_bytes - - def _get_compressed_kv_slot_bytes(self): - block_token_num = self.compressed_kv_token_alignment - # 整块按 KV 数据行对齐，同时能均分到 token 槽位。 - block_bytes_alignment = math.lcm(block_token_num, self.kv_data_bytes) - slot_bytes_list = [] - for ratio in self.compress_rates: - if ratio == 0: - continue - compressed_token_num = block_token_num // ratio - block_bytes = self.get_packed_kv_bytes(compressed_token_num) - if ratio == self.short_compress_ratio: - block_bytes += compressed_token_num * self.indexer_bytes - block_bytes = triton.cdiv(block_bytes, block_bytes_alignment) * block_bytes_alignment - slot_bytes_list.append(block_bytes // block_token_num) - return tuple(slot_bytes_list) - - def get_swa_state_shape(self): - return (self.layer_num, self.sliding_window, self.kv_bytes) - - def get_c4_state_shape(self): - return (self.n_c4, self.short_compress_ratio, 4 * self.head_dim) - - def get_indexer_state_shape(self): - return (self.n_c4, self.short_compress_ratio, 4 * self.indexer_head_dim) - - def get_runtime_state_shapes(self, mtp_step): - """每个请求的完整 ring；P/D 使用相同布局，可直接复制槽位。""" - swa_tokens = triton.cdiv(self.sliding_window + 2 * mtp_step, self.swa_page_token_num) * self.swa_page_token_num - c4_tokens = 2 * self.short_compress_ratio + mtp_step - c128_tokens = self.long_compress_ratio + mtp_step - return ( - (self.layer_num, swa_tokens, self.swa_slot_bytes), - (self.n_c4, c4_tokens, 4 * self.head_dim), - (self.n_c4, c4_tokens, 4 * self.indexer_head_dim), - (self.n_c128, c128_tokens, 2 * self.head_dim), - ) - - def get_cpu_cache_full_att_bytes(self): - args = get_env_start_args() - return args.cpu_cache_token_page_size * sum(self.compressed_kv_slot_bytes) - - def get_cpu_cache_big_page_bytes(self): - # MLA KV and compressor state are replicated across TP ranks. - swa_bytes = math.prod(self.get_swa_state_shape()) - c4_bytes = math.prod(self.get_c4_state_shape()) * torch.float32.itemsize - indexer_bytes = math.prod(self.get_indexer_state_shape()) * torch.float32.itemsize - return self.get_cpu_cache_full_att_bytes() + swa_bytes + c4_bytes + indexer_bytes - - -class DeepseekV4StateCacheManager(StateCacheManager): - """Pinned continuation checkpoints; compressed KV stays in token pages.""" - - def __init__(self, size, cache_config: DeepseekV4CacheConfig, keep_num=0): - super().__init__(size, keep_num) - self.cache_config = cache_config - self.swa_shape = cache_config.get_swa_state_shape() - self.c4_shape = cache_config.get_c4_state_shape() - self.indexer_shape = cache_config.get_indexer_state_shape() - self.swa_nbytes = math.prod(self.swa_shape) - self.c4_nbytes = math.prod(self.c4_shape) * torch.float32.itemsize - indexer_nbytes = math.prod(self.indexer_shape) * torch.float32.itemsize - self.buffer = torch.empty( - (size, self.swa_nbytes + self.c4_nbytes + indexer_nbytes), dtype=torch.uint8, pin_memory=True - ) - - def get_state_cache(self, buffer_idx): - data = self.buffer[buffer_idx] - swa = data[: self.swa_nbytes].view(self.swa_shape) - indexer_start = self.swa_nbytes + self.c4_nbytes - c4 = data[self.swa_nbytes : indexer_start].view(torch.float32).view(self.c4_shape) - indexer = data[indexer_start:].view(torch.float32).view(self.indexer_shape) - return swa, c4, indexer diff --git a/lightllm/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py deleted file mode 100644 index f0928c2c18..0000000000 --- a/lightllm/common/state_cache_manager/glm5_next.py +++ /dev/null @@ -1,64 +0,0 @@ -import dataclasses -import torch - -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args -from lightllm.utils.torch_dtype_utils import get_torch_dtype - -from .linear_att import LinearAttCacheConfig - - -@dataclasses.dataclass -class Glm5NextCacheConfig(LinearAttCacheConfig): - """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" - - index_kpool: int = 4 - index_head_dim: int = 128 - - INDEXER_K_BYTES = 132 - # FlashMLA V4 layout: 448 FP8 values, 64 BF16 values, seven E8M0 scales and one padding byte. - FP8_MLA_KV_BYTES = 584 - - @classmethod - def from_model_config(cls, config, args): - config = config.get("text_config", config) - linear = config["linear_attn_config"] - tp = args.tp // args.dp - assert linear["num_heads"] % tp == 0 - dtype = get_torch_dtype(args.data_type) - use_fp8_kv = args.llm_kv_type == "fp8kv_dsa" - full_att_dtype = torch.uint8 if use_fp8_kv else dtype - mla_kv_bytes = cls.FP8_MLA_KV_BYTES if use_fp8_kv else config["kv_lora_rank"] * dtype.itemsize - # FlashMLA loads 16-byte vectors; keep every packed token row aligned. - packed_kv_bytes = ((mla_kv_bytes + cls.INDEXER_K_BYTES + 15) // 16) * 16 - # Raw index keys and compression scores only live in request tails. - return cls( - tp_world_size=tp, - full_att_all_num_kv_heads=1, - full_att_dtype=full_att_dtype, - full_att_num_kv_heads=1, - full_att_head_dim=packed_kv_bytes // full_att_dtype.itemsize, - global_linear_k_heads=linear["num_heads"], - global_linear_v_heads=linear["num_heads"], - num_linear_k_heads=linear["num_heads"] // tp, - num_linear_v_heads=linear["num_heads"] // tp, - head_linear_k_dim=linear["head_dim"], - head_linear_v_dim=linear["head_dim"], - conv_kernel_size=linear["short_conv_kernel_size"], - conv_state_dtype=dtype, - ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), - layer_types=config["layer_types"], - draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if args.mtp_mode is not None else 0, - index_kpool=config["index_kpool"], - index_head_dim=config["index_head_dim"], - ) - - def get_cpu_cache_full_att_bytes(self): - args = get_env_start_args() - page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num - assert page_tokens == args.cpu_cache_token_page_size - return ( - self.full_att_head_dim - * self.full_att_dtype.itemsize - * self.get_full_att_kv_layer_num_with_draft_model() - * page_tokens - ) diff --git a/lightllm/common/state_cache_manager/layer_cache.py b/lightllm/common/state_cache_manager/layer_cache.py deleted file mode 100644 index 9ac72acca3..0000000000 --- a/lightllm/common/state_cache_manager/layer_cache.py +++ /dev/null @@ -1,37 +0,0 @@ -from typing import Tuple -import torch -import numpy as np -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -class LayerCache: - def __init__( - self, - size: int, - dtype: torch.dtype, - shape: Tuple[int, ...], - layer_num: int, - device: torch.device, - size_first: bool = False, - ): - self.size_first = size_first - self.size = size - self.dtype = dtype - self.shape = shape - self.layer_num = layer_num - self.device = device - if not self.size_first: - if device == "cpu": - self.buffer = torch.zeros((self.layer_num, size, *shape), dtype=dtype, device="cpu", pin_memory=True) - else: - self.buffer = torch.zeros((self.layer_num, size, *shape), dtype=dtype, device=device) - else: - if device == "cpu": - self.buffer = torch.zeros((size, self.layer_num, *shape), dtype=dtype, device="cpu", pin_memory=True) - else: - self.buffer = torch.zeros((size, self.layer_num, *shape), dtype=dtype, device=device) - - def get_cell_size(self): - return np.prod(self.shape) * self.layer_num * torch._utils._element_size(self.dtype) diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py deleted file mode 100644 index 496791d489..0000000000 --- a/lightllm/common/state_cache_manager/linear_att.py +++ /dev/null @@ -1,198 +0,0 @@ -import torch -import dataclasses -import triton - -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args -from lightllm.utils.log_utils import init_logger -from lightllm.utils.torch_dtype_utils import get_torch_dtype - -from .base import StateCacheManager -from .layer_cache import LayerCache - - -logger = init_logger(__name__) - - -@dataclasses.dataclass -class LinearAttCacheConfig: - tp_world_size: int - # full att 的参数 - full_att_all_num_kv_heads: int - full_att_dtype: torch.dtype - full_att_num_kv_heads: int # 这个是 tp 后的head头数量 - full_att_head_dim: int - - # linear att 的参数 - global_linear_k_heads: int - global_linear_v_heads: int - num_linear_k_heads: int - num_linear_v_heads: int - head_linear_k_dim: int - head_linear_v_dim: int - conv_kernel_size: int - conv_state_dtype: torch.dtype - ssm_state_dtype: torch.dtype - layer_types: tuple[str, ...] - draft_full_att_kv_layer_num: int = 0 - linear_layer_num: int = dataclasses.field(init=False) - all_layer_num: int = dataclasses.field(init=False) - - def __post_init__(self): - self.layer_types = tuple(self.layer_types) - indices = [] - full_count = linear_count = 0 - for kind in self.layer_types: - if kind == "linear_attention": - indices.append(linear_count) - linear_count += 1 - else: - indices.append(full_count) - full_count += 1 - # MTP draft layers only use full-attention KV slots after the main model. - indices.extend(range(full_count, full_count + self.draft_full_att_kv_layer_num)) - self._layer_indices = tuple(indices) - self.linear_layer_num = linear_count - self.all_layer_num = len(self.layer_types) - - def get_linear_state_layer_index(self, layer_index: int) -> int: - return self._layer_indices[layer_index] - - def get_conv_dim(self): - # 第一项对应q的参数，第二项对应k的参数，第三项对应v的参数 - # 由于 k_dim = q_dim, k_heads = q_heads, 所以第一项和第二项的计算 - # 形式相同，但是实际内在含义是不同的。 - return ( - self.head_linear_k_dim * self.num_linear_k_heads - + self.head_linear_k_dim * self.num_linear_k_heads - + self.head_linear_v_dim * self.num_linear_v_heads - ) - - def get_main_model_full_att_layer_num(self): - full_att_layer_num = self.all_layer_num - self.linear_layer_num - return full_att_layer_num - - def get_full_att_kv_layer_num_with_draft_model(self): - return self.get_main_model_full_att_layer_num() + self.draft_full_att_kv_layer_num - - 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.""" - - return self._layer_indices[layer_index] - - 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) - - def get_mtp_conv_state_shape(self, mtp_step: int): - # Working state with room for S speculative tokens before acceptance. - return (self.get_conv_dim(), (self.conv_kernel_size - 1) + mtp_step) - - def get_ssm_state_shape(self): - return (self.num_linear_v_heads, self.head_linear_k_dim, self.head_linear_v_dim) - - def get_conv_state_bytes_per_layer(self): - return self.get_conv_dim() * (self.conv_kernel_size - 1) * self.conv_state_dtype.itemsize - - def get_ssm_state_bytes_per_layer(self): - return self.num_linear_v_heads * self.head_linear_k_dim * self.head_linear_v_dim * self.ssm_state_dtype.itemsize - - def get_cpu_cache_big_page_bytes(self): - a = self.get_cpu_cache_full_att_bytes() - b = self.get_cpu_cache_conv_bytes() - c = self.get_cpu_cache_ssm_bytes() - - return triton.cdiv(a + b + c, 16) * 16 - - def get_cpu_cache_full_att_bytes(self): - big_page_token_num = ( - get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size - ) - assert big_page_token_num == get_env_start_args().cpu_cache_token_page_size - full_att_bytes = 2 * self.full_att_all_num_kv_heads * self.full_att_head_dim * self.full_att_dtype.itemsize - a = full_att_bytes * self.get_full_att_kv_layer_num_with_draft_model() * big_page_token_num - return a - - def get_cpu_cache_conv_bytes(self): - b = self.get_conv_state_bytes_per_layer() * self.linear_layer_num * self.tp_world_size - return b - - def get_cpu_cache_ssm_bytes(self): - c = self.get_ssm_state_bytes_per_layer() * self.linear_layer_num * self.tp_world_size - return c - - @staticmethod - def load_from_args() -> "LinearAttCacheConfig": - args = get_env_start_args() - model_path = args.model_dir - from transformers.configuration_utils import PretrainedConfig - - model_cfg, _ = PretrainedConfig.get_config_dict(model_path) - llm_config = model_cfg.get("text_config", model_cfg) - layer_types = llm_config["layer_types"] - if model_cfg["model_type"] == "sensenova_loop": - from lightllm.models.sensenova_loop.config import get_loop_layer_types - - layer_types = get_loop_layer_types(model_cfg["loop_settings"], llm_config) - 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) - ) - return LinearAttCacheConfig( - tp_world_size=tp_world_size, - full_att_all_num_kv_heads=llm_config["num_key_value_heads"], - full_att_dtype=full_att_dtype, - full_att_num_kv_heads=max(1, llm_config["num_key_value_heads"] // tp_world_size), - full_att_head_dim=llm_config["head_dim"], - global_linear_k_heads=llm_config["linear_num_key_heads"], - global_linear_v_heads=llm_config["linear_num_value_heads"], - num_linear_k_heads=max(1, llm_config["linear_num_key_heads"] // tp_world_size), - num_linear_v_heads=max(1, llm_config["linear_num_value_heads"] // tp_world_size), - head_linear_k_dim=llm_config["linear_key_head_dim"], - head_linear_v_dim=llm_config["linear_value_head_dim"], - conv_kernel_size=llm_config["linear_conv_kernel_dim"], - conv_state_dtype=get_torch_dtype(args.data_type), - ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), - layer_types=layer_types, - draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num(), - ) - - -class LinearAttCacheManager(StateCacheManager): - """CPU pinned conv/SSM checkpoint，两个 buffer 均保持 size-first 布局。""" - - def __init__( - self, - size: int, - linear_config: LinearAttCacheConfig, - keep_num: int = 0, # 用于记录需要保留的缓存数量，用于支持含有 linear_att 的如qwen3.5 模型的cpu cache的碎页处理。 - ): - super().__init__(size, keep_num) - self.linear_config = linear_config - # init the layer cache - self.conv_state_cache = LayerCache( - size=self.size, - dtype=self.linear_config.conv_state_dtype, - shape=self.linear_config.get_conv_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cpu", - size_first=True, - ) - self.ssm_state_cache = LayerCache( - size=self.size, - dtype=self.linear_config.ssm_state_dtype, - shape=self.linear_config.get_ssm_state_shape(), - layer_num=self.linear_config.linear_layer_num, - device="cpu", - size_first=True, - ) - self.clear_to_init_state() - return - - def get_state_cache(self, buffer_idx: int): - return self.conv_state_cache.buffer[buffer_idx, ...], self.ssm_state_cache.buffer[buffer_idx, ...] - - def clear_to_init_state(self): - self.conv_state_cache.buffer.zero_() - self.ssm_state_cache.buffer.zero_() - super().clear_to_init_state() - return diff --git a/lightllm/distributed/flashinfer_all_reduce.py b/lightllm/distributed/flashinfer_all_reduce.py index 35cb0aaedf..1029795fd5 100644 --- a/lightllm/distributed/flashinfer_all_reduce.py +++ b/lightllm/distributed/flashinfer_all_reduce.py @@ -102,6 +102,9 @@ def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool: hidden_dim=hidden_dim, dtype=dtype, comm_backend=TorchDistBackend(group=self.group), + # TRT-LLM symmetric-memory rendezvous must use the same TP + # subgroup as world_size/rank; otherwise it defaults to WORLD. + group=self.group, ) except Exception as e: logger.warning("FlashInferAllReduce workspace init failed: %s. Disabling.", e) diff --git a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py index c8a78e6b62..5ec00d542a 100755 --- a/lightllm/models/bloom/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/bloom/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.full_att import FullAttentionState import torch from typing import Tuple from lightllm.common.basemodel import TransformerLayerInferTpl @@ -28,7 +29,13 @@ def _context_attention_kernel( layer_weight: BloomTransformerLayerWeight, out=None, ) -> torch.Tensor: - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, @@ -43,7 +50,13 @@ def _context_attention_kernel( def _token_attention_kernel( self, q: torch.Tensor, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight, out=None ) -> torch.Tensor: - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( q=_q, diff --git a/lightllm/models/bloom/model.py b/lightllm/models/bloom/model.py index 7368ed1dc3..8e963d7f84 100644 --- a/lightllm/models/bloom/model.py +++ b/lightllm/models/bloom/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING from lightllm.models.bloom.layer_infer.pre_layer_infer import BloomPreLayerInfer from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer from lightllm.models.bloom.layer_infer.transformer_layer_infer import BloomTransformerLayerInfer @@ -7,6 +8,12 @@ from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class BloomTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = BloomPreAndPostLayerWeight @@ -24,6 +31,39 @@ def __init__(self, kvargs): super().__init__(kvargs) return + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", + ) + return [state] + def _init_config(self): super()._init_config() # rename key diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 3254031056..986b8a703f 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.full_att import FullAttentionState import os import torch import triton @@ -103,7 +104,13 @@ def _token_attention_kernel( ): q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :] q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1) - kv = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + kv = state.get_layer_view(logical_layer_id=self.layer_num_) out = infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), @@ -119,7 +126,13 @@ def _decompress_kv( infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, ): - compressed_kv = infer_state.mem_manager.kv_buffer[self.layer_num_] + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + compressed_kv = state.get_layer_view(logical_layer_id=self.layer_num_) total_token_num = infer_state.total_token_num sampled_compressed_kv_nope = self.alloc_tensor( diff --git a/lightllm/models/deepseek2/model.py b/lightllm/models/deepseek2/model.py index 26a7ecd81f..e06c18ea16 100644 --- a/lightllm/models/deepseek2/model.py +++ b/lightllm/models/deepseek2/model.py @@ -1,18 +1,24 @@ +from typing import TYPE_CHECKING import torch from typing import final from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer 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.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.utils.envs_utils import enable_env_vars, get_env_start_args from lightllm.distributed.communication_op import dist_group_manager from lightllm.common.basemodel.attention import get_mla_decode_att_backend_class, get_mla_prefill_att_backend_class logger = init_logger(__name__) +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class Deepseek2TpPartModel(LlamaTpPartModel): # weight class transformer_weight_class = Deepseek2TransformerLayerWeight @@ -57,18 +63,36 @@ def _init_custom(self): def _verify_params(self): return super()._verify_params() - def _init_mem_manager(self): - manager_class = select_mem_manager_class() - - self.mem_manager = manager_class( - self.max_total_token_num, - dtype=self.data_type, - head_num=1, - head_dim=self.config["kv_lora_rank"] + self.config["qk_rope_head_dim"], - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """DeepSeek MLA 保存 latent 与 RoPE，TP 副本规则由 MLA 创建入口处理。""" + from lightllm.common.attention_state.full_att import MLAConfig, create_mla_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_mla_state( + config=MLAConfig( + logical_layer_groups=logical_layer_groups, + latent_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], + latent_dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + page_size=args.page_size, + device="meta", ) - return + return [state] def _init_to_get_yarn_rotary(self): from lightllm.models.llama.yarn_rotary_utils import find_correction_range, linear_ramp_mask, get_deepseek_mscale diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py index b9be73e278..795f75459f 100644 --- a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py +++ b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_fp8.py @@ -24,7 +24,7 @@ def gqa_token_decode_attention_flash_decoding_fp8( softmax_scale, out=None, alloc_tensor_func=torch.empty, - **run_config + **run_config, ): batch_size = infer_state.batch_size max_kv_seq_len = infer_state.max_kv_seq_len @@ -73,7 +73,7 @@ def gqa_token_decode_attention_flash_decoding_fp8( mid_o_logexpsum, softmax_scale, get_sm_count=True, - **run_config + **run_config, ) if not hasattr(infer_state, "decode_att_block_seq"): @@ -122,7 +122,7 @@ def gqa_token_decode_attention_flash_decoding_fp8( mid_o_logexpsum, softmax_scale, get_sm_count=False, - **run_config + **run_config, ) flash_decode_stage2( @@ -132,7 +132,7 @@ def gqa_token_decode_attention_flash_decoding_fp8( mid_o_logexpsum, infer_state.b_seq_len, o_tensor.view(calcu_shape1), - **run_config + **run_config, ) return o_tensor @@ -194,7 +194,7 @@ def _fwd_kernel_calcu_index_and_block_seq( infer_state.batch_size = Z infer_state.max_kv_seq_len = N_CTX infer_state.total_token_num = Z * N_CTX - infer_state.req_manager = ReqManager(Z, N_CTX, None) + infer_state.req_manager = ReqManager(Z, N_CTX) infer_state.req_manager.req_to_token_indexs = req_to_token_indexs infer_state.b_req_idx = b_req_idx infer_state.b_seq_len = b_seq_len diff --git a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py index d6eaebe2fd..e026d1ea8e 100644 --- a/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek3_2/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.full_att import FullAttentionState import torch from typing import Any from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo @@ -99,9 +100,15 @@ def _context_attention_kernel( }, ) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) mla_out = infer_state.prefill_att_state.prefill_att( q=q_all, - k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, v=None, att_control=att_control, ) @@ -141,9 +148,15 @@ def _token_attention_kernel( }, ) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) o_tensor = infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), - k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, v=None, att_control=att_control, ) @@ -196,11 +209,17 @@ def _get_indices( q_fp8, q_scale = act_quant(q, self.block_size, self.scale_fmt) k_fp8, k_scale = act_quant(k, self.block_size, self.scale_fmt) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", + logical_layer_id=self.layer_idx_, + state_class=FullAttentionState, + ) destindex_copy_indexer_ks( K_fp8=k_fp8, K_scale=k_scale, DestLoc=infer_state.mem_index, - O_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), + O_buffer=state.get_layer_view(logical_layer_id=self.layer_idx_).indexer_k, ) weights = layer_weight.weights_proj_.mm(hidden_states) * self.index_n_heads_scale @@ -216,7 +235,7 @@ def _get_indices( mtp_step = get_env_start_args().mtp_step # Use efficient Triton kernel to extract FP8 keys and scales from buffer k_fp8_, k_scale_ = extract_indexer_ks( - I_buffer=infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx_), + I_buffer=state.get_layer_view(logical_layer_id=self.layer_idx_).indexer_k, b_seq_len=infer_state.b_seq_len, b_req_idx=infer_state.b_req_idx, req_to_token_indexs=infer_state.req_manager.req_to_token_indexs, diff --git a/lightllm/models/deepseek3_2/model.py b/lightllm/models/deepseek3_2/model.py index 57b6739e6f..b691d986cc 100644 --- a/lightllm/models/deepseek3_2/model.py +++ b/lightllm/models/deepseek3_2/model.py @@ -1,9 +1,17 @@ +from typing import TYPE_CHECKING +import torch from lightllm.models.deepseek2.model import Deepseek2TpPartModel from lightllm.models.deepseek3_2.layer_weights.transformer_layer_weight import Deepseek3_2TransformerLayerWeight from lightllm.models.deepseek3_2.layer_infer.transformer_layer_infer import Deepseek3_2TransformerLayerInfer from lightllm.common.basemodel.attention import get_nsa_prefill_att_backend_class, get_nsa_decode_att_backend_class +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class Deepseek3_2TpPartModel(Deepseek2TpPartModel): # weight class @@ -12,6 +20,38 @@ class Deepseek3_2TpPartModel(Deepseek2TpPartModel): # infer class transformer_layer_infer_class = Deepseek3_2TransformerLayerInfer + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """DeepSeek V3.2 使用包含 indexer 字段的 packed MLA，编码由 MLA 创建入口选择。""" + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=logical_layer_groups, + latent_dim=config["kv_lora_rank"] + config["qk_rope_head_dim"], + tp_rank=tp_rank, + tp_world_size=tp_world_size, + format="packed_group128_fp8" if args.llm_kv_type == "fp8kv_dsa" else "packed_bf16", + indexer_dim=config.get("indexer_dim", 128), + ), + page_size=args.page_size, + device="meta", + ) + return [state] + 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/deepseek_mtp/model.py b/lightllm/models/deepseek_mtp/model.py index e2b2a56137..c151d8efcd 100644 --- a/lightllm/models/deepseek_mtp/model.py +++ b/lightllm/models/deepseek_mtp/model.py @@ -28,14 +28,6 @@ def _init_custom(self): self._sin_cached = self.main_model._sin_cached return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None self.pre_post_weight = self.pre_and_post_weight_class( diff --git a/lightllm/models/deepseek_v4/config.py b/lightllm/models/deepseek_v4/config.py new file mode 100644 index 0000000000..a5dafbf108 --- /dev/null +++ b/lightllm/models/deepseek_v4/config.py @@ -0,0 +1,85 @@ +"""模型选择 V4 的逻辑层和压缩方案，State 配置只记录确定的几何参数。""" + + +from dataclasses import dataclass +from typing import ClassVar + +from lightllm.common.attention_state.bounded_att.sliding_window.fp8_e8m0_swa_state import FP8E8M0SlidingWindowConfig +from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import CompressedAttentionConfig + + +@dataclass(frozen=True) +class DSV4AttentionConfig: + # 当前模型拥有的逻辑层 ID，顺序对应 SWA 的物理 layer 轴；main/draft 分别创建。 + logical_layer_ids: tuple[int, ...] + # 与 logical_layer_ids 逐项对应；0 为纯 SWA，4/128 为同时启用对应压缩器。 + compress_ratios: tuple[int, ...] + # 未编码的共享 KV 向量宽度，包含以 BF16 存储的 RoPE 部分。 + head_dim: int + # RoPE 特征宽度；其余特征按每 64 个元素一组编码为 FP8。 + rope_head_dim: int + # C4 indexer key 的特征宽度；每个 key 另有一个 FP32 scale。 + indexer_head_dim: int + # 每份可恢复 SWA 快照包含的逻辑 token 数，不包含 MTP 候选位置。 + sliding_window: int + # 存储在各 TP rank 完整复制，只有 rank 0 向共享的 CPU/PD 区域写出。 + tp_rank: int = 0 + tp_world_size: int = 1 + + # 压缩 KV 按原始 token 的 256-token 块存储，KV 与 indexer 不得跨块拆开。 + compressed_kv_token_alignment: ClassVar[int] = 256 + + def __post_init__(self) -> None: + assert self.logical_layer_ids and len(set(self.logical_layer_ids)) == len(self.logical_layer_ids) + assert len(self.logical_layer_ids) == len(self.compress_ratios) + assert all(ratio in (0, 4, 128) for ratio in self.compress_ratios) + assert self.head_dim > self.rope_head_dim > 0 and (self.head_dim - self.rope_head_dim) % 64 == 0 + assert self.indexer_head_dim > 0 and self.sliding_window > 0 + assert 0 <= self.tp_rank < self.tp_world_size + + def create_swa_config(self) -> FP8E8M0SlidingWindowConfig: + """模型选择窗口及向量维度，具体编码和物理页布局由原子配置负责。""" + return FP8E8M0SlidingWindowConfig( + logical_layer_ids=self.logical_layer_ids, + window_size=self.sliding_window, + kv_dim=self.head_dim, + rope_dim=self.rope_head_dim, + tp_rank=self.tp_rank, + tp_world_size=self.tp_world_size, + ) + + def create_compressed_config(self, compress_ratio: int) -> CompressedAttentionConfig: + """按模型层组合创建一种压缩历史配置，供 State 创建和 backend 使用。""" + return CompressedAttentionConfig( + logical_layer_groups=tuple( + (layer,) + for layer, ratio in zip(self.logical_layer_ids, self.compress_ratios) + if ratio == compress_ratio + ), + compress_ratio=compress_ratio, + kv_dim=self.head_dim, + rope_dim=self.rope_head_dim, + block_token_size=self.compressed_kv_token_alignment, + indexer_dim=self.indexer_head_dim if compress_ratio == 4 else 0, + tp_rank=self.tp_rank, + tp_world_size=self.tp_world_size, + ) + + +def create_attention_config( + model_config: dict[str, object], + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, +) -> DSV4AttentionConfig: + ratios = model_config["compress_ratios"] + return DSV4AttentionConfig( + logical_layer_ids=model_logical_layer_ids, + compress_ratios=tuple(ratios[layer] if layer < len(ratios) else 0 for layer in model_logical_layer_ids), + head_dim=model_config["head_dim"], + rope_head_dim=model_config["qk_rope_head_dim"], + indexer_head_dim=model_config["index_head_dim"], + sliding_window=model_config["sliding_window"], + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) diff --git a/lightllm/models/deepseek_v4/infer_struct.py b/lightllm/models/deepseek_v4/infer_struct.py index 14f751394e..6f90fa1207 100644 --- a/lightllm/models/deepseek_v4/infer_struct.py +++ b/lightllm/models/deepseek_v4/infer_struct.py @@ -1,9 +1,13 @@ import dataclasses +from typing import Literal import torch from lightllm.common.basemodel import InferStateInfo -from lightllm.common.req_manager import DeepseekV4ReqManager -from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager +from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState, WindowCompressorState +from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import ( + CompressedAttentionState, + CompressedAttentionLayerView, +) C4_PREFILL_LOGITS_BUDGET_BYTES = 512 * 1024 * 1024 @@ -23,11 +27,39 @@ class DeepseekV4InferStateInfo(InferStateInfo): """Per-forward inputs, derived indices and attention metadata. Indices and metadata are built for this batch and reused by its layers. - Request-owned SWA and compressor rings remain in the request manager. + SWA and compressor rings are owned by the model's bounded State. """ - req_manager: DeepseekV4ReqManager - mem_manager: DeepseekV4MemoryManager + @property + def swa_state(self) -> FP8E8M0SlidingWindowState: + return self.att_state_manager.get_state( + model_role=self.att_state_role, + state_type_filter="swa_", + logical_layer_id=self.first_logical_layer_id, + state_class=FP8E8M0SlidingWindowState, + ) + + def get_compressor_state( + self, logical_layer_id: int, target: Literal["kv", "indexer"] = "kv" + ) -> WindowCompressorState: + """按当前层和用途选择独立的窗口压缩器。""" + return self.att_state_manager.get_state( + model_role=self.att_state_role, + state_type_filter=f"compressor_{target}_", + logical_layer_id=logical_layer_id, + state_class=WindowCompressorState, + ) + + def get_compressed_kv_state(self, logical_layer_id: int) -> CompressedAttentionState: + return self.att_state_manager.get_state( + model_role=self.att_state_role, + state_type_filter="compressed", + logical_layer_id=logical_layer_id, + state_class=CompressedAttentionState, + ) + + def get_compressed_kv_view(self, logical_layer_id: int) -> CompressedAttentionLayerView: + return self.get_compressed_kv_state(logical_layer_id).get_layer_view(logical_layer_id) def __init__(self): super().__init__() @@ -56,9 +88,21 @@ def __init__(self): self.c4_indexer_metadata: C4IndexerMetadata | None = None def init_some_extra_state(self, model): + self.first_logical_layer_id = model.layers_infer[0].layer_num_ + self.first_c4_logical_layer_id = None + has_c128_layer = False + ratios = model.config["compress_ratios"] + for layer in model.layers_infer: + compress_ratio = ratios[layer.layer_num_] if layer.layer_num_ < len(ratios) else 0 + if compress_ratio == 4 and self.first_c4_logical_layer_id is None: + self.first_c4_logical_layer_id = layer.layer_num_ + if compress_ratio == 128: + has_c128_layer = True self.c4_indexer_metadata = None # Rebuild for each eager forward or graph capture. self.dsv4_c4_indices = None self.dsv4_c4_lengths = None + self.dsv4_c128_indices = None + self.dsv4_c128_lengths = None super().init_some_extra_state(model) # sets position_ids, b_q_seq_len, b_q_start_loc (prefill) if self.is_prefill: self.prefill_att_state.prepare_swa_workspace(self) @@ -84,14 +128,18 @@ def init_some_extra_state(self, model): self.dsv4_workspace = model.dsv4_workspace self._init_swa_metadata(model) + # 当前模型没有 C128 层时，不分配或构建这组索引；draft 只检查自身拥有的层。 + if not has_c128_layer: + return + from lightllm.models.deepseek_v4.triton_kernel.build_compressed_kv_indices import build_compressed_kv_indices - long_ratio = self.req_manager.cache_config.long_compress_ratio + c128_compress_ratio = 128 max_kv_seq_len = self.max_kv_seq_len # Replay copies C128 indices into the captured state, so their widths must match. if not self.is_prefill and model.graph is not None and model.graph.can_run(self.batch_size, max_kv_seq_len): max_kv_seq_len = model.graph.graph_max_len_in_batch - c128_entries = max(1, max_kv_seq_len // long_ratio) + c128_entries = max(1, max_kv_seq_len // c128_compress_ratio) c128_cap = ((c128_entries + 63) // 64) * 64 self.dsv4_c128_indices = torch.empty((pos.numel(), c128_cap), dtype=torch.int32, device=pos.device) self.dsv4_c128_lengths = torch.empty_like(pos, dtype=torch.int32) @@ -99,7 +147,7 @@ def init_some_extra_state(self, model): self.dsv4_sparse_req_idx, self.position_ids, self.req_manager.req_to_token_indexs, - long_ratio, + c128_compress_ratio, self.dsv4_c128_indices, self.dsv4_c128_lengths, ) @@ -113,7 +161,7 @@ def _init_swa_metadata(self, model): pos = self.position_ids has_vision = model.has_vision - swa_width = self.req_manager.cache_config.sliding_window + swa_width = self.swa_state.config.window_size if has_vision and self.is_prefill: swa_width += model.config["vision_max_n_token"] self.dsv4_swa_indices = torch.empty((pos.numel(), swa_width), dtype=torch.int32, device=pos.device) @@ -141,11 +189,11 @@ def _init_swa_metadata(self, model): build_swa_indices( token_req_idx=self.dsv4_sparse_req_idx, position_ids=self.position_ids, - req_manager=self.req_manager, + swa_state=self.swa_state, swa_indices=self.dsv4_swa_indices, swa_lengths=self.dsv4_swa_lengths, swa_write_slots=self.dsv4_swa_write_slots, - sliding_window=self.req_manager.cache_config.sliding_window, + sliding_window=self.swa_state.config.window_size, image_left=self.dsv4_image_left, image_right=self.dsv4_image_right, prefill_state=self.prefill_att_state if self.is_prefill else None, @@ -180,7 +228,7 @@ def get_c4_indexer_metadata(self, c4_cap): self.req_manager.req_to_token_indexs, self.req_manager.HOLD_REQUEST_ID, ) - page_size = self.req_manager.cache_config.compressed_kv_token_alignment // 4 + page_size = self.get_compressed_kv_state(self.first_c4_logical_layer_id).config.block_token_size // 4 rows_per_chunk = max(1, self.position_ids.numel()) if self.is_prefill: # DeepGEMM rounds the logits width to a multiple of 256. diff --git a/lightllm/models/deepseek_v4/layer_infer/compressor.py b/lightllm/models/deepseek_v4/layer_infer/compressor.py index f9841302cf..4938d1f823 100644 --- a/lightllm/models/deepseek_v4/layer_infer/compressor.py +++ b/lightllm/models/deepseek_v4/layer_infer/compressor.py @@ -18,7 +18,9 @@ def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int, is_ self.layer_idx_ = layer_idx self.tp_world_size_ = tp_world_size self.is_in_indexer = is_in_indexer - self.compress_ratio = network_config["compress_ratios"][layer_idx] + ratios = network_config["compress_ratios"] + # 额外的 MTP 执行层只使用 SWA，checkpoint 可只记录一个共享的 MTP 权重层。 + self.compress_ratio = ratios[layer_idx] if layer_idx < len(ratios) else 0 self.head_dim = network_config["index_head_dim" if is_in_indexer else "head_dim"] self.qk_rope_head_dim = network_config["qk_rope_head_dim"] self.eps = network_config["rms_norm_eps"] @@ -92,35 +94,25 @@ def fused_compress( is_in_indexer: bool = False, out_buffer: torch.Tensor = None, ): - mem_manager = infer_state.mem_manager - req_manager = infer_state.req_manager + state = infer_state.get_compressor_state(logical_layer_id=layer_idx, target="indexer" if is_in_indexer else "kv") + req_to_compressor_state = state.get_layer_view(logical_layer_id=layer_idx) compressed_kv_slots = torch.where( (infer_state.position_ids + 1) % compress_ratio == 0, infer_state.mem_index.reshape(-1) // compress_ratio, -1, ) if is_in_indexer: - req_to_compressor_state = req_manager.req_to_c4_indexer_state[ - req_manager.cache_config.layer_to_c4_idx[layer_idx] - ] if out_buffer is None: out_buffer = torch.empty( - (infer_state.mem_index.numel(), req_manager.cache_config.indexer_head_dim), + (infer_state.mem_index.numel(), head_dim), dtype=torch.bfloat16, device=infer_state.mem_index.device, ) out_page_size = 1 else: - if compress_ratio == 4: - req_to_compressor_state = req_manager.req_to_c4_state[req_manager.cache_config.layer_to_c4_idx[layer_idx]] - else: - req_to_compressor_state = req_manager.req_to_c128_state[ - req_manager.cache_config.layer_to_c128_idx[layer_idx] - ] - out_page_size = req_manager.cache_config.compressed_kv_token_alignment // compress_ratio - from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_compressed_kv_buffer - - out_buffer = get_compressed_kv_buffer(mem_manager, layer_idx) + compressed = infer_state.get_compressed_kv_state(logical_layer_id=layer_idx) + out_page_size = compressed.config.block_token_size // compress_ratio + out_buffer = compressed.get_layer_view(logical_layer_id=layer_idx).kv if kv_score.shape[0] == 0: return out_buffer @@ -160,7 +152,7 @@ def fused_compress( CANDIDATE_DIM=candidate_dim, STATE_DIM=state_dim, COMPRESS_RATIO=compress_ratio, - WINDOW_SIZE=compress_ratio * (2 if is_c4 else 1), + WINDOW_SIZE=state.config.window_size, IS_C4=is_c4, IS_PREFILL=infer_state.is_prefill, STATE_CAPACITY_TOKENS=state_capacity_tokens, diff --git a/lightllm/models/deepseek_v4/layer_infer/indexer.py b/lightllm/models/deepseek_v4/layer_infer/indexer.py index cdec83bfdb..177c5529d6 100644 --- a/lightllm/models/deepseek_v4/layer_infer/indexer.py +++ b/lightllm/models/deepseek_v4/layer_infer/indexer.py @@ -18,7 +18,9 @@ class DeepseekV4IndexInfer(BaseLayerInfer): def __init__(self, layer_idx: int, network_config: dict, tp_world_size: int): self.layer_idx_ = layer_idx - self.compress_ratio = network_config["compress_ratios"][layer_idx] + ratios = network_config["compress_ratios"] + # 额外的 MTP 执行层只使用 SWA，checkpoint 可只记录一个共享的 MTP 权重层。 + self.compress_ratio = ratios[layer_idx] if layer_idx < len(ratios) else 0 self.index_topk = network_config["index_topk"] self.index_head_dim = network_config["index_head_dim"] self.index_n_heads = network_config["index_n_heads"] @@ -64,7 +66,6 @@ def write_indexer_k( if use_custom_tensor_manager: hadamard_out = self.alloc_tensor(scratch.shape, scratch.dtype, device=scratch.device) scratch = hadamard_transform(scratch, scale=self.index_head_dim ** -0.5, out=hadamard_out) - from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_indexer_k_buffer from lightllm.models.deepseek_v4.triton_kernel.destindex_copy_indexer_k import ( destindex_copy_indexer_k, ) @@ -73,8 +74,9 @@ def write_indexer_k( scratch, infer_state.mem_index.reshape(-1), infer_state.position_ids, - get_indexer_k_buffer(infer_state.mem_manager, self.layer_idx_), - infer_state.mem_manager.cache_config.compressed_kv_token_alignment // self.compress_ratio, + infer_state.get_compressed_kv_view(logical_layer_id=self.layer_idx_).indexer_k, + infer_state.get_compressed_kv_state(logical_layer_id=self.layer_idx_).config.block_token_size + // self.compress_ratio, ) def build_metadata( @@ -170,11 +172,12 @@ def _c4_indices(self, infer_state: DeepseekV4InferStateInfo, idx_q_fp8, weights) return slots.unsqueeze(1), lengths paged_meta = infer_state.get_c4_indexer_metadata(c4_cap) - mem_manager = infer_state.mem_manager - page_size = mem_manager.cache_config.compressed_kv_token_alignment // self.compress_ratio - from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_indexer_k_buffer + page_size = ( + infer_state.get_compressed_kv_state(logical_layer_id=self.layer_idx_).config.block_token_size + // self.compress_ratio + ) - packed = get_indexer_k_buffer(mem_manager, self.layer_idx_) + packed = infer_state.get_compressed_kv_view(logical_layer_id=self.layer_idx_).indexer_k indexer_k_cache = packed.view( packed.shape[0], page_size, diff --git a/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py index 8a923e97ab..2e09c9d3b3 100644 --- a/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/pre_layer_infer.py @@ -23,6 +23,7 @@ def context_forward(self, input_ids, infer_state: DeepseekV4InferStateInfo, laye def token_forward(self, input_ids, infer_state: DeepseekV4InferStateInfo, layer_weight): input_embdings = super().token_forward(input_ids, infer_state, layer_weight) - if not infer_state.is_cuda_graph: + # 捕获时保留 graph 输入引用；普通 decode 及时释放未使用的 draft 特征。 + if not infer_state.is_cuda_graph_capture: infer_state.mtp_draft_input_hiddens = None return hc_expand(input_embdings, self.hc_mult) diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py index f52cd019fb..c5fe0495ee 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py @@ -30,7 +30,9 @@ def __init__(self, layer_num, network_config): self.hc_mult = network_config["hc_mult"] self.sinkhorn_iters = network_config["hc_sinkhorn_iters"] self.hc_eps = network_config["hc_eps"] - self.compress_ratio = network_config["compress_ratios"][layer_num] + ratios = network_config["compress_ratios"] + # 额外的 MTP 执行层只使用 SWA，checkpoint 可只记录一个共享的 MTP 权重层。 + self.compress_ratio = ratios[layer_num] if layer_num < len(ratios) else 0 self.is_hash = layer_num < network_config["num_hash_layers"] self.has_vision = network_config.get("vision_n_layers", 0) > 0 self.vocab_size = network_config["vocab_size"] @@ -334,7 +336,7 @@ def _select_rope(self, infer_state: DeepseekV4InferStateInfo): def _norm_rope_and_cache_kv(self, kv, infer_state, layer_weight): from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa - req_manager = infer_state.req_manager + state = infer_state.swa_state att_state = infer_state.prefill_att_state if infer_state.is_prefill else infer_state.decode_att_state swa, _ = att_state.backend.get_att_input_params(infer_state, self.layer_num_) # 一次完成 RMSNorm、RoPE、FP8 打包和 SWA 写入。 @@ -347,8 +349,8 @@ def _norm_rope_and_cache_kv(self, kv, infer_state, layer_weight): swa_write_slots=infer_state.dsv4_swa_write_slots, swa_runtime_write_slots=infer_state.dsv4_swa_runtime_write_slots, swa_kv_buffer=swa, - swa_page_token_num=req_manager.cache_config.swa_page_token_num, - req_to_swa=req_manager.req_to_swa[self.layer_num_], + swa_page_token_num=state.config.page_token_size, + req_to_swa=state.get_layer_view(logical_layer_id=self.layer_num_), ) def _get_qkv( diff --git a/lightllm/models/deepseek_v4/model.py b/lightllm/models/deepseek_v4/model.py index 9a6b75cc2d..76758da0cc 100644 --- a/lightllm/models/deepseek_v4/model.py +++ b/lightllm/models/deepseek_v4/model.py @@ -1,10 +1,10 @@ +from __future__ import annotations + import os +from typing import TYPE_CHECKING import torch from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.common.req_manager import DeepseekV4ReqManager -from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager -from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4CacheConfig from lightllm.models.deepseek_v4.layer_weights.pre_and_post_layer_weight import ( DeepseekV4PreAndPostLayerWeight, ) @@ -37,6 +37,12 @@ from lightllm.utils.log_utils import init_logger from lightllm.distributed.communication_op import dist_group_manager +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att.base_state import FullAttentionState + from lightllm.common.attention_state.bounded_att.base_state import BoundedAttentionState + logger = init_logger(__name__) @@ -61,8 +67,7 @@ def add(self, layer_index: int, hidden): class DeepseekV4TpPartModel(LlamaTpPartModel): - req_manager: DeepseekV4ReqManager - mem_manager: DeepseekV4MemoryManager + is_hybrid_att_model = True has_vision = False pre_and_post_weight_class = DeepseekV4PreAndPostLayerWeight @@ -90,30 +95,120 @@ def _verify_params(self): raise ValueError("DeepSeek-V4 preselected MoE routing does not yet support redundant experts") return - def _init_req_manager(self): - cache_config = DeepseekV4CacheConfig.from_model_config(self.config, self.args) - self.req_manager = DeepseekV4ReqManager( - self.max_req_num, - self.max_seq_length, - None, - cache_config, - mtp_step=self.args.mtp_step, + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: StartArgs | EasyDict, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list[FullAttentionState]: + """分别创建 C4/C128 定义；逻辑索引保留原始 token 域，不提前分配存储。""" + from lightllm.common.attention_state.full_att.compressed_att.create_compressed_att_state import ( + create_compressed_att_state, ) + from .config import create_attention_config - def _init_mem_manager(self): - DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(self) - cache_config = self.req_manager.cache_config - self.mem_manager = DeepseekV4MemoryManager( - self.max_total_token_num, - dtype=self.data_type, - head_num=1, - head_dim=cache_config.head_dim, - layer_num=cache_config.layer_num, - mem_fraction=self.mem_fraction, - cache_config=cache_config, + assert args.llm_kv_type in ("None", "fp8kv_dsa"), f"unsupported V4 KV encoding: {args.llm_kv_type}" + config = create_attention_config( + model_config=model_config, + model_logical_layer_ids=model_logical_layer_ids, + tp_rank=tp_rank, + tp_world_size=tp_world_size, ) - self.req_manager.bind_mem_manager(self.mem_manager) - return + states = [] + for ratio in (4, 128): + if ratio not in config.compress_ratios: + continue + states.append( + create_compressed_att_state( + config=config.create_compressed_config(compress_ratio=ratio), + page_size=args.page_size, + device="meta", + ) + ) + return states + + @classmethod + def create_bounded_att_states( + cls, + model_config: dict[str, object], + args: StartArgs | EasyDict, + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", + ) -> list[BoundedAttentionState]: + """创建本模型的 SWA/compressor 请求运行态，大小页容量由管理器后续确定。""" + from lightllm.common.attention_state.bounded_att import ( + WindowCompressorStateConfig, + create_window_compressor_state, + create_sliding_window_state, + ) + from .config import create_attention_config + + config = create_attention_config( + model_config=model_config, + model_logical_layer_ids=model_logical_layer_ids, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + states = [ + create_sliding_window_state( + config=config.create_swa_config(), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + ] + # 模型决定窗口组合：CSA 使用 W8/S4，HCA 使用 W128/S128；State 本身不识别 C4/C128。 + for ratio, window_size in ((4, 8), (128, 128)): + layers = tuple( + layer + for layer, compress_ratio in zip(model_logical_layer_ids, config.compress_ratios) + if compress_ratio == ratio + ) + if not layers: + continue + targets = ( + (("kv", config.head_dim), ("indexer", config.indexer_head_dim)) + if ratio == 4 + else (("kv", config.head_dim),) + ) + for target, head_dim in targets: + states.append( + create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=layers, + window_size=window_size, + compress_stride=ratio, + candidate_dim=head_dim * (2 if window_size > ratio else 1), + target=target, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + ) + return states + + def prepare_attention_states(self) -> None: + super().prepare_attention_states() + # 在统一容量规划前预留长 prefill 的 SWA 临时区，避免 KV 分配挤占计算所需空间。 + DeepseekV4FlashMlaFp8SparseAttBackend.get_swa_prefill_workspace(self) + + def _init_quant(self): + if not self.is_mtp_draft_model: + return super()._init_quant() + from lightllm.common.quantization import Quantcfg + + # 权重量化按 checkpoint 层号查询；draft 的资源层数仍只描述自身。 + config = dict(self.config, n_layer=len(self.config["compress_ratios"])) + self.quant_cfg = Quantcfg(config, self.quant_type, self.quant_cfg_path, self.expert_dtype) def _init_att_backend(self): args = get_env_start_args() diff --git a/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py b/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py index ed2e129396..6a48b024cd 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py +++ b/lightllm/models/deepseek_v4/triton_kernel/build_dspark_swa_indices.py @@ -15,7 +15,7 @@ def _build_dspark_swa_indices_kernel( HOLD_REQUEST_ID: tl.constexpr, HOLD_SWA_SLOT: tl.constexpr, SLIDING_WINDOW: tl.constexpr, - BLOCK_SIZE: tl.constexpr, + block_size, WIDTH: tl.constexpr, BLOCK_W: tl.constexpr, ): @@ -24,14 +24,16 @@ def _build_dspark_swa_indices_kernel( position = tl.load(position_ids + token_idx).to(tl.int64) is_hold = req_idx == HOLD_REQUEST_ID - block_offset = token_idx % BLOCK_SIZE + # draft proposal 长度随请求变化，不参与编译期分支；保持运行时参数避免 + # 每种 MTP 接受长度都生成一份新的 kernel。 + block_offset = token_idx % block_size block_start = tl.maximum(position - block_offset, 0) history_len = tl.minimum(block_start, SLIDING_WINDOW) column = tl.arange(0, BLOCK_W) in_width = column < WIDTH is_history = column < history_len - is_block = (column >= history_len) & (column < history_len + BLOCK_SIZE) + is_block = (column >= history_len) & (column < history_len + block_size) history_position = block_start - 1 - column block_position = block_start + column - history_len source_position = tl.where(is_history, history_position, block_position) @@ -43,7 +45,7 @@ def _build_dspark_swa_indices_kernel( output = tl.where(is_hold, HOLD_SWA_SLOT, output).to(tl.int32) tl.store(swa_indices + token_idx * swa_indices_stride0 + column, output, mask=in_width) - length = tl.where(is_hold, 1, history_len + BLOCK_SIZE).to(tl.int32) + length = tl.where(is_hold, 1, history_len + block_size).to(tl.int32) tl.store(swa_lengths + token_idx, length) write_slot = req_idx * SWA_CAPACITY_TOKENS + position % SWA_CAPACITY_TOKENS write_slot = tl.where(is_hold, HOLD_SWA_SLOT, write_slot).to(tl.int32) @@ -53,14 +55,11 @@ def _build_dspark_swa_indices_kernel( def build_dspark_swa_indices( token_req_idx: torch.Tensor, position_ids: torch.Tensor, - req_manager, + swa_state, swa_indices: torch.Tensor, swa_lengths: torch.Tensor, swa_write_slots: torch.Tensor, - sliding_window: int, block_size: int, - hold_req_id: int, - hold_swa_slot: int, ): """Build ``history SWA + complete draft block`` indices for every DSpark query row.""" @@ -76,11 +75,11 @@ def build_dspark_swa_indices( swa_indices.stride(0), swa_lengths, swa_write_slots, - SWA_CAPACITY_TOKENS=req_manager.swa_capacity_tokens, - HOLD_REQUEST_ID=hold_req_id, - HOLD_SWA_SLOT=hold_swa_slot, - SLIDING_WINDOW=sliding_window, - BLOCK_SIZE=block_size, + SWA_CAPACITY_TOKENS=swa_state.swa_capacity_tokens, + HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, + HOLD_SWA_SLOT=swa_state.HOLD_SWA_SLOT, + SLIDING_WINDOW=swa_state.config.window_size, + block_size=block_size, WIDTH=width, BLOCK_W=triton.next_power_of_2(width), num_warps=4, diff --git a/lightllm/models/deepseek_v4/triton_kernel/build_swa_indices.py b/lightllm/models/deepseek_v4/triton_kernel/build_swa_indices.py index 11704de3c9..f46eaaf1cb 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/build_swa_indices.py +++ b/lightllm/models/deepseek_v4/triton_kernel/build_swa_indices.py @@ -138,7 +138,7 @@ def _build_swa_indices_kernel( def build_swa_indices( token_req_idx: torch.Tensor, position_ids: torch.Tensor, - req_manager, + swa_state, swa_indices: torch.Tensor, swa_lengths: torch.Tensor, swa_write_slots: torch.Tensor, @@ -173,14 +173,14 @@ def build_swa_indices( token_to_batch_idx, b_seq_len, WRITE_SWA_RUNTIME=swa_runtime_write_slots is not None, - SWA_PAGE_TOKEN_NUM=req_manager.cache_config.swa_page_token_num, + SWA_PAGE_TOKEN_NUM=swa_state.config.page_token_size, SLIDING_WINDOW=sliding_window, WIDTH=width, HAS_IMAGE=image_left is not None, BLOCK_W=triton.next_power_of_2(width), - SWA_CAPACITY_TOKENS=req_manager.swa_capacity_tokens, - HOLD_REQUEST_ID=req_manager.HOLD_REQUEST_ID, - HOLD_SWA_SLOT=prefill_state.HOLD_SWA_SLOT if prefill_state is not None else req_manager.HOLD_SWA_SLOT, + SWA_CAPACITY_TOKENS=swa_state.swa_capacity_tokens, + HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, + HOLD_SWA_SLOT=prefill_state.HOLD_SWA_SLOT if prefill_state is not None else swa_state.HOLD_SWA_SLOT, IS_PREFILL=prefill_state is not None, num_warps=4, ) diff --git a/lightllm/models/deepseek_v4/triton_kernel/cache_io.py b/lightllm/models/deepseek_v4/triton_kernel/cache_io.py deleted file mode 100644 index c012ca1b92..0000000000 --- a/lightllm/models/deepseek_v4/triton_kernel/cache_io.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Model-side packed views and token-indexed cache transfers.""" - -import triton -import triton.language as tl - -from lightllm.utils.envs_utils import get_env_start_args - - -def view_swa_kv(swa_buffer, config): - return swa_buffer.view(-1, config.get_packed_kv_bytes(config.swa_page_token_num)) - - -def get_compressed_kv_buffer(mem_manager, layer_index): - kv_buffer = mem_manager.get_att_input_params(layer_index) - if kv_buffer is None: - return None - config = mem_manager.cache_config - compress_ratio = config.compress_rates[layer_index] - packed_kv_pages = kv_buffer.view(-1, config.compressed_kv_token_alignment * kv_buffer.shape[-1]) - return packed_kv_pages[:, : config.get_packed_kv_bytes(config.compressed_kv_token_alignment // compress_ratio)] - - -def get_indexer_k_buffer(mem_manager, layer_index): - config = mem_manager.cache_config - kv_buffer = mem_manager.get_att_input_params(layer_index) - packed_kv_pages = kv_buffer.view(-1, config.compressed_kv_token_alignment * kv_buffer.shape[-1]) - entries = config.compressed_kv_token_alignment // config.short_compress_ratio - begin = config.get_packed_kv_bytes(entries) - return packed_kv_pages[:, begin : begin + entries * config.indexer_bytes] - - -@triton.jit -def _copy_compressed_kv_kernel( - kv_buffer, - kv_move_buffer, - mem_indexes, - page_indexes, - layer_byte_offset, - TOKEN_NUM: tl.constexpr, - COPY_NUM: tl.constexpr, - PAGE_TOKENS: tl.constexpr, - PAGE_STRIDE: tl.constexpr, - SLOT_BYTES: tl.constexpr, - TO_PAGES: tl.constexpr, - ROUND_TAIL: tl.constexpr, - BLOCK_BYTES: tl.constexpr, -): - kv_buffer = kv_buffer.to(tl.pointer_type(tl.uint8)) - tokens = tl.program_id(0) * 32 + tl.arange(0, 32) - slot = tl.load(mem_indexes + tokens, tokens < TOKEN_NUM, -1).to(tl.int64) - if ROUND_TAIL: - # Packed page tails include scales for earlier tokens. PD also carries - # the unused token slots of its final allocated compression block. - last = tl.load(mem_indexes + TOKEN_NUM - 1).to(tl.int64) - slot = tl.where(tokens < TOKEN_NUM, slot, last + tokens - TOKEN_NUM + 1) - page = tokens // PAGE_TOKENS - valid = (tokens < COPY_NUM) & (slot >= 0) - page_id = tl.load(page_indexes + page, tokens < COPY_NUM, 0).to(tl.int64) - columns = tl.arange(0, BLOCK_BYTES) - kv_ptr = kv_buffer + slot[:, None] * SLOT_BYTES + columns[None, :] - page_ptr = ( - kv_move_buffer - + page_id[:, None] * PAGE_STRIDE - + layer_byte_offset - + (tokens % PAGE_TOKENS)[:, None] * SLOT_BYTES - + columns[None, :] - ) - mask = valid[:, None] & (columns[None, :] < SLOT_BYTES) - if TO_PAGES: - tl.store(page_ptr, tl.load(kv_ptr, mask, 0), mask) - else: - tl.store(kv_ptr, tl.load(page_ptr, mask, 0), mask) - - -def copy_compressed_kv_pages( - mem_manager, mem_indexes, page_indexes, kv_move_buffer, *, to_pages, page_size, round_tail=False -): - if not mem_indexes.numel(): - return - config = mem_manager.cache_config - copy_num = ( - triton.cdiv(mem_indexes.numel(), config.compressed_kv_token_alignment) * config.compressed_kv_token_alignment - if round_tail - else mem_indexes.numel() - ) - layer_byte_offset = 0 - for kv_buffer in mem_manager.kv_buffer: - slot_bytes = kv_buffer.shape[-1] - _copy_compressed_kv_kernel[(triton.cdiv(copy_num, 32),)]( - kv_buffer.data_ptr(), - kv_move_buffer, - mem_indexes, - page_indexes, - layer_byte_offset, - TOKEN_NUM=mem_indexes.numel(), - COPY_NUM=copy_num, - PAGE_TOKENS=page_size, - PAGE_STRIDE=kv_move_buffer.stride(0), - SLOT_BYTES=slot_bytes, - TO_PAGES=to_pages, - ROUND_TAIL=round_tail, - BLOCK_BYTES=triton.next_power_of_2(slot_bytes), - num_warps=4, - ) - layer_byte_offset += page_size * slot_bytes - - -@triton.jit -def _copy_cpu_cache_kernel( - kv_buffer, - big_page_state_buffer, - mem_indexes, - big_page_buffer_ids, - page_indexes, - page_readies, - cpu_cache_tensor, - TOKEN_NUM: tl.constexpr, - PAGE_TOKENS: tl.constexpr, - PAGE_BYTES: tl.constexpr, - STATE_BYTES: tl.constexpr, - SLOT_BYTES: tl.constexpr, - TO_CPU_CACHE: tl.constexpr, - HAS_READIES: tl.constexpr, - BLOCK: tl.constexpr, -): - page = tl.program_id(0) - if HAS_READIES: - if tl.load(page_readies + page): - return - - cpu_page = cpu_cache_tensor + tl.load(page_indexes + page).to(tl.int64) * PAGE_BYTES - block_start, block_step = tl.program_id(1), tl.num_programs(1) - layer_offset = 0 - # CPU 页前半段按层保存压缩 KV 和 indexer K，碎页只复制实际存在的 token 槽位。 - for layer_index in tl.static_range(len(SLOT_BYTES)): - slot_bytes = SLOT_BYTES[layer_index] - layer_bytes = PAGE_TOKENS * slot_bytes - for block in range(block_start, tl.cdiv(layer_bytes, BLOCK), block_step): - offsets = block * BLOCK + tl.arange(0, BLOCK) - tokens = page * PAGE_TOKENS + offsets // slot_bytes - mask = (offsets < layer_bytes) & (tokens < TOKEN_NUM) - slots = tl.load(mem_indexes + tokens, mask, -1).to(tl.int64) - mask &= slots >= 0 - gpu = kv_buffer[layer_index] + slots * slot_bytes + offsets % slot_bytes - cpu = cpu_page + layer_offset + offsets - if TO_CPU_CACHE: - tl.store(cpu, tl.load(gpu, mask, 0), mask) - else: - tl.store(gpu, tl.load(cpu, mask, 0), mask) - layer_offset += layer_bytes - - # 页尾的 SWA、C4 和 indexer 状态已连续存储，作为一整块复制。 - state_id = tl.load(big_page_buffer_ids + page).to(tl.int64) - state = big_page_state_buffer + state_id * STATE_BYTES - for block in range(block_start, tl.cdiv(STATE_BYTES, BLOCK), block_step): - offsets = block * BLOCK + tl.arange(0, BLOCK) - mask = offsets < STATE_BYTES - cpu = cpu_page + layer_offset + offsets - if TO_CPU_CACHE: - tl.store(cpu, tl.load(state + offsets, mask, 0), mask) - else: - tl.store(state + offsets, tl.load(cpu, mask, 0), mask) - - -def _copy_cpu_cache_pages( - mem_manager, mem_indexes, big_page_buffer_ids, page_indexes, cpu_cache_tensor, *, to_cpu_cache, page_readies=None -): - assert len(big_page_buffer_ids) == len(page_indexes) - cpu_cache_pages = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1) - _copy_cpu_cache_kernel[(len(page_indexes), 12)]( - tuple(mem_manager.kv_buffer), - mem_manager.big_page_buffers.buffer, - mem_indexes, - big_page_buffer_ids, - page_indexes, - page_readies, - cpu_cache_pages, - TOKEN_NUM=len(mem_indexes), - PAGE_TOKENS=get_env_start_args().cpu_cache_token_page_size, - PAGE_BYTES=cpu_cache_pages.stride(0), - STATE_BYTES=mem_manager.big_page_buffers.buffer.shape[-1], - SLOT_BYTES=mem_manager.cache_config.compressed_kv_slot_bytes, - TO_CPU_CACHE=to_cpu_cache, - HAS_READIES=page_readies is not None, - BLOCK=4096, - num_warps=4, - ) - - -def copy_kv_buffer_to_cpu_cache( - mem_manager, mem_indexes, page_indexes, page_readies, big_page_buffer_ids, cpu_cache_tensor -): - _copy_cpu_cache_pages( - mem_manager, - mem_indexes, - big_page_buffer_ids, - page_indexes, - cpu_cache_tensor, - to_cpu_cache=True, - page_readies=page_readies, - ) - - -def copy_cpu_cache_to_kv_buffer(mem_manager, mem_indexes, big_page_buffer_ids, page_indexes, cpu_cache_tensor): - _copy_cpu_cache_pages( - mem_manager, mem_indexes, big_page_buffer_ids, page_indexes, cpu_cache_tensor, to_cpu_cache=False - ) - - -@triton.jit -def _copy_state_cache_kernel(src_state, dst_state, NBYTES: tl.constexpr, BLOCK: tl.constexpr): - offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - mask = offsets < NBYTES - tl.store(dst_state + offsets, tl.load(src_state + offsets, mask, 0), mask) - - -def copy_state_cache(src_state, dst_state): - """在当前 CUDA stream 上整块复制 pinned 状态，保持与 offload/load 的执行顺序。""" - assert src_state.shape == dst_state.shape - _copy_state_cache_kernel[(triton.cdiv(src_state.numel(), 4096),)]( - src_state, dst_state, NBYTES=src_state.numel(), BLOCK=4096, num_warps=4 - ) diff --git a/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py b/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py index e94e92c381..63d936afc0 100644 --- a/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py +++ b/lightllm/models/deepseek_v4/triton_kernel/swa_cache.py @@ -1,4 +1,4 @@ -"""SWA 物理索引、prefill 临时区，以及大小页 SWA 状态的保存和恢复。""" +"""SWA 物理索引与 prefill workspace 的布局准备和前缀搬运。""" import torch import triton @@ -27,21 +27,19 @@ def get_swa_token_index( return tl.where(req_idx == HOLD_REQUEST_ID, HOLD_SWA_SLOT + position % SWA_PAGE_TOKEN_NUM, slot) -def get_swa_token_indices(req_manager, req_idx, position_ids, *, prefill_state=None): - config = req_manager.cache_config - swa_token_indices = req_idx * req_manager.swa_capacity_tokens + position_ids % req_manager.swa_capacity_tokens - hold_slot = req_manager.HOLD_SWA_SLOT +def get_swa_token_indices(swa_state, req_idx, position_ids, *, prefill_state=None): + config = swa_state.config + swa_token_indices = req_idx * swa_state.swa_capacity_tokens + position_ids % swa_state.swa_capacity_tokens + hold_slot = swa_state.HOLD_SWA_SLOT if prefill_state is not None: ready = prefill_state.prefill_token_ranges[req_idx, 0] - first = ( - (ready - config.sliding_window + 1).clamp_min(0) // config.swa_page_token_num * config.swa_page_token_num - ) + first = (ready - config.window_size + 1).clamp_min(0) // config.page_token_size * config.page_token_size workspace = prefill_state.swa_prefill_workspace_offsets[req_idx] swa_token_indices = workspace + position_ids - first hold_slot = prefill_state.HOLD_SWA_SLOT return torch.where( - torch.as_tensor(req_idx, device=position_ids.device) == req_manager.HOLD_REQUEST_ID, - hold_slot + position_ids % config.swa_page_token_num, + torch.as_tensor(req_idx, device=position_ids.device) == swa_state.HOLD_REQUEST_ID, + hold_slot + position_ids % config.page_token_size, swa_token_indices, ) @@ -53,22 +51,24 @@ def _prepare_swa_prefill_workspace_kernel( b_req_idx, b_ready_cache_len, b_seq_len, - BATCH: tl.constexpr, + batch_size, SLIDING_WINDOW: tl.constexpr, SWA_PAGE_TOKEN_NUM: tl.constexpr, HOLD_REQUEST_ID: tl.constexpr, - SWA_WORKSPACE_BASE: tl.constexpr, + swa_workspace_base, BLOCK: tl.constexpr, ): batch = tl.arange(0, BLOCK) - valid = batch < BATCH + valid = batch < batch_size req_idx = tl.load(b_req_idx + batch, valid, HOLD_REQUEST_ID) ready = tl.load(b_ready_cache_len + batch, valid, 0) end = tl.load(b_seq_len + batch, valid, 0) first = tl.maximum(0, ready - SLIDING_WINDOW + 1) // SWA_PAGE_TOKEN_NUM * SWA_PAGE_TOKEN_NUM length = tl.cdiv(end - first, SWA_PAGE_TOKEN_NUM) * SWA_PAGE_TOKEN_NUM tokens = tl.where((req_idx != HOLD_REQUEST_ID) & valid, length, 0) - workspace = SWA_WORKSPACE_BASE + tl.cumsum(tokens) - tokens + # 请求数和 workspace 起点随每个 prefill microbatch 变化，保持为运行时参数， + # 避免按 batch 大小或 microbatch 索引重复编译同一份布局 kernel。 + workspace = swa_workspace_base + tl.cumsum(tokens) - tokens tl.store(prefill_token_ranges + req_idx * 2, ready, valid) tl.store(prefill_token_ranges + req_idx * 2 + 1, end, valid) tl.store(swa_prefill_workspace_offsets + req_idx, workspace, valid) @@ -76,8 +76,8 @@ def _prepare_swa_prefill_workspace_kernel( @triton.jit def _copy_swa_prefix_to_workspace_kernel( - req_to_swa, - swa_buffer, + runtime_swa_buffer, + workspace_swa_buffer, swa_prefill_workspace_offsets, b_req_idx, b_ready_cache_len, @@ -85,7 +85,7 @@ def _copy_swa_prefix_to_workspace_kernel( SLIDING_WINDOW: tl.constexpr, SWA_PAGE_TOKEN_NUM: tl.constexpr, HOLD_REQUEST_ID: tl.constexpr, - REQ_LAYER_STRIDE: tl.constexpr, + RUNTIME_LAYER_STRIDE: tl.constexpr, WORKSPACE_LAYER_STRIDE: tl.constexpr, PACKED_BYTES: tl.constexpr, BLOCK: tl.constexpr, @@ -101,125 +101,56 @@ def _copy_swa_prefix_to_workspace_kernel( if position >= ready: return source_slot = req_idx * SWA_CAPACITY_TOKENS + position % SWA_CAPACITY_TOKENS - source = req_to_swa + layer * REQ_LAYER_STRIDE + source_slot // SWA_PAGE_TOKEN_NUM * PACKED_BYTES - target = swa_buffer + layer * WORKSPACE_LAYER_STRIDE + (target_slot // SWA_PAGE_TOKEN_NUM + block) * PACKED_BYTES + source = runtime_swa_buffer + layer * RUNTIME_LAYER_STRIDE + source_slot // SWA_PAGE_TOKEN_NUM * PACKED_BYTES + target = ( + workspace_swa_buffer + + layer * WORKSPACE_LAYER_STRIDE + + (target_slot // SWA_PAGE_TOKEN_NUM + block) * PACKED_BYTES + ) lanes = tl.arange(0, BLOCK) for chunk in range(0, PACKED_BYTES, BLOCK): offset = chunk + lanes tl.store(target + offset, tl.load(source + offset, offset < PACKED_BYTES, 0), offset < PACKED_BYTES) -def prepare_swa_prefill_workspace(req_manager, prefill_state, b_req_idx, b_ready_cache_len, b_seq_len, workspace_base): +def prepare_swa_prefill_workspace(swa_state, prefill_state, b_req_idx, b_ready_cache_len, b_seq_len, workspace_base): if not b_req_idx.numel(): return - config = req_manager.cache_config + config = swa_state.config _prepare_swa_prefill_workspace_kernel[(1,)]( prefill_state.prefill_token_ranges, prefill_state.swa_prefill_workspace_offsets, b_req_idx, b_ready_cache_len, b_seq_len, - BATCH=b_req_idx.numel(), - SLIDING_WINDOW=config.sliding_window, - SWA_PAGE_TOKEN_NUM=config.swa_page_token_num, - HOLD_REQUEST_ID=req_manager.HOLD_REQUEST_ID, - SWA_WORKSPACE_BASE=workspace_base, + batch_size=b_req_idx.numel(), + SLIDING_WINDOW=config.window_size, + SWA_PAGE_TOKEN_NUM=config.page_token_size, + HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, + swa_workspace_base=workspace_base, BLOCK=triton.next_power_of_2(b_req_idx.numel()), num_warps=4, ) - swa_buffer = prefill_state.swa_buffer + workspace_swa_buffer = prefill_state.swa_buffer _copy_swa_prefix_to_workspace_kernel[ ( b_req_idx.numel(), - triton.cdiv(config.sliding_window - 1, config.swa_page_token_num) + 1, - swa_buffer.shape[0], + triton.cdiv(config.window_size - 1, config.page_token_size) + 1, + workspace_swa_buffer.shape[0], ) ]( - req_manager.req_to_swa, - swa_buffer, - prefill_state.swa_prefill_workspace_offsets, - b_req_idx, - b_ready_cache_len, - SWA_CAPACITY_TOKENS=req_manager.swa_capacity_tokens, - SLIDING_WINDOW=config.sliding_window, - SWA_PAGE_TOKEN_NUM=config.swa_page_token_num, - HOLD_REQUEST_ID=req_manager.HOLD_REQUEST_ID, - REQ_LAYER_STRIDE=req_manager.req_to_swa.stride(0), - WORKSPACE_LAYER_STRIDE=swa_buffer.stride(0), - PACKED_BYTES=config.get_packed_kv_bytes(config.swa_page_token_num), + runtime_swa_buffer=swa_state.swa_buffer, + workspace_swa_buffer=workspace_swa_buffer, + swa_prefill_workspace_offsets=prefill_state.swa_prefill_workspace_offsets, + b_req_idx=b_req_idx, + b_ready_cache_len=b_ready_cache_len, + SWA_CAPACITY_TOKENS=swa_state.swa_capacity_tokens, + SLIDING_WINDOW=config.window_size, + SWA_PAGE_TOKEN_NUM=config.page_token_size, + HOLD_REQUEST_ID=swa_state.HOLD_REQUEST_ID, + RUNTIME_LAYER_STRIDE=swa_state.swa_buffer.stride(0), + WORKSPACE_LAYER_STRIDE=workspace_swa_buffer.stride(0), + PACKED_BYTES=config.get_packed_kv_bytes(config.page_token_size), BLOCK=4096, num_warps=4, ) - - -@triton.jit -def _copy_swa_state_kernel( - swa_buffer, - swa_state, - swa_token_indices, - position_ids, - TOKEN_COUNT: tl.constexpr, - SWA_LAYER_STRIDE: tl.constexpr, - KV_BYTES: tl.constexpr, - SWA_PAGE_TOKEN_NUM: tl.constexpr, - SWA_PAGE_BYTES: tl.constexpr, - KV_DATA_BYTES: tl.constexpr, - GATHER: tl.constexpr, - BLOCK: tl.constexpr, -): - token, layer = tl.program_id(0), tl.program_id(1) - slot = tl.load(swa_token_indices + token).to(tl.int64) - position = tl.load(position_ids + token) - byte = tl.arange(0, BLOCK) - offset = slot // SWA_PAGE_TOKEN_NUM * SWA_PAGE_BYTES - packed_offset = tl.where( - byte < KV_DATA_BYTES, - slot % SWA_PAGE_TOKEN_NUM * KV_DATA_BYTES + byte, - SWA_PAGE_TOKEN_NUM * KV_DATA_BYTES - + slot % SWA_PAGE_TOKEN_NUM * (KV_BYTES - KV_DATA_BYTES) - + byte - - KV_DATA_BYTES, - ) - packed = swa_buffer + layer * SWA_LAYER_STRIDE + offset + packed_offset - dense = swa_state + (layer * TOKEN_COUNT + token) * KV_BYTES + byte - valid = (byte < KV_BYTES) & (position >= 0) - if GATHER: - tl.store(dense, tl.load(packed, valid, 0), byte < KV_BYTES) - else: - tl.store(packed, tl.load(dense, valid, 0), valid) - - -def _copy_swa_state(req_manager, req_idx, position_ids, swa_state, *, gather): - config = req_manager.cache_config - with torch.cuda.device(req_manager.req_to_swa.device): - swa_token_indices = get_swa_token_indices(req_manager, req_idx, position_ids) - _copy_swa_state_kernel[(position_ids.numel(), config.layer_num)]( - req_manager.req_to_swa, - swa_state, - swa_token_indices, - position_ids, - TOKEN_COUNT=position_ids.numel(), - SWA_LAYER_STRIDE=req_manager.req_to_swa.stride(0), - KV_BYTES=config.kv_bytes, - SWA_PAGE_TOKEN_NUM=config.swa_page_token_num, - SWA_PAGE_BYTES=config.get_packed_kv_bytes(config.swa_page_token_num), - KV_DATA_BYTES=config.kv_data_bytes, - GATHER=gather, - BLOCK=triton.next_power_of_2(config.kv_bytes), - num_warps=4, - ) - - -def gather_swa_state(req_manager, req_idx, position_ids): - """将 SWA packed 存储中的窗口读为 [层, token, kv_bytes]，供大小页保存。""" - swa_state = torch.empty( - (req_manager.cache_config.layer_num, position_ids.numel(), req_manager.cache_config.kv_bytes), - dtype=torch.uint8, - device=req_manager.req_to_swa.device, - ) - _copy_swa_state(req_manager, req_idx, position_ids, swa_state, gather=True) - return swa_state - - -def scatter_swa_state(req_manager, req_idx, position_ids, swa_state): - _copy_swa_state(req_manager, req_idx, position_ids, swa_state, gather=False) diff --git a/lightllm/models/deepseek_v4_dspark/infer_struct.py b/lightllm/models/deepseek_v4_dspark/infer_struct.py index 42d7eaf37e..5c4fe88d26 100644 --- a/lightllm/models/deepseek_v4_dspark/infer_struct.py +++ b/lightllm/models/deepseek_v4_dspark/infer_struct.py @@ -19,7 +19,7 @@ def _init_swa_metadata(self, model): return super()._init_swa_metadata(model) pos = self.position_ids - swa_width = ((self.req_manager.cache_config.sliding_window + model.block_size + 127) // 128) * 128 + swa_width = ((self.swa_state.config.window_size + model.block_size + 127) // 128) * 128 self.dsv4_swa_indices = torch.empty((pos.numel(), swa_width), dtype=torch.int32, device=pos.device) self.dsv4_swa_lengths = torch.empty_like(pos, dtype=torch.int32) self.dsv4_swa_write_slots = torch.empty_like(pos, dtype=torch.int32) @@ -27,12 +27,9 @@ def _init_swa_metadata(self, model): build_dspark_swa_indices( token_req_idx=self.dsv4_sparse_req_idx, position_ids=self.position_ids, - req_manager=self.req_manager, + swa_state=self.swa_state, swa_indices=self.dsv4_swa_indices, swa_lengths=self.dsv4_swa_lengths, swa_write_slots=self.dsv4_swa_write_slots, - sliding_window=model.config["sliding_window"], block_size=model.block_size, - hold_req_id=self.req_manager.HOLD_REQUEST_ID, - hold_swa_slot=self.req_manager.HOLD_SWA_SLOT, ) diff --git a/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py index 5b696e5f30..24bc22aab5 100644 --- a/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4_dspark/layer_infer/transformer_layer_infer.py @@ -14,9 +14,8 @@ class DeepseekV4DSparkTransformerLayerInfer(DeepseekV4TransformerLayerInfer): def __init__(self, layer_num, network_config): super().__init__(layer_num, network_config) - final_layer = network_config["n_layer"] + network_config["dspark_layer_num"] - 1 - self.stage_id = layer_num - network_config["n_layer"] - self.is_last_layer = layer_num == final_layer + self.stage_id = layer_num - network_config["draft_layer_start_id"] + self.is_last_layer = self.stage_id == network_config["n_layer"] - 1 def context_forward( self, diff --git a/lightllm/models/deepseek_v4_dspark/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/deepseek_v4_dspark/layer_weights/pre_and_post_layer_weight.py index c0a7efc8fa..64c79b4e40 100644 --- a/lightllm/models/deepseek_v4_dspark/layer_weights/pre_and_post_layer_weight.py +++ b/lightllm/models/deepseek_v4_dspark/layer_weights/pre_and_post_layer_weight.py @@ -24,8 +24,8 @@ def __init__(self, data_type, network_config, quant_cfg: Quantcfg): hc_mult = network_config["hc_mult"] markov_rank = network_config["dspark_markov_rank"] target_hidden_size = hidden * len(network_config["dspark_target_layer_ids"]) - first_layer_idx = network_config["n_layer"] - last_stage = network_config["dspark_layer_num"] - 1 + first_layer_idx = network_config["draft_layer_start_id"] + last_stage = network_config["n_layer"] - 1 last_prefix = f"mtp.{last_stage}" # Replaced with the already-loaded target weights by the model. diff --git a/lightllm/models/deepseek_v4_dspark/layer_weights/transformer_layer_weight.py b/lightllm/models/deepseek_v4_dspark/layer_weights/transformer_layer_weight.py index 1d37b95601..02ba7d055a 100644 --- a/lightllm/models/deepseek_v4_dspark/layer_weights/transformer_layer_weight.py +++ b/lightllm/models/deepseek_v4_dspark/layer_weights/transformer_layer_weight.py @@ -8,5 +8,5 @@ class DeepseekV4DSparkTransformerLayerWeight(DeepseekV4TransformerLayerWeight): def _parse_config(self): super()._parse_config() - stage_id = self.layer_num_ - self.network_config_["n_layer"] + stage_id = self.layer_num_ - self.network_config_["draft_layer_start_id"] self.prefix = f"mtp.{stage_id}" diff --git a/lightllm/models/deepseek_v4_dspark/model.py b/lightllm/models/deepseek_v4_dspark/model.py index ee7f40b42b..b2f324e75d 100644 --- a/lightllm/models/deepseek_v4_dspark/model.py +++ b/lightllm/models/deepseek_v4_dspark/model.py @@ -53,21 +53,25 @@ def __init__(self, kvargs: dict): assert not self.mtp_previous_draft_models, "DeepSeek-V4 DSpark uses one parallel draft model" super().__init__(kvargs) - def _init_config(self): - super()._init_config() - draft_layer_num = len(self.config["compress_ratios"]) - self.config["n_layer"] + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + config = super().normalize_model_config(model_config) + # draft 接在 main 的逻辑层之后；实际 checkpoint 权重使用 mtp.{stage_id} 前缀。 + config.setdefault("draft_layer_start_id", config["num_hidden_layers"]) + draft_layer_num = len(config["compress_ratios"]) - config["draft_layer_start_id"] assert draft_layer_num > 0 - assert self.config["compress_ratios"][-draft_layer_num:] == [0] * draft_layer_num - assert len(self.config["dspark_target_layer_ids"]) == draft_layer_num - self.config["dspark_layer_num"] = draft_layer_num + assert config["compress_ratios"][-draft_layer_num:] == [0] * draft_layer_num + assert len(config["dspark_target_layer_ids"]) == draft_layer_num + config.update(num_hidden_layers=draft_layer_num, n_layer=draft_layer_num) # Aliases consumed by the shared block/Markov implementation. - self.config["block_size"] = int(self.config["dspark_block_size"]) - self.config["mask_token_id"] = int(self.config["dspark_noise_token_id"]) - self.config["markov_rank"] = int(self.config["dspark_markov_rank"]) - self.config["markov_head_type"] = "vanilla" - self.config["enable_confidence_head"] = True - self.config["confidence_head_with_markov"] = True + config["block_size"] = int(config["dspark_block_size"]) + config["mask_token_id"] = int(config["dspark_noise_token_id"]) + config["markov_rank"] = int(config["dspark_markov_rank"]) + config["markov_head_type"] = "vanilla" + config["enable_confidence_head"] = True + config["confidence_head_with_markov"] = True + return config def _verify_params(self): super()._verify_params() @@ -82,8 +86,8 @@ def _verify_params(self): def _init_quant(self): super()._init_quant() - expert_quant_type = self.quant_cfg.get_quant_type(self.config["n_layer"] - 1, "fused_moe") - for layer_idx in range(self.config["n_layer"], len(self.config["compress_ratios"])): + expert_quant_type = self.quant_cfg.get_quant_type(self.config["draft_layer_start_id"] - 1, "fused_moe") + for layer_idx in range(self.config["draft_layer_start_id"], len(self.config["compress_ratios"])): self.quant_cfg.quant_cfg[layer_idx]["fused_moe"] = expert_quant_type def _init_weights(self, start_layer_index=None): @@ -94,7 +98,7 @@ def _init_weights(self, start_layer_index=None): ) self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ - first_layer = self.config["n_layer"] + first_layer = self.config["draft_layer_start_id"] self.trans_layers_weight = [ self.transformer_weight_class( first_layer + stage_id, @@ -102,7 +106,7 @@ def _init_weights(self, start_layer_index=None): network_config=self.config, quant_cfg=self.quant_cfg, ) - for stage_id in range(self.config["dspark_layer_num"]) + for stage_id in range(self.config["n_layer"]) ] self.context_wkv_weight = ROWMMWeight( in_dim=self.config["hidden_size"], @@ -114,12 +118,6 @@ def _init_weights(self, start_layer_index=None): tp_world_size=1, ) - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_size: int): # A DSpark logical request always occupies one complete physical block. # Keep the generic HOLD padding semantics, but make its requested size @@ -134,16 +132,12 @@ def _create_padded_decode_model_input(self, model_input: ModelInput, new_batch_s def _init_infer_layer(self, start_layer_index=None): self.pre_infer = self.pre_layer_infer_class(network_config=self.config) self.post_infer = self.post_layer_infer_class(network_config=self.config) - first_layer = self.config["n_layer"] + first_layer = self.config["draft_layer_start_id"] self.layers_infer = [ self.transformer_layer_infer_class(first_layer + stage_id, network_config=self.config) - for stage_id in range(self.config["dspark_layer_num"]) + for stage_id in range(self.config["n_layer"]) ] - def _init_some_value(self): - super()._init_some_value() - self.layers_num = self.config["dspark_layer_num"] - def _init_custom(self): self._freqs_cis_sliding = self.main_model._freqs_cis_sliding self._freqs_cis_compress = self.main_model._freqs_cis_compress diff --git a/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py index bd1a5a3929..a8f124b6f7 100644 --- a/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py +++ b/lightllm/models/deepseek_v4_mtp/layer_weights/pre_and_post_layer_weight.py @@ -20,7 +20,7 @@ def __init__(self, data_type, network_config, quant_cfg: Quantcfg): hidden = network_config["hidden_size"] vocab = network_config["vocab_size"] hc_mult = network_config["hc_mult"] - layer_idx = network_config["n_layer"] + layer_idx = network_config["draft_weight_layer_id"] prefix = "mtp.0" self.wte_weight_ = EmbeddingWeight( diff --git a/lightllm/models/deepseek_v4_mtp/model.py b/lightllm/models/deepseek_v4_mtp/model.py index 95c3f183fb..36ac375cef 100644 --- a/lightllm/models/deepseek_v4_mtp/model.py +++ b/lightllm/models/deepseek_v4_mtp/model.py @@ -39,6 +39,14 @@ def __init__(self, kvargs: dict): super().__init__(kvargs) return + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + config = super().normalize_model_config(model_config) + # 权重层号用于量化配置及权重构造，与各 draft 实例的执行逻辑层号分别记录。 + config.setdefault("draft_weight_layer_id", config["num_hidden_layers"]) + config.update(num_hidden_layers=1, n_layer=1) + return config + def _init_custom(self): self._freqs_cis_sliding = self.main_model._freqs_cis_sliding self._freqs_cis_compress = self.main_model._freqs_cis_compress @@ -53,16 +61,8 @@ def _init_custom(self): layer.sin_compress_table = self._sin_cached_compress return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_weights(self, start_layer_index=None): - mtp_layer_index = self.config["n_layer"] + draft_weight_layer_id = self.config["draft_weight_layer_id"] self.pre_post_weight = self.pre_and_post_weight_class( self.data_type, network_config=self.config, quant_cfg=self.quant_cfg ) @@ -70,7 +70,7 @@ def _init_weights(self, start_layer_index=None): self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ self.trans_layers_weight = [ self.transformer_weight_class( - mtp_layer_index, + draft_weight_layer_id, self.data_type, network_config=self.config, quant_cfg=self.quant_cfg, @@ -89,11 +89,6 @@ def _init_infer_layer(self, start_layer_index=None): assert self.layers_infer[0].compress_ratio == 0, "DeepSeek-V4 MTP draft layer must be SWA-only" return - def _init_some_value(self): - super()._init_some_value() - self.layers_num = 1 - return - def _gen_special_model_input(self, token_num: int): return { "mtp_draft_input_hiddens": torch.randn( diff --git a/lightllm/models/gemma3/model.py b/lightllm/models/gemma3/model.py index eabbfe95e0..1267979882 100644 --- a/lightllm/models/gemma3/model.py +++ b/lightllm/models/gemma3/model.py @@ -1,9 +1,8 @@ +from typing import TYPE_CHECKING import os import json import torch from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.models.gemma3.infer_struct import Gemma3InferStateInfo from lightllm.models.gemma3.layer_infer.post_layer_infer import Gemma3PostLayerInfer from lightllm.models.gemma3.layer_infer.pre_layer_infer import Gemma3PreLayerInfer @@ -17,6 +16,12 @@ from transformers import AutoConfig from lightllm.utils.log_utils import init_logger +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + logger = init_logger(__name__) @@ -138,16 +143,37 @@ def _init_custom(self): self.head_dim_ = 256 self._init_to_get_rotary() - def _init_mem_manager(self): - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=torch.bfloat16, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=256, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """Gemma3 使用固定 256 维 BF16 KV，不随权重精度改变。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=256, + dtype=torch.bfloat16, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) - return + return [state] def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: diff --git a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py index 2f0c01dbf6..d1ebd30e30 100644 --- a/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.full_att import FullAttentionState import math import torch import torch.nn as nn @@ -21,7 +22,7 @@ class Gemma4TransformerLayerInfer(LlamaTransformerLayerInfer): time. The KV cache layout is uniform (sliding shape: num_kv_heads=16, head_dim=256); full-attention layers pack their (4, 512) tensor into the first 8 heads of the 16-head slot at cache-write time, then reshape on - read. See Gemma4TpPartModel._init_mem_manager for context. + read. See Gemma4TpPartModel.create_full_att_states for context. """ def __init__(self, layer_num, network_config): @@ -84,17 +85,6 @@ def __init__(self, layer_num, network_config): kv_shared_count = network_config.get("num_kv_shared_layers") or 0 total_layers = network_config["num_hidden_layers"] self.is_kv_shared_ = kv_shared_count > 0 and layer_num >= total_layers - kv_shared_count - self.kv_share_target_layer_ = None - if self.is_kv_shared_: - cutoff = total_layers - kv_shared_count - for j in range(cutoff - 1, -1, -1): - if network_config["layer_types"][j] == layer_type: - self.kv_share_target_layer_ = j - break - assert self.kv_share_target_layer_ is not None, ( - f"layer {layer_num} ({layer_type}) is KV-shared but no earlier non-shared " - f"layer of the same type found below cutoff={cutoff}" - ) # Always 1.0: NoPE dims for full-attn layers are zero-padded into # cos/sin (cos=1, sin=0 → identity), so the kernel walks the whole @@ -189,13 +179,18 @@ def _post_cache_kv(self, cache_kv, infer_state, layer_weight): def _att_control(self): if self.is_sliding and self.sliding_window_ > 0: w = self.sliding_window_ - 1 - return AttControl(use_sliding_window=True, sliding_window=(w, 0)) - return AttControl(use_sliding_window=False, sliding_window=(-1, -1)) + return AttControl(use_sliding_window=True, sliding_window=(w, 0), logical_layer_id=self.layer_num_) + return AttControl(use_sliding_window=False, sliding_window=(-1, -1), logical_layer_id=self.layer_num_) def _get_layer_kv(self, infer_state: InferStateInfo): - # KV-shared layers read from the target layer's cache slot. - layer_idx = self.kv_share_target_layer_ if self.is_kv_shared_ else self.layer_num_ - _k_raw, _v_raw = infer_state.mem_manager.get_att_input_params(layer_index=layer_idx) + # 直接用当前逻辑层号访问；共享层由 State 的 logical_layer_groups 映射到生产层的物理 KV。 + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + _k_raw, _v_raw = state.get_layer_view(logical_layer_id=self.layer_num_) # _k_raw / _v_raw shape (S, cache_slot_num, cache_slot_dim). Use .view # (not .reshape) so any non-contiguous layout from a future mem_manager # backend fails loudly instead of silently copying — slice + view is diff --git a/lightllm/models/gemma4/model.py b/lightllm/models/gemma4/model.py index b68a1a4399..4f998ea25e 100644 --- a/lightllm/models/gemma4/model.py +++ b/lightllm/models/gemma4/model.py @@ -1,8 +1,8 @@ +from typing import TYPE_CHECKING import os import json import torch from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class from lightllm.common.build_utils import repair_config from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.gemma4.infer_struct import Gemma4InferStateInfo @@ -11,13 +11,19 @@ from lightllm.models.gemma4.layer_infer.transformer_layer_infer import Gemma4TransformerLayerInfer from lightllm.models.gemma4.layer_weights.pre_and_post_layer_weight import Gemma4PreAndPostLayerWeight from lightllm.models.gemma4.layer_weights.transformer_layer_weight import Gemma4TransformerLayerWeight -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger from lightllm.distributed.communication_op import dist_group_manager logger = init_logger(__name__) +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class Gemma4TpPartModel(LlamaTpPartModel): pre_and_post_weight_class = Gemma4PreAndPostLayerWeight transformer_weight_class = Gemma4TransformerLayerWeight @@ -31,7 +37,7 @@ class Gemma4TpPartModel(LlamaTpPartModel): def __init__(self, kvargs): # head_dim_ is used by the default _init_to_get_rotary which we # override; still set it to the sliding-layer head_dim for consistency - # with the mem manager and any generic helpers. + # with the full-attention storage and any generic helpers. self.head_dim_ = 256 super().__init__(kvargs) return @@ -80,31 +86,57 @@ def _verify_params(self): ) return - def _init_mem_manager(self): - # Uniform per-layer KV cache layout. The per-layer cache slot must fit - # whichever layer type has the largest per-token K/V width: sliding - # (num_key_value_heads * head_dim) or full - # (num_global_kv * global_head_dim). Keep cache_slot_dim = head_dim - # and pick cache_slot_num = max-width / head_dim. For 31B this - # collapses to num_key_value_heads; for E4B the full-attn shape wins - # (2*512 > 2*256), so it uses 4 storage slots of 256 dims. - # Gemma4TransformerLayerInfer.__init__ computes the same value and - # uses it to pack/unpack K/V at write/read time. - head_dim = self.config["head_dim"] - num_global_kv = self.config.get("num_global_key_value_heads") or self.config["num_key_value_heads"] - sliding_total = self.config["num_key_value_heads"] * self.config["head_dim"] - full_total = num_global_kv * self.config["global_head_dim"] + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按 Gemma4 的 attention 宽度和共享关系创建物理 KV 层。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + # 每层的 KV 槽位须容纳 sliding/full 两类 attention 的最大宽度。 + head_dim = config["head_dim"] + num_global_kv = config.get("num_global_key_value_heads") or config["num_key_value_heads"] + sliding_total = config["num_key_value_heads"] * head_dim + full_total = num_global_kv * config["global_head_dim"] per_token_k_width = max(sliding_total, full_total) - head_num_per_rank = (per_token_k_width // head_dim) // self.tp_world_size_ - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=head_num_per_rank, - head_dim=head_dim, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, + assert per_token_k_width % head_dim == 0, "Gemma4 KV width must contain complete storage heads" + + shared_layer_count = config.get("num_kv_shared_layers") or 0 + storage_layer_count = len(model_logical_layer_ids) - shared_layer_count + assert 0 <= shared_layer_count < len(model_logical_layer_ids), "invalid KV-shared layer count" + layer_groups = [[layer_id] for layer_id in model_logical_layer_ids[:storage_layer_count]] + # 后部共享层读取同类型的最后一个生产层，只增加逻辑映射，不增加物理存储行。 + for local_layer in range(storage_layer_count, len(model_logical_layer_ids)): + for source_layer in range(storage_layer_count - 1, -1, -1): + if config["layer_types"][source_layer] == config["layer_types"][local_layer]: + layer_groups[source_layer].append(model_logical_layer_ids[local_layer]) + break + else: + raise AssertionError("KV-shared layer requires an earlier producer of the same attention type") + logical_layer_groups = tuple(tuple(group) for group in layer_groups) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=per_token_k_width // head_dim, + head_dim=head_dim, + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) - return + return [state] def _init_att_backend(self): # Gemma-4 has per-layer heterogeneous attention: sliding layers use @@ -178,7 +210,16 @@ def _init_to_get_rotary_gemma4(self): sliding_partial = sliding_params.get("partial_rotary_factor", 1.0) sliding_rot_dim = int(sliding_head_dim * sliding_partial) inv_freq_sliding = 1.0 / ( - sliding_theta ** (torch.arange(0, sliding_rot_dim, 2, dtype=torch.float32) / sliding_rot_dim) + sliding_theta + ** ( + torch.arange( + start=0, + end=sliding_rot_dim, + step=2, + dtype=torch.float32, + ) + / sliding_rot_dim + ) ) freqs_s = torch.outer(t, inv_freq_sliding) self._cos_cached_sliding = torch.cos(freqs_s).to(self.data_type).cuda() @@ -198,7 +239,16 @@ def _init_to_get_rotary_gemma4(self): if rope_type == "proportional": rope_angles = int(full_partial * full_head_dim // 2) inv_freq_rot = 1.0 / ( - full_theta ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float32) / full_head_dim) + full_theta + ** ( + torch.arange( + start=0, + end=2 * rope_angles, + step=2, + dtype=torch.float32, + ) + / full_head_dim + ) ) nope_angles = full_head_dim // 2 - rope_angles if nope_angles > 0: @@ -207,7 +257,18 @@ def _init_to_get_rotary_gemma4(self): inv_freq_full = inv_freq_rot else: full_rot_dim = int(full_head_dim * full_partial) - inv_freq_full = 1.0 / (full_theta ** (torch.arange(0, full_rot_dim, 2, dtype=torch.float32) / full_rot_dim)) + inv_freq_full = 1.0 / ( + full_theta + ** ( + torch.arange( + start=0, + end=full_rot_dim, + step=2, + dtype=torch.float32, + ) + / full_rot_dim + ) + ) freqs_f = torch.outer(t, inv_freq_full) self._cos_cached_full = torch.cos(freqs_f).to(self.data_type).cuda() diff --git a/lightllm/models/gemma_2b/model.py b/lightllm/models/gemma_2b/model.py index b31f79e502..e16d08cfce 100644 --- a/lightllm/models/gemma_2b/model.py +++ b/lightllm/models/gemma_2b/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING from lightllm.models.gemma_2b.layer_weights.transformer_layer_weight import Gemma_2bTransformerLayerWeight from lightllm.models.gemma_2b.layer_weights.pre_and_post_layer_weight import Gemma_2bPreAndPostLayerWeight from lightllm.models.gemma_2b.layer_infer.pre_layer_infer import Gemma_2bPreLayerInfer @@ -5,8 +6,12 @@ from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num + + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState class Gemma_2bTpPartModel(LlamaTpPartModel): @@ -36,13 +41,36 @@ def _verify_params(self): assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 return - def _init_mem_manager(self): - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"], - head_dim=self.config["hidden_size"] // self.config["num_attention_heads"], - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """Gemma 2B 保留复制 KV 的布局及本模型定义的 head 维度。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config["hidden_size"] // config["num_attention_heads"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + replicated=True, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) - return + return [state] diff --git a/lightllm/models/glm4_moe_lite_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm4_moe_lite_mtp/layer_weights/pre_and_post_layer_weight.py index 205b8d6026..5d9150fea2 100644 --- a/lightllm/models/glm4_moe_lite_mtp/layer_weights/pre_and_post_layer_weight.py +++ b/lightllm/models/glm4_moe_lite_mtp/layer_weights/pre_and_post_layer_weight.py @@ -14,7 +14,7 @@ def __init__(self, data_type, network_config, quant_cfg: Quantcfg, layer_idx=Non self.quant_cfg: Quantcfg = quant_cfg if layer_idx is None: - layer_idx = network_config["num_hidden_layers"] + layer_idx = network_config["mtp_checkpoint_layer_start"] hidden_size = network_config["hidden_size"] self.eh_proj_weight_ = ROWMMWeight( diff --git a/lightllm/models/glm4_moe_lite_mtp/model.py b/lightllm/models/glm4_moe_lite_mtp/model.py index 2e4ba5c86b..18ceb80e78 100644 --- a/lightllm/models/glm4_moe_lite_mtp/model.py +++ b/lightllm/models/glm4_moe_lite_mtp/model.py @@ -24,20 +24,23 @@ def _pre_init(self, kvargs: dict): self.main_model: TpPartBaseModel = kvargs.pop("main_model") self.mtp_previous_draft_models: List[TpPartBaseModel] = kvargs.pop("mtp_previous_draft_models") + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + config = super().normalize_model_config(model_config) + # checkpoint 中 MTP 权重接在 main 层之后；重复规范化不能覆盖这个起点。 + config.setdefault("mtp_checkpoint_layer_start", config["num_hidden_layers"]) + count = config.get("num_nextn_predict_layers", 1) + config.update(num_hidden_layers=count, n_layer=count) + return config + def _init_custom(self): self._cos_cached = self.main_model._cos_cached self._sin_cached = self.main_model._sin_cached - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - def _init_weights(self, start_layer_index=None): assert start_layer_index is None - mtp_layer_start = self.config["num_hidden_layers"] + mtp_layer_start = self.config["mtp_checkpoint_layer_start"] num_mtp_layers = self.config.get("num_nextn_predict_layers", 1) self.pre_post_weight = self.pre_and_post_weight_class( diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py new file mode 100644 index 0000000000..0e52e736a9 --- /dev/null +++ b/lightllm/models/glm5_next/config.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +import dataclasses + +from lightllm.utils.torch_dtype_utils import get_torch_dtype + +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + + +@dataclasses.dataclass +class Glm5NextResourceConfig(Qwen3NextResourceConfig): + """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" + + # Indexer 池化步长，单位为原始 token；同时决定 K-pool tail 的基准 ring 长度和恢复边界。 + index_kpool: int = 4 + # 单个 indexer key 的特征维度；原始 key/gate tail 的 channel 宽度为该值的两倍。 + index_head_dim: int = 128 + + @classmethod + def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict) -> "Glm5NextResourceConfig": + config = config.get("text_config", config) + linear = config["linear_attn_config"] + tp = args.tp // args.dp + assert linear["num_heads"] % tp == 0 + layers = config["num_hidden_layers"] + assert config["layer_types"] == [ + "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) + ] + dtype = get_torch_dtype(args.data_type) + return cls( + tp_world_size=tp, + full_att_all_num_kv_heads=1, + full_att_head_dim=config["kv_lora_rank"], + global_linear_k_heads=linear["num_heads"], + global_linear_v_heads=linear["num_heads"], + head_linear_k_dim=linear["head_dim"], + head_linear_v_dim=linear["head_dim"], + conv_kernel_size=linear["short_conv_kernel_size"], + conv_state_dtype=dtype, + ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), + full_attention_interval=4, + all_layer_num=layers, + index_kpool=config["index_kpool"], + index_head_dim=config["index_head_dim"], + ) diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index cf13ab83de..5df510e5de 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.full_att import FullAttentionState import torch import triton @@ -16,10 +17,11 @@ def __init__(self, layer_idx, network_config, tp_world_size): self.heads = network_config["index_n_heads"] self.dim = network_config["index_head_dim"] self.eps = network_config["rms_norm_eps"] + self.pool_size = network_config["index_kpool"] def select_topk_indices(self, logits, lengths, indices): """Select row-relative indices, with valid entries before -1 padding.""" - if HAS_VLLM: + if HAS_VLLM and hasattr(vllm_ops, "top_k_per_row_decode"): # next_n=1 treats each query as an independent row, including prefill. # The decode entry splits long rows before merging their candidates. # Tested with vLLM 0.22.1; persistent_topk can drop candidates (#51782). @@ -36,8 +38,19 @@ 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) - packed_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx) + role = infer_state.att_state_role + tail = infer_state.att_state_manager.get_state( + model_role=role, + state_type_filter=f"compressor_indexer_w{self.pool_size}_s{self.pool_size}", + logical_layer_id=self.layer_idx, + ).get_layer_view(logical_layer_id=self.layer_idx) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", + logical_layer_id=self.layer_idx, + state_class=FullAttentionState, + ) + packed_buffer = state.get_layer_view(logical_layer_id=self.layer_idx).indexer_k compress_pools( raw=raw, tail=tail, diff --git a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index 620554bc06..be7a2fa781 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -2,6 +2,8 @@ from __future__ import annotations +from lightllm.common.attention_state.full_att import FullAttentionState + import torch from lightllm.common.basemodel import TransformerLayerInferTpl @@ -159,9 +161,15 @@ def _context_attention_kernel(self, q, kv, infer_state, layer_weight, out=None): layer_weight=layer_weight, ) del infer_state.get_topk_indices_params + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) return infer_state.prefill_att_state.prefill_att( q=q, - k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, v=None, att_control=AttControl( nsa_prefill=True, @@ -240,9 +248,15 @@ def _token_attention_kernel(self, q, infer_state, layer_weight, out=None): ) del infer_state.get_topk_indices_params q_rope = q_nope[..., :0] + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="mla", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) return infer_state.decode_att_state.decode_att( q=(q_nope, q_rope), - k=infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_), + k=state.get_layer_view(logical_layer_id=self.layer_num_).kv, v=None, att_control=AttControl( nsa_decode=True, diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 6c5ffa6280..5b5235c76e 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -1,17 +1,14 @@ import json import os +from typing import TYPE_CHECKING import torch import triton -from lightllm.common.build_utils import repair_config 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.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager -from lightllm.common.req_manager import Glm5NextReqManager -from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.distributed.communication_op import dist_group_manager from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from .layer_infer.pre_layer_infer import Glm5NextPreLayerInfer @@ -20,34 +17,117 @@ from .layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + + class Glm5NextTpPartModel(TpPartBaseModel): + is_hybrid_att_model = True pre_and_post_weight_class = Glm5NextPreAndPostLayerWeight transformer_weight_class = Glm5NextTransformerLayerWeight pre_layer_infer_class = Glm5NextPreLayerInfer post_layer_infer_class = LlamaPostLayerInfer transformer_layer_infer_class = Glm5NextTransformerLayerInfer + @classmethod + def create_bounded_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", + ) -> list["BoundedAttentionState"]: + """分别创建 KDA 运行态和索引 key 的无重叠池化窗口，不分配大小页。""" + from lightllm.common.attention_state.bounded_att import ( + ConvSSMStateConfig, + WindowCompressorStateConfig, + create_kda_state, + create_window_compressor_state, + ) + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + assert not ( + args.mtp_step and args.enable_prefill_decode_mixed + ), "MTP linear attention does not support mixed prefill/decode" + cfg = Glm5NextResourceConfig.from_model_config(config, args) + states = [] + linear_layers = tuple( + layer for i, layer in enumerate(model_logical_layer_ids) if (i + 1) % cfg.full_attention_interval + ) + if linear_layers: + states.append( + create_kda_state( + config=ConvSSMStateConfig( + logical_layer_ids=linear_layers, + global_k_heads=cfg.global_linear_k_heads, + global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, + value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, + conv_dtype=cfg.conv_state_dtype, + ssm_dtype=cfg.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + ) + full_layers = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) + if full_layers: + assert ( + args.linear_att_hash_page_size % config["index_kpool"] == 0 + ), "GLM cache pages must align to index_kpool" + states.append( + create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=full_layers, + window_size=config["index_kpool"], + compress_stride=config["index_kpool"], + candidate_dim=config["index_head_dim"], + target="indexer", + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + ) + return states + + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """保留外层量化配置，并统一 GLM 文本配置的别名及 autotune 层数。""" + config = super().normalize_model_config(model_config) + if "quantization_config" in model_config: + config["quantization_config"] = dict(model_config["quantization_config"]) + config["autotune_layer_num"] = 4 + return config + def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json")) as f: - outer_config = json.load(f) - self.config = dict(outer_config.get("text_config", outer_config)) - if "quantization_config" in outer_config: - self.config["quantization_config"] = dict(outer_config["quantization_config"]) - self.config["autotune_layer_num"] = 4 - repair_config(self.config, same_names=["num_attention_heads", "n_head"]) - repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) + self.config = self.normalize_model_config(json.load(f)) def _verify_params(self): super()._verify_params() assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): - raise ValueError("GLM-5.3 Flash supports only --llm_kv_type None or fp8kv_dsa") - if self.args.llm_kv_type == "fp8kv_dsa" and ( - self.config["kv_lora_rank"] != 512 or self.data_type != torch.bfloat16 - ): - raise ValueError("GLM-5.3 Flash fp8kv_dsa requires kv_lora_rank=512 and --data_type bfloat16") + raise ValueError("GLM-5.3 Flash supports only unquantized KV or dynamic fp8kv_dsa") + if self.args.llm_kv_type == "fp8kv_dsa": + # group64/E8M0 编码与稀疏 attention 算子要求 BF16 输入和固定的 512 维 latent。 + if self.data_type != torch.bfloat16 or self.config["kv_lora_rank"] != 512: + raise ValueError("GLM-5.3 Flash FP8 KV requires BF16 computation and kv_lora_rank=512") def autotune_layers(self): return 4 @@ -65,28 +145,42 @@ 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): - 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, - ) + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype - def _init_mem_manager(self): - mem_manager_class = FP8Glm5NextMemManager if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextMemManager - self.mem_manager = mem_manager_class( - size=self.max_total_token_num, - dtype=self.linear_config.full_att_dtype, - num_kv_heads=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, - mem_fraction=self.mem_fraction, + config = cls.normalize_model_config(model_config) + cfg = Glm5NextResourceConfig.from_model_config(config, args) + logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_layer_indices()) + if not logical_layer_groups: + return [] + state = create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=logical_layer_groups, + latent_dim=cfg.full_att_head_dim, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + format="packed_group64_e8m0" + if args.llm_kv_type == "fp8kv_dsa" + else ("packed_fp16" if get_torch_dtype(args.data_type) == torch.float16 else "packed_bf16"), + ), + page_size=args.page_size, + device="meta", ) + return [state] def _init_att_backend(self): + from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + backend_class = ( Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend ) diff --git a/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py index 06e5576f16..40260a3be6 100644 --- a/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py +++ b/lightllm/models/glm5_next_mtp/layer_weights/pre_and_post_layer_weight.py @@ -7,7 +7,7 @@ class Glm5NextMTPPreAndPostLayerWeight(PreAndPostLayerWeight): def __init__(self, data_type, network_config, quant_cfg): super().__init__(data_type, network_config) hidden_size = network_config["hidden_size"] - prefix = f"model.layers.{network_config['num_hidden_layers']}" + prefix = f"model.layers.{network_config['mtp_checkpoint_layer_start']}" self.eh_proj_weight_ = ROWMMWeight( in_dim=2 * hidden_size, out_dims=[hidden_size], diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index c82962fda7..a59c6f1ab1 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -1,4 +1,6 @@ +from typing import TYPE_CHECKING from copy import deepcopy +import torch from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend @@ -10,6 +12,13 @@ from .layer_weights.pre_and_post_layer_weight import Glm5NextMTPPreAndPostLayerWeight +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + + class Glm5NextMTPModel(TpPartBaseModel): is_mtp_draft_model = True pre_and_post_weight_class = Glm5NextMTPPreAndPostLayerWeight @@ -25,9 +34,8 @@ def __init__(self, kvargs): def _init_config(self): # Reuse the validated target configuration without changing its mHC layout. - self.config = deepcopy(self.main_model.config) + self.config = self.normalize_model_config(deepcopy(self.main_model.config)) assert self.config.get("num_nextn_predict_layers") == 1, "GLM NextN requires one native MTP block" - self.config["mhc"] = False def _init_weights(self, start_layer_index=None): assert start_layer_index is None @@ -36,7 +44,12 @@ def _init_weights(self, start_layer_index=None): self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ self.pre_post_weight.main_norm_weight_ = self.main_model.pre_post_weight.final_norm_weight_ self.trans_layers_weight = [ - self.transformer_weight_class(self.config["num_hidden_layers"], self.data_type, self.config, self.quant_cfg) + self.transformer_weight_class( + layer_num=self.config["mtp_checkpoint_layer_start"], + data_type=self.data_type, + network_config=self.config, + quant_cfg=self.quant_cfg, + ) ] def _init_infer_layer(self, start_layer_index=None): @@ -59,12 +72,86 @@ 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): - self.req_manager = self.main_model.req_manager - self.linear_config = self.main_model.linear_config + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """原生 GLM draft 的自有层采用 NoPE packed MLA，与 main 分别持有存储。""" + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=logical_layer_groups, + latent_dim=config["kv_lora_rank"], + tp_rank=tp_rank, + tp_world_size=tp_world_size, + format="packed_group64_e8m0" + if args.llm_kv_type == "fp8kv_dsa" + else ("packed_fp16" if get_torch_dtype(args.data_type) == torch.float16 else "packed_bf16"), + ), + page_size=args.page_size, + device="meta", + ) + return [state] + + @classmethod + def create_bounded_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", + ) -> list["BoundedAttentionState"]: + """直接选择本模型的 K-pool 窗口层，不依赖 full State 创建过程。""" + from lightllm.common.attention_state.bounded_att import ( + create_window_compressor_state, + WindowCompressorStateConfig, + ) + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_ids = model_logical_layer_ids + if not logical_layer_ids: + return [] + assert args.linear_att_hash_page_size % config["index_kpool"] == 0, "GLM cache pages must align to index_kpool" + state = create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=logical_layer_ids, + window_size=config["index_kpool"], + compress_stride=config["index_kpool"], + candidate_dim=config["index_head_dim"], + target="indexer", + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + return [state] - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """区分 draft 自有层数与 checkpoint 中的 MTP 权重层起点。""" + config = super().normalize_model_config(model_config) + # checkpoint 中 MTP 权重接在 main 层之后；重复规范化不能覆盖这个起点。 + config.setdefault("mtp_checkpoint_layer_start", config["num_hidden_layers"]) + config.update(num_hidden_layers=1, n_layer=1, mhc=False) + return config def _init_att_backend(self): backend_class = ( diff --git a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index 490d2dc4c5..b26025c753 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.full_att import FullAttentionState import torch from lightllm.models.gpt_oss.layer_weights.transformer_layer_weight import GptOssTransformerLayerWeight from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer, LlamaInferStateInfo @@ -72,7 +73,13 @@ def _context_attention_kernel( window_size = (-1, -1) use_sliding_window = False - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, @@ -99,7 +106,13 @@ def _token_attention_kernel( window_size = (-1, -1) use_sliding_window = False - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.decode_att_state.decode_att( q=_q, diff --git a/lightllm/models/llama/layer_infer/transformer_layer_infer.py b/lightllm/models/llama/layer_infer/transformer_layer_infer.py index 69acffaa4d..0b3bb2e867 100644 --- a/lightllm/models/llama/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/llama/layer_infer/transformer_layer_infer.py @@ -1,3 +1,5 @@ +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.common.basemodel.attention import AttControl import torch import triton import torch.distributed as dist @@ -44,12 +46,19 @@ def _context_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) o_tensor = infer_state.prefill_att_state.prefill_att( q=_q, k=_k, v=_v, + att_control=AttControl(logical_layer_id=self.layer_num_), alloc_func=self.alloc_tensor, ) o_tensor = o_tensor.view(q.shape) @@ -61,9 +70,21 @@ def _token_attention_kernel( infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight, ) -> torch.Tensor: - _k, _v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + _k, _v = state.get_layer_view(logical_layer_id=self.layer_num_) _q = q.view(-1, self.tp_q_head_num_, self.head_dim_) - o_tensor = infer_state.decode_att_state.decode_att(q=_q, k=_k, v=_v, alloc_func=self.alloc_tensor) + o_tensor = infer_state.decode_att_state.decode_att( + q=_q, + k=_k, + v=_v, + att_control=AttControl(logical_layer_id=self.layer_num_), + alloc_func=self.alloc_tensor, + ) return o_tensor.view(q.shape) def _att_norm( diff --git a/lightllm/models/llama/model.py b/lightllm/models/llama/model.py index 2eab67f807..13036b0b83 100644 --- a/lightllm/models/llama/model.py +++ b/lightllm/models/llama/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import os import json import torch @@ -9,14 +10,18 @@ from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.common.basemodel import TpPartBaseModel -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args logger = init_logger(__name__) +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class LlamaTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = LlamaPreAndPostLayerWeight @@ -34,6 +39,39 @@ def __init__(self, kvargs): super().__init__(kvargs) return + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", + ) + return [state] + def _init_config(self): super()._init_config() # rename key @@ -52,19 +90,6 @@ def _verify_params(self): assert self.config["num_attention_heads"] % self.tp_world_size_ == 0 return - def _init_mem_manager(self): - head_dim_ = self.config["hidden_size"] // self.config["num_attention_heads"] - head_dim_ = self.config.get("head_dim", head_dim_) - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=head_dim_, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return - def _init_custom(self): """ 模型特殊的一些初始化 diff --git a/lightllm/models/mistral/model.py b/lightllm/models/mistral/model.py index 32b8c538d3..f906270d5c 100644 --- a/lightllm/models/mistral/model.py +++ b/lightllm/models/mistral/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import torch from lightllm.common.basemodel import TpPartBaseModel from lightllm.models.llama.layer_weights.transformer_layer_weight import LlamaTransformerLayerWeight @@ -6,11 +7,15 @@ from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.mistral.layer_infer.transformer_layer_infer import MistralTransformerLayerInfer -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.utils.envs_utils import get_env_start_args +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class MistralTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = LlamaPreAndPostLayerWeight @@ -28,6 +33,39 @@ def __init__(self, kvargs): super().__init__(kvargs) return + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", + ) + return [state] + def _verify_params(self): assert self.load_way in ["HF"], "mistral only supports HF format to load Now!" assert self.config["num_key_value_heads"] % self.tp_world_size_ == 0 @@ -38,20 +76,6 @@ def _init_custom(self): self._init_to_get_rotary() return - def _init_mem_manager(self): - # Dealing with head_dim_!=n_embed // num_attention_heads scenarios, such as mistral 13B - head_dim = self.config["hidden_size"] // self.config["num_attention_heads"] - head_dim = self.config.get("head_dim", head_dim) - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=head_dim, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return - def _init_to_get_rotary(self, default_base=10000): if self.config.get("rope_scaling", {}) is None: rope_scaling_factor = 1.0 diff --git a/lightllm/models/mistral_mtp/model.py b/lightllm/models/mistral_mtp/model.py index f17bc0a383..e8458846f9 100644 --- a/lightllm/models/mistral_mtp/model.py +++ b/lightllm/models/mistral_mtp/model.py @@ -41,14 +41,6 @@ def _init_custom(self): self._sin_cached = self.main_model._sin_cached return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None self.config["n_layer"] = 1 diff --git a/lightllm/models/mixtral/model.py b/lightllm/models/mixtral/model.py index 8eaeec4018..b749698baa 100644 --- a/lightllm/models/mixtral/model.py +++ b/lightllm/models/mixtral/model.py @@ -1,7 +1,7 @@ +from typing import TYPE_CHECKING import torch import numpy as np from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight @@ -14,6 +14,12 @@ logger = init_logger(__name__) +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class MixtralTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = LlamaPreAndPostLayerWeight @@ -35,6 +41,39 @@ def _init_config(self): super()._init_config() return + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", + ) + return [state] + def _verify_params(self): assert self.load_way in ["HF"], "mixtral only supports HF format to load Now!" assert self.config["num_key_value_heads"] % self.tp_world_size_ == 0 @@ -45,18 +84,6 @@ def _init_custom(self): self._init_to_get_rotary() return - def _init_mem_manager(self): - self.mem_manager = MemoryManager( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=self.config["hidden_size"] // self.config["num_attention_heads"], - layer_num=self.config["num_hidden_layers"], - always_copy=False, - mem_fraction=self.mem_fraction, - ) - return - def _init_to_get_rotary(self, default_base=10000): if self.config.get("rope_scaling", {}) is None: rope_scaling_factor = 1.0 diff --git a/lightllm/models/qwen2/model.py b/lightllm/models/qwen2/model.py index 82f16bc2f1..250cbbe461 100644 --- a/lightllm/models/qwen2/model.py +++ b/lightllm/models/qwen2/model.py @@ -1,8 +1,6 @@ from lightllm.models.qwen2.layer_weights.pre_and_post_layer_weight import Qwen2PreAndPostLayerWeight from lightllm.models.qwen2.layer_weights.transformer_layer_weight import Qwen2TransformerLayerWeight from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num class Qwen2TpPartModel(LlamaTpPartModel): @@ -35,18 +33,3 @@ def _init_some_value(self): self.layers_num = self.config["n_layer"] self.vocab_size = self.config["vocab_size"] return - - def _init_mem_manager(self): - head_dim_ = self.config["hidden_size"] // self.config["num_attention_heads"] - head_dim_ = self.config.get("head_dim", head_dim_) - tp_k_head_num_ = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=tp_k_head_num_, - head_dim=head_dim_, - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return diff --git a/lightllm/models/qwen3_5/model.py b/lightllm/models/qwen3_5/model.py index 3315831f53..1dac4c1f23 100644 --- a/lightllm/models/qwen3_5/model.py +++ b/lightllm/models/qwen3_5/model.py @@ -15,7 +15,6 @@ Qwen35TransformerLayerInfer, ) from lightllm.models.qwen3_5.infer_struct import Qwen35InferStateInfo -from lightllm.common.build_utils import repair_config from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -57,35 +56,37 @@ class Qwen3_5TpPartModel(Qwen3NextTpPartModel): infer_state_class = Qwen35InferStateInfo + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """展开文本配置并补齐 Qwen3.5 的 rotary 与 MoE 默认值。""" + config = super().normalize_model_config(model_config) + rope_parameters = config.get("rope_parameters") + if isinstance(rope_parameters, dict): + if "rope_theta" in rope_parameters and "rope_theta" not in config: + config["rope_theta"] = rope_parameters["rope_theta"] + if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in config: + config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] + if "rope_scaling" not in config: + config["rope_scaling"] = rope_parameters + + # MoE routing parameters - set defaults for Qwen3.5 compatibility + if "norm_topk_prob" not in config: + config["norm_topk_prob"] = True + + return config + def _init_config(self): config_path = os.path.join(self.weight_dir_, "config.json") with open(config_path, "r") as json_file: all_config = json.load(json_file) - self.config = all_config["text_config"] + self.config = self.normalize_model_config(all_config) self.vision_config = all_config.get("vision_config", None) if self.vision_config is None: logger.warning("No vision_config found in checkpoint. " "Multimodal features may not work correctly.") - repair_config(self.config, same_names=["num_attention_heads", "n_head"]) - repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(self.config, same_names=["num_hidden_layers", "n_layer"]) - - rope_parameters = self.config.get("rope_parameters") - if isinstance(rope_parameters, dict): - if "rope_theta" in rope_parameters and "rope_theta" not in self.config: - self.config["rope_theta"] = rope_parameters["rope_theta"] - if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in self.config: - self.config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] - if "rope_scaling" not in self.config: - self.config["rope_scaling"] = rope_parameters - - # MoE routing parameters - set defaults for Qwen3.5 compatibility - if "norm_topk_prob" not in self.config: - self.config["norm_topk_prob"] = True - if self.finetune_config: self.config["vocab_size"] = self.finetune_config.vocab_size self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) diff --git a/lightllm/models/qwen3_5_dflash/model.py b/lightllm/models/qwen3_5_dflash/model.py index 422ccf61cf..c51c3f4d90 100644 --- a/lightllm/models/qwen3_5_dflash/model.py +++ b/lightllm/models/qwen3_5_dflash/model.py @@ -10,17 +10,22 @@ class Qwen3_5DFlashModel(Qwen3DFlashModel): pre_and_post_weight_class = Qwen35DFlashPreAndPostLayerWeight - def _init_config(self): - super()._init_config() - self.config.update(self.config.get("dflash_config", {})) - - rope_parameters = self.config["rope_parameters"] - if "rope_theta" in rope_parameters and "rope_theta" not in self.config: - self.config["rope_theta"] = rope_parameters["rope_theta"] - if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in self.config: - self.config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] - if "rope_scaling" not in self.config: - self.config["rope_scaling"] = rope_parameters + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """合并 draft 专有配置，再同步别名及 rotary 参数；启动与 CPU cache 使用同一结果。""" + config = model_config.get("text_config", model_config).copy() + dflash_config = config.get("dflash_config", {}) + config.update(dflash_config) + config = super().normalize_model_config(config) + + rope_parameters = config["rope_parameters"] + if "rope_theta" in rope_parameters and "rope_theta" not in config: + config["rope_theta"] = rope_parameters["rope_theta"] + if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in config: + config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] + if "rope_scaling" not in config: + config["rope_scaling"] = rope_parameters + return config def _init_custom(self): # Draft and target use different rotary shapes, so the draft owns its rotary cache. @@ -28,19 +33,6 @@ def _init_custom(self): self.block_size = self.config["block_size"] self.mask_token_id = self.config["mask_token_id"] - 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, - target_mem_manager.head_dim, - ) - assert draft_kv_shape == target_kv_shape, ( - "Qwen3.5 parallel block drafter requires matching draft and target KV shapes, " - f"got draft={draft_kv_shape}, target={target_kv_shape}." - ) - super()._init_mem_manager() - def _init_weights(self, start_layer_index=None): super()._init_weights(start_layer_index=start_layer_index) self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ diff --git a/lightllm/models/qwen3_5_dspark/model.py b/lightllm/models/qwen3_5_dspark/model.py index 896949d1cf..d81d5a9b2a 100644 --- a/lightllm/models/qwen3_5_dspark/model.py +++ b/lightllm/models/qwen3_5_dspark/model.py @@ -5,35 +5,27 @@ class Qwen3_5DSparkModel(Qwen3DSparkModel): """Adapter for the current Qwen3 DSpark checkpoint with a Qwen3.5 target.""" - def _init_config(self): - super()._init_config() - self.config.update(self.config.get("dflash_config", {})) + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """合并 draft 专有配置，再同步别名及 rotary 参数；启动与 CPU cache 使用同一结果。""" + config = model_config.get("text_config", model_config).copy() + dflash_config = config.get("dflash_config", {}) + config.update(dflash_config) + config = super().normalize_model_config(config) - rope_parameters = self.config["rope_parameters"] - if "rope_theta" in rope_parameters and "rope_theta" not in self.config: - self.config["rope_theta"] = rope_parameters["rope_theta"] + rope_parameters = config["rope_parameters"] + if "rope_theta" in rope_parameters and "rope_theta" not in config: + config["rope_theta"] = rope_parameters["rope_theta"] # The draft is Qwen3-style and owns a 1D rotary cache. Released # checkpoints rotate the full head, while target-shaped custom # checkpoints can retain Qwen3.5's partial rotary layout. - self.config["rope_scaling"] = rope_parameters - self.config["partial_rotary_factor"] = rope_parameters.get("partial_rotary_factor", 1.0) + config["rope_scaling"] = rope_parameters + config["partial_rotary_factor"] = rope_parameters.get("partial_rotary_factor", 1.0) + return config def _init_custom(self): # Draft and target use different rotary shapes, so the draft owns its rotary cache. LlamaTpPartModel._init_custom(self) self.block_size = self.config["block_size"] self.mask_token_id = self.config["mask_token_id"] - - 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, - target_mem_manager.head_dim, - ) - assert draft_kv_shape == target_kv_shape, ( - "Qwen3.5 parallel block drafter requires matching draft and target KV shapes, " - f"got draft={draft_kv_shape}, target={target_kv_shape}." - ) - super()._init_mem_manager() diff --git a/lightllm/models/qwen3_5_mtp/model.py b/lightllm/models/qwen3_5_mtp/model.py index b0f55af7f3..3a0ef250fc 100644 --- a/lightllm/models/qwen3_5_mtp/model.py +++ b/lightllm/models/qwen3_5_mtp/model.py @@ -14,7 +14,7 @@ class Qwen3_5MTPModel(Qwen3_5TpPartModel): transformer_weight_class = Qwen3_5MTPTransformerLayerWeight transformer_layer_infer_class = Qwen35TransformerLayerInfer - # MTP draft model: reuses the main model's req/mem managers and rope caches, and is + # MTP draft model: reuses the main model's request/attention resources and rope caches, and is # marked so the decode CUDA-graph / padding paths detect it (is_mtp_draft_model). is_mtp_draft_model = True @@ -29,22 +29,12 @@ def _init_custom(self): self._sin_cached = self.main_model._sin_cached return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - - def _init_config(self): - super()._init_config() - # MTP draft model: reuses the main model's config, but overrides the following: - # 因为 qwen3.5 的 mtp 和 main 是存储在一起的，所以需要进行修复。 - self.config["full_attention_interval"] = 1 - self.config["num_hidden_layers"] = 1 - self.config["n_layer"] = 1 - return + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """原生 MTP 只包含一个 full-attention 层，与 main 共用 checkpoint。""" + config = super().normalize_model_config(model_config) + config.update(full_attention_interval=1, num_hidden_layers=1, n_layer=1, layer_types=["full_attention"]) + return config def _init_some_value(self): super()._init_some_value() diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 9c4a42fffc..f0683dd882 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -46,13 +46,6 @@ def _init_custom(self): self.block_size = int(self.config["block_size"]) self.mask_token_id = int(self.config["mask_token_id"]) - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - - def _init_mem_manager(self): - # Draft KV uses target-owned cache slots. - self.mem_manager = self.main_model.mem_manager - def _init_att_backend(self): super()._init_att_backend() # FA3 is currently the only backend that supports non-causal block decode. @@ -98,7 +91,9 @@ def _decode(self, model_input: ModelInput) -> ModelOutput: infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens infer_state.position_cos = torch.index_select(self._cos_cached, 0, position_ids) infer_state.position_sin = torch.index_select(self._sin_cached, 0, position_ids) - infer_state.mem_manager = self.mem_manager + infer_state.att_state_manager = self.att_state_manager + infer_state.att_state_role = self.att_state_role + infer_state.is_draft_model = True infer_state.mem_index = self._select_mem_indexes(model_input) # 这里调用的是 DFlash 覆写后的 context_forward：pre layer 只投影 target hidden， diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index b03fd77cc2..7a631622d6 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.common.basemodel.attention import AttControl from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer from lightllm.models.qwen3_dflash2.triton_kernel import grouped_dynamic_conv @@ -13,7 +14,13 @@ def __init__(self, layer_num, network_config): self.sliding_window_ = int(network_config.get("sliding_window", 0) or 0) def _token_attention_kernel(self, q, infer_state, layer_weight): - k, v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + k, v = state.get_layer_view(logical_layer_id=self.layer_num_) k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) q = q.view(-1, self.tp_q_head_num_, self.head_dim_) @@ -26,6 +33,7 @@ def _token_attention_kernel(self, q, infer_state, layer_weight): k=k, v=v, att_control=AttControl( + logical_layer_id=self.layer_num_, use_sliding_window=use_sliding_window, sliding_window=(window, window) if use_sliding_window else (-1, -1), ), @@ -34,8 +42,14 @@ def _token_attention_kernel(self, q, infer_state, layer_weight): return output.view(-1, self.tp_q_head_num_ * self.head_dim_) def _post_cache_kv(self, cache_kv, infer_state, layer_weight): - storage_head_num = infer_state.mem_manager.head_num - storage_head_dim = infer_state.mem_manager.head_dim + state = infer_state.att_state_manager.get_state( + model_role=infer_state.att_state_role, + state_type_filter="gqa", + logical_layer_id=self.layer_num_, + state_class=FullAttentionState, + ) + storage_head_num = state.config.head_num + storage_head_dim = state.config.head_dim storage_kv = cache_kv.view(cache_kv.shape[0], 2 * storage_head_num, storage_head_dim) return super()._post_cache_kv(storage_kv, infer_state, layer_weight) diff --git a/lightllm/models/qwen3_dflash2/model.py b/lightllm/models/qwen3_dflash2/model.py index 3542b37188..e43a534b1f 100644 --- a/lightllm/models/qwen3_dflash2/model.py +++ b/lightllm/models/qwen3_dflash2/model.py @@ -18,20 +18,24 @@ class Qwen3DFlash2Model(Qwen3DFlashModel): post_layer_infer_class = Qwen3DFlash2PostLayerInfer transformer_layer_infer_class = Qwen3DFlash2TransformerLayerInfer - def _init_config(self): - super()._init_config() - dflash_config = self.config.get("dflash_config", {}) + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + """合并 draft 专有配置，再同步别名及 rotary 参数；启动与 CPU cache 使用同一结果。""" + config = model_config.get("text_config", model_config).copy() + dflash_config = config.get("dflash_config", {}) if not isinstance(dflash_config, dict): raise ValueError("dflash_config must be an object in the DFlash2 checkpoint config") - self.config.update(dflash_config) + config.update(dflash_config) + config = super().normalize_model_config(config) - rope_parameters = self.config.get("rope_parameters", {}) - if "rope_theta" in rope_parameters and "rope_theta" not in self.config: - self.config["rope_theta"] = rope_parameters["rope_theta"] - if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in self.config: - self.config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] - if rope_parameters and "rope_scaling" not in self.config: - self.config["rope_scaling"] = rope_parameters + rope_parameters = config.get("rope_parameters", {}) + if "rope_theta" in rope_parameters and "rope_theta" not in config: + config["rope_theta"] = rope_parameters["rope_theta"] + if "partial_rotary_factor" in rope_parameters and "partial_rotary_factor" not in config: + config["partial_rotary_factor"] = rope_parameters["partial_rotary_factor"] + if rope_parameters and "rope_scaling" not in config: + config["rope_scaling"] = rope_parameters + return config def _verify_params(self): LlamaTpPartModel._verify_params(self) @@ -61,18 +65,6 @@ def _init_custom(self): self.block_size = int(self.config["block_size"]) self.mask_token_id = int(self.config["mask_token_id"]) - def _init_mem_manager(self): - main_mem_manager = self.main_model.mem_manager - draft_head_num = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - draft_width = draft_head_num * self.config["head_dim"] - target_width = main_mem_manager.head_num * main_mem_manager.head_dim - assert draft_width == target_width, ( - "DFlash2 draft and target KV must have equal flat widths: " - f"draft=({draft_head_num}, {self.config['head_dim']}), " - f"target=({main_mem_manager.head_num}, {main_mem_manager.head_dim})" - ) - self.mem_manager = main_mem_manager - def _init_weights(self, start_layer_index=None): super()._init_weights(start_layer_index=start_layer_index) self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ diff --git a/lightllm/models/qwen3_eagle/model.py b/lightllm/models/qwen3_eagle/model.py index d6257d3ebf..42d273b31f 100644 --- a/lightllm/models/qwen3_eagle/model.py +++ b/lightllm/models/qwen3_eagle/model.py @@ -46,12 +46,6 @@ def _init_custom(self): self._cos_cached = self.main_model._cos_cached self._sin_cached = self.main_model._sin_cached - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - def _init_weights(self, start_layer_index=None): self.pre_post_weight = self.pre_and_post_weight_class( self.data_type, network_config=self.config, quant_cfg=self.quant_cfg diff --git a/lightllm/models/qwen3_moe_mtp/model.py b/lightllm/models/qwen3_moe_mtp/model.py index d9854250e2..da6a1dc76f 100644 --- a/lightllm/models/qwen3_moe_mtp/model.py +++ b/lightllm/models/qwen3_moe_mtp/model.py @@ -33,14 +33,6 @@ def _init_custom(self): self._sin_cached = self.main_model._sin_cached return - def _init_req_manager(self): - self.req_manager = self.main_model.req_manager - return - - def _init_mem_manager(self): - self.mem_manager = self.main_model.mem_manager - return - def _init_weights(self, start_layer_index=None): assert start_layer_index is None mtp_index = len(self.mtp_previous_draft_models) diff --git a/lightllm/models/qwen3next/__init__.py b/lightllm/models/qwen3next/__init__.py index a9d22c6643..482161743d 100644 --- a/lightllm/models/qwen3next/__init__.py +++ b/lightllm/models/qwen3next/__init__.py @@ -1,3 +1,9 @@ -from lightllm.models.qwen3next.model import Qwen3NextTpPartModel - __all__ = ["Qwen3NextTpPartModel"] + + +def __getattr__(name): + if name == "Qwen3NextTpPartModel": + from .model import Qwen3NextTpPartModel + + return Qwen3NextTpPartModel + raise AttributeError(name) diff --git a/lightllm/models/qwen3next/config.py b/lightllm/models/qwen3next/config.py new file mode 100644 index 0000000000..d647096564 --- /dev/null +++ b/lightllm/models/qwen3next/config.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +import torch +import dataclasses + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.torch_dtype_utils import get_torch_dtype + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + + +@dataclasses.dataclass +class Qwen3NextResourceConfig: + # 单个模型的 TP 组大小，通常为 args.tp // args.dp，不包含其他 DP 组。 + tp_world_size: int + + # TP 切分前的全局 KV head 数，K/V 各有这么多个，不相加；GLM 的 MLA latent 记为 1。 + full_att_all_num_kv_heads: int + # 未打包的 KV 特征维度：GQA 为单个 K/V head 的宽度，GLM MLA 为 latent 宽度。 + # 量化、indexer 字段及行 padding 由具体 full-attention State 确定，不计入本字段。 + full_att_head_dim: int + + # TP 切分前的线性 attention K head 数；Q 使用相同的 head 数和维度。 + global_linear_k_heads: int + # TP 切分前的 V head 数；GDN 可与 K 不同，GLM KDA 与 K 相同。 + global_linear_v_heads: int + # 每个 Q/K head 的特征维度，对应 SSM 的 key 轴，不随 TP 切分。 + head_linear_k_dim: int + # 每个 V head 的特征维度，对应 SSM 的 value 轴；KDA 中与 key 维度相等。 + head_linear_v_dim: int + # 原始卷积核宽度，单位为 token；基准 Conv 历史保存 kernel_size - 1 个位置。 + # MTP 的额外候选位置由 State.mtp_step 决定，不计入本字段。 + conv_kernel_size: int + # Conv 运行态和大小页快照的元素 dtype，不代表算子的内部累加精度。 + conv_state_dtype: torch.dtype + # SSM 运行态和大小页快照的元素 dtype，可独立于 Conv 精度配置。 + ssm_state_dtype: torch.dtype + # 混合层排列周期：每隔这么多层放一个 full-attention 层，零起始层号为 interval - 1。 + full_attention_interval: int + # 当前模型总层数，包含线性与 full-attention 两类层，不合并其他模型的层。 + all_layer_num: int + # 显式执行层类型，支持 Loop 等非周期排列；为空时沿用周期配置。 + layer_types: tuple[str, ...] = () + + def get_full_att_layer_indices(self) -> tuple[int, ...]: + """返回当前模型内部的 full-attention 层序号，不包含其他模型的层。 + + 调用方通过 model_logical_layer_ids[index] 转换为实际逻辑层 ID； + 物理存储层映射由 State 管理，draft 使用自身配置独立创建 State。 + """ + if self.layer_types: + return tuple(index for index, kind in enumerate(self.layer_types) if kind != "linear_attention") + return tuple(range(self.full_attention_interval - 1, self.all_layer_num, self.full_attention_interval)) + + @classmethod + def load_from_args(cls) -> "Qwen3NextResourceConfig": + """从启动参数加载模型配置，供 attention 后端校验使用。""" + args = get_env_start_args() + model_path = args.model_dir + from transformers.configuration_utils import PretrainedConfig + + model_cfg, _ = PretrainedConfig.get_config_dict(model_path) + return cls.from_model_config(model_cfg, args) + + @classmethod + def from_model_config( + cls, + model_cfg: dict[str, object], + args: StartArgs | EasyDict, + ) -> "Qwen3NextResourceConfig": + """解析当前模型的全局几何参数，TP 分片与实际存储由 State 确定。""" + llm_config = model_cfg.get("text_config", model_cfg) + n_layer = llm_config["num_hidden_layers"] + + tp_world_size = args.tp // args.dp + return cls( + tp_world_size=tp_world_size, + full_att_all_num_kv_heads=llm_config["num_key_value_heads"], + full_att_head_dim=llm_config["head_dim"], + global_linear_k_heads=llm_config["linear_num_key_heads"], + global_linear_v_heads=llm_config["linear_num_value_heads"], + head_linear_k_dim=llm_config["linear_key_head_dim"], + head_linear_v_dim=llm_config["linear_value_head_dim"], + conv_kernel_size=llm_config["linear_conv_kernel_dim"], + conv_state_dtype=get_torch_dtype(args.data_type), + ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), + full_attention_interval=llm_config["full_attention_interval"], + all_layer_num=n_layer, + layer_types=tuple(llm_config.get("layer_types", ())), + ) 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 444ed1f393..f1a3a65748 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -1,5 +1,5 @@ import torch -from typing import Optional +from typing import Optional, TYPE_CHECKING import triton from lightllm.models.qwen3_moe.model import Qwen3MOEModel from lightllm.models.qwen3next.layer_weights.transformer_layer_weight import ( @@ -11,20 +11,20 @@ ) from lightllm.models.qwen3next.infer_struct import Qwen3NextInferStateInfo from lightllm.utils.log_utils import init_logger -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( - FP8StaticPerHeadQuantQwen3NextMemManager, - FP8StaticPerTensorQuantQwen3NextMemManager, - Qwen3NextMemManager, -) -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.common.req_manager import ReqManagerForMamba -from lightllm.common.state_cache_manager import LinearAttCacheConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState logger = init_logger(__name__) class Qwen3NextTpPartModel(Qwen3MOEModel): + is_hybrid_att_model = True + # weight class pre_and_post_weight_class = Qwen3NextPreAndPostLayerWeight transformer_weight_class = Qwen3NextTransformerLayerWeight @@ -39,6 +39,51 @@ def __init__(self, kvargs) -> None: self._init_triton() super().__init__(kvargs) + @classmethod + def create_bounded_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + request_capacity: int = 1, + tp_rank: int = 0, + tp_world_size: int = 1, + device: str = "meta", + ) -> list["BoundedAttentionState"]: + """创建当前模型的 bounded 运行态；CPU 布局查询传入 meta 设备。""" + from lightllm.common.attention_state.bounded_att import create_gdn_state, ConvSSMStateConfig + + config = cls.normalize_model_config(model_config) + assert config["num_attention_heads"] % tp_world_size == 0 + assert not ( + args.mtp_step and args.enable_prefill_decode_mixed + ), "MTP linear attention does not support mixed prefill/decode" + cfg = Qwen3NextResourceConfig.from_model_config(config, args) + full_layer_indices = set(cfg.get_full_att_layer_indices()) + logical_layer_ids = tuple( + layer for i, layer in enumerate(model_logical_layer_ids) if i not in full_layer_indices + ) + if not logical_layer_ids: + return [] + state = create_gdn_state( + config=ConvSSMStateConfig( + logical_layer_ids=logical_layer_ids, + global_k_heads=cfg.global_linear_k_heads, + global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, + value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, + conv_dtype=cfg.conv_state_dtype, + ssm_dtype=cfg.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=args.mtp_step, + device=device, + ) + return [state] + def _init_triton(self): def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: return torch.empty(size, device="cuda", dtype=torch.int8) @@ -56,45 +101,43 @@ 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() - if start_args.llm_kv_type == "None": - mem_manager_class = Qwen3NextMemManager - elif start_args.llm_kv_type == "fp8kv_sph": - mem_manager_class = FP8StaticPerHeadQuantQwen3NextMemManager - elif start_args.llm_kv_type == "fp8kv_spt": - mem_manager_class = FP8StaticPerTensorQuantQwen3NextMemManager - else: - raise ValueError(f"Qwen3Next does not support llm_kv_type={start_args.llm_kv_type}") - self.mem_manager = mem_manager_class( - size=self.max_total_token_num, - dtype=self.data_type, - num_kv_heads=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, - mem_fraction=self.mem_fraction, + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按 Qwen 混合层周期选择 full-attention 层，保留全局逻辑层号。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + assert args.llm_kv_type in ("None", "fp8kv_sph", "fp8kv_spt"), "unsupported Qwen3Next KV encoding" + cfg = Qwen3NextResourceConfig.from_model_config(config, args) + logical_layer_groups = tuple((model_logical_layer_ids[index],) for index in cfg.get_full_att_layer_indices()) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=cfg.full_att_all_num_kv_heads, + head_dim=cfg.full_att_head_dim, + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", ) - - def _init_req_manager(self): - self.linear_config = LinearAttCacheConfig.load_from_args() - 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=self.linear_config - ) - return + return [state] def _init_att_backend1(self): - # MTP draft models only keep full-attention layers and reuse the main - # model's managers, so they do not need LinearAttBackend. - if getattr(self, "is_mtp_draft_model", False): + role = self.att_state_role + if not self.att_state_manager.get_bounded_attention_states(role): self.prefill_att_backend1 = None self.decode_att_backend1 = None return diff --git a/lightllm/models/sensenova_loop/model.py b/lightllm/models/sensenova_loop/model.py index ce6d36a7d5..ea882ba450 100644 --- a/lightllm/models/sensenova_loop/model.py +++ b/lightllm/models/sensenova_loop/model.py @@ -1,7 +1,6 @@ """SenseNova Loop: shared checkpoint weights and independent execution caches.""" from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel -from lightllm.utils.config_utils import get_config_json from .config import get_loop_layer_indices, get_loop_layer_types from .layer_infer.transformer_layer_infer import SenseNovaLoopTransformerLayerInfer from .layer_weights.transformer_layer_weight import SenseNovaLoopTransformerLayerWeight @@ -13,20 +12,26 @@ class SenseNovaLoopTpPartModel(Qwen3_5TpPartModel): transformer_weight_class = SenseNovaLoopTransformerLayerWeight transformer_layer_infer_class = SenseNovaLoopTransformerLayerInfer + @classmethod + def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, object]: + config = super().normalize_model_config(model_config) + if "loop_physical_indices" not in config: + settings = model_config["loop_settings"] + indices = get_loop_layer_indices(settings, config["num_hidden_layers"]) + config.update( + layer_types=get_loop_layer_types(settings, config), + n_layer=len(indices), + num_hidden_layers=len(indices), + loop_physical_indices=indices, + ) + if model_config.get("mhc") is not None: + config["mhc"] = model_config["mhc"] + return config + def _init_config(self): super()._init_config() - model_config = get_config_json(self.weight_dir_) - settings = model_config["loop_settings"] - self.num_physical_layers = self.config["n_layer"] - self.loop_layer_indices = get_loop_layer_indices(settings, self.num_physical_layers) - self.config.update( - layer_types=get_loop_layer_types(settings, self.config), - n_layer=len(self.loop_layer_indices), - num_hidden_layers=len(self.loop_layer_indices), - loop_physical_indices=self.loop_layer_indices, - ) - if model_config.get("mhc") is not None: - self.config["mhc"] = model_config["mhc"] + self.loop_layer_indices = self.config["loop_physical_indices"] + self.num_physical_layers = max(self.loop_layer_indices) + 1 def _init_weights(self, start_layer_index=0): # Allocate and preprocess each checkpoint layer once, then alias weights. diff --git a/lightllm/models/starcoder/model.py b/lightllm/models/starcoder/model.py index de7eeac78b..39a90415bc 100644 --- a/lightllm/models/starcoder/model.py +++ b/lightllm/models/starcoder/model.py @@ -1,15 +1,20 @@ +from typing import TYPE_CHECKING from lightllm.models.starcoder.layer_infer.transformer_layer_infer import StarcoderTransformerLayerInfer from lightllm.models.starcoder.layer_infer.pre_layer_infer import StarcoderPreLayerInfer from lightllm.models.starcoder.layer_weights.transformer_layer_weight import StarcoderTransformerLayerWeight from lightllm.models.starcoder.layer_weights.pre_and_post_layer_weight import StarcoderPreAndPostLayerWeight from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer from lightllm.common.build_utils import repair_config -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel import InferStateInfo +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class StarcoderTpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = StarcoderPreAndPostLayerWeight @@ -24,6 +29,40 @@ class StarcoderTpPartModel(TpPartBaseModel): def __init__(self, kvargs): super().__init__(kvargs) + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """StarCoder 的单份 KV 在各 TP rank 复制，逻辑层由本模型直接定义。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=1, + head_dim=config["hidden_size"] // config["num_attention_heads"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + replicated=True, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", + ) + return [state] + def _init_config(self): super()._init_config() # rename key @@ -39,17 +78,6 @@ def _reset_num_key_value_heads(self): def _verify_params(self): assert self.load_way == "HF", "StarCoder only support HF format to load Now!" - def _init_mem_manager(self): - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"], - head_dim=self.config["hidden_size"] // self.config["num_attention_heads"], - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return - def _init_some_value(self): super()._init_some_value() self.tp_k_head_num_ = self.config["num_key_value_heads"] diff --git a/lightllm/models/starcoder2/model.py b/lightllm/models/starcoder2/model.py index 09107103a4..50fc740e04 100644 --- a/lightllm/models/starcoder2/model.py +++ b/lightllm/models/starcoder2/model.py @@ -1,3 +1,4 @@ +from typing import TYPE_CHECKING import torch from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.models.starcoder2.layer_weights.pre_and_post_layer_weight import Starcoder2PreAndPostLayerWeight @@ -7,11 +8,15 @@ from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer from lightllm.common.build_utils import repair_config -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num from lightllm.common.basemodel import TpPartBaseModel +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.common.attention_state.full_att import FullAttentionState + + class Starcoder2TpPartModel(TpPartBaseModel): # weight class pre_and_post_weight_class = Starcoder2PreAndPostLayerWeight @@ -27,6 +32,39 @@ def __init__(self, kvargs): super().__init__(kvargs) return + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: "StartArgs | EasyDict", + *, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list["FullAttentionState"]: + """按本模型的全局逻辑层创建 GQA 布局，量化存储由算法入口选择。""" + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + config = cls.normalize_model_config(model_config) + logical_layer_groups = tuple((layer_id,) for layer_id in model_logical_layer_ids) + if not logical_layer_groups: + return [] + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=logical_layer_groups, + global_head_num=config.get("num_key_value_heads", config["num_attention_heads"]), + head_dim=config.get("head_dim", config["hidden_size"] // config["num_attention_heads"]), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=args.llm_kv_type, + page_size=args.page_size, + device="meta", + ) + return [state] + def _init_config(self): super()._init_config() repair_config(self.config, same_names=["norm_epsilon", "rms_norm_eps", "layer_norm_epsilon"]) @@ -45,17 +83,6 @@ def _init_custom(self): self._init_to_get_rotary() return - def _init_mem_manager(self): - self.mem_manager = select_mem_manager_class()( - self.max_total_token_num, - dtype=self.data_type, - head_num=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=self.config["hidden_size"] // self.config["num_attention_heads"], - layer_num=self.config["num_hidden_layers"] + get_added_mtp_kv_layer_num(), - mem_fraction=self.mem_fraction, - ) - return - def _init_some_value(self): super()._init_some_value() return diff --git a/lightllm/server/multi_level_kv_cache/cache_placement_controller.py b/lightllm/server/multi_level_kv_cache/cache_placement_controller.py index 2bd67e4f93..43b3eea996 100644 --- a/lightllm/server/multi_level_kv_cache/cache_placement_controller.py +++ b/lightllm/server/multi_level_kv_cache/cache_placement_controller.py @@ -229,7 +229,7 @@ def create_cache_placement_controller( if args.enable_disk_cache: disk_capacity_bytes = int(args.disk_cache_storage_size * (1024 ** 3)) - disk_capacity_tokens = disk_capacity_bytes // cpu_cache_meta.calcu_one_page_size() * page_size_tokens + disk_capacity_tokens = disk_capacity_bytes // cpu_cache_meta.page_bytes * page_size_tokens else: disk_capacity_tokens = 0 diff --git a/lightllm/server/multi_level_kv_cache/cpu_cache_client.py b/lightllm/server/multi_level_kv_cache/cpu_cache_client.py index fb9e74c47b..8138b798f6 100644 --- a/lightllm/server/multi_level_kv_cache/cpu_cache_client.py +++ b/lightllm/server/multi_level_kv_cache/cpu_cache_client.py @@ -1,4 +1,5 @@ import ctypes +import torch from lightllm.utils.envs_utils import get_env_start_args, get_disk_cache_prompt_limit_length from typing import List, Optional, Tuple from lightllm.utils.log_utils import init_logger @@ -27,14 +28,11 @@ def __init__(self, only_create_meta_data: bool, init_shm_data: bool): tensor_spec = CpuCacheTensorSpec( shm_key=self.args.cpu_kv_cache_shm_id, shape=( - self.kv_cache_tensor_meta.page_num, - self.kv_cache_tensor_meta.layer_num, - self.kv_cache_tensor_meta.token_page_size, - self.kv_cache_tensor_meta.num_heads, - self.kv_cache_tensor_meta.get_merged_head_dim(), + self.page_num, + self.kv_cache_tensor_meta.page_bytes, ), - dtype=self.kv_cache_tensor_meta.data_type, - size_bytes=self.kv_cache_tensor_meta.calcu_size(), + dtype=torch.uint8, + size_bytes=self.kv_cache_tensor_meta.size_bytes, ) tensor_creator = CpuCacheCreator(tensor_spec=tensor_spec) self.cpu_kv_cache_tensor = tensor_creator.create_or_attach( diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index 95f0e4daec..52e27992ec 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -192,8 +192,9 @@ class PDChunckedTransTask: transfer_time_out_secs: int = 66 # kv: 通过 mem_indexes 传输 [start_kv_index, end_kv_index) 的 token KV。 # att_state: 混合注意力模型的请求运行态 buffer（如 linear attention 的 conv/SSM 状态）。 - # start_kv_index == end_kv_index 标记状态对应的 token 位置，mem_indexes 为空。 - # 通过本地 req_idx 定位运行态 buffer，具体打包和恢复由模型 mem_manager 负责。 + # start_kv_index == end_kv_index 标记 prefill 状态对应的 token 位置，mem_indexes 为空。 + # tail PD 读写用该位置在固定逻辑窗口与本地 MTP ring 之间转换，不另存长度字段。 + # 通过本地 req_idx 定位运行态 buffer，具体打包和恢复由模型 AttentionStateManager 负责。 page_kind: str = "kv" # Only valid for the local task owner; remote notify copies may carry the sender-local req_idx. req_idx: Optional[int] = None 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..29898f122a 100644 --- a/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/hybrid_att_radix_cache.py @@ -2,7 +2,6 @@ import numpy as np from typing import Tuple, Dict, Set, List, Optional, Sequence from sortedcontainers import SortedSet, SortedDict -from lightllm.common.state_cache_manager import StateCacheManager from lightllm.utils.multimodal_chunk_utils import clamp_mm_cache_len from .shared_arr import SharedArray from .radix_cache import time_gen @@ -117,11 +116,8 @@ def __init__( rank_in_node: int, hash_page_size: int, big_page_num: int, - kv_cache_mem_manager=None, - small_page_buffers=None, + att_state_manager, ): - from lightllm.common.kv_cache_mem_manager import MemoryManager - assert hash_page_size >= 1, "hash_page_size must be >= 1" assert big_page_num >= 1, "big_page_num must be >= 1" @@ -130,9 +126,7 @@ def __init__( self.big_page_tokens = hash_page_size * big_page_num self.total_token_num = total_token_num - self.mem_manager: MemoryManager = kv_cache_mem_manager - - self.big_page_buffers: StateCacheManager = self.mem_manager.big_page_buffers + self.att_state_manager = att_state_manager self._key_dtype = torch.int64 self._value_dtype = torch.int64 @@ -151,7 +145,15 @@ def __init__( self.refed_tokens_num.arr[0] = 0 self.tree_total_tokens_num = SharedArray(f"tree_total_tokens_num_{rank_in_node}", (1,), dtype=np.int64) self.tree_total_tokens_num.arr[0] = 0 - self.small_page_buffers: StateCacheManager = small_page_buffers + + def _release_pages(self, page_kind, page_slot_ids): + resources = self.att_state_manager + allocator = ( + resources.bounded_att_big_page_allocator + if page_kind == "big" + else resources.bounded_att_small_page_allocator + ) + allocator.free(page_slot_ids) def _discard_node(self, node: HybridAttPagedTreeNode): if node.is_leaf(): @@ -248,7 +250,7 @@ def _insert_helper( 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: # 因为节点已经存在，所以无法插入，但是要释放对应的buffer_id 节点 - self.big_page_buffers.free_state_cache([new_big_page_buffer_id]) + self._release_pages("big", [new_big_page_buffer_id]) # 已经存在了 sub_prefix_len, ans_node = self._insert_helper( @@ -303,7 +305,7 @@ def _insert_helper( self._add_node(child) else: # 节点已有 checkpoint，释放本次重复申请的槽位。 - self.small_page_buffers.free_state_cache(free_indexes=[block_state_idxs[0]]) + self._release_pages("small", page_slot_ids=[block_state_idxs[0]]) sub_prefix_len, ans_node = self._insert_helper( child, @@ -585,9 +587,8 @@ def _print_helper(self, node: HybridAttPagedTreeNode, indent): return def free_radix_cache_to_get_enough_token(self, need_token_num): - assert self.mem_manager is not None - if need_token_num > self.mem_manager.allocator.can_use_mem_size: - need_evict_token_num = need_token_num - self.mem_manager.allocator.can_use_mem_size + if need_token_num > self.att_state_manager.full_att_allocator.can_use_mem_size: + need_evict_token_num = need_token_num - self.att_state_manager.full_att_allocator.can_use_mem_size release_mems = [] small_page_buffer_ids = [] @@ -598,16 +599,17 @@ def release_mem(mem_index, small_page_buffer_id): self._evict(need_evict_token_num, release_mem) mem_index = torch.concat(release_mems) - self.mem_manager.free(mem_index) + self.att_state_manager.full_att_allocator.free(mem_index) small_page_buffer_ids = [idx for idx in small_page_buffer_ids if idx is not None] if len(small_page_buffer_ids) > 0: - self.small_page_buffers.free_state_cache(small_page_buffer_ids) + self._release_pages("small", small_page_buffer_ids) return def free_one_small_page_buffer(self): - if self.small_page_buffers is None: + allocator = self.att_state_manager.bounded_att_small_page_allocator + if allocator is None: return - if self.small_page_buffers.get_free_cache_num() > 0: + if allocator.free_count > 0: return if len(self._evict_tree_set_for_state_cache) == 0: return @@ -616,7 +618,7 @@ def free_one_small_page_buffer(self): self._discard_node(node) assert node.small_page_buffer_idx is not None - self.small_page_buffers.free_state_cache(free_indexes=[node.small_page_buffer_idx]) + self._release_pages("small", page_slot_ids=[node.small_page_buffer_idx]) node.small_page_buffer_idx = None self._add_node(node) @@ -639,7 +641,7 @@ def _evict(self, need_remove_tokens, evict_callback): if node.is_big_page_node(): assert node.big_page_buffer_idx is not None - self.big_page_buffers.free_state_cache([node.big_page_buffer_idx]) + self._release_pages("big", [node.big_page_buffer_idx]) evict_callback(node.token_mem_index_value, node.small_page_buffer_idx) self.tree_total_tokens_num.arr[0] -= len(node.token_mem_index_value) diff --git a/lightllm/server/router/dynamic_prompt/radix_cache.py b/lightllm/server/router/dynamic_prompt/radix_cache.py index 25df6d5ef7..41a292a8a0 100644 --- a/lightllm/server/router/dynamic_prompt/radix_cache.py +++ b/lightllm/server/router/dynamic_prompt/radix_cache.py @@ -110,19 +110,16 @@ def match(t1: torch.Tensor, t2: torch.Tensor) -> int: class RadixCache: - def __init__(self, total_token_num, rank_in_node, mem_manager=None, page_size: int = 1): - from lightllm.common.kv_cache_mem_manager import MemoryManager - + def __init__(self, total_token_num, rank_in_node, att_state_manager=None, page_size: int = 1): self.total_token_num = total_token_num - self.mem_manager: MemoryManager = mem_manager + self.att_state_manager = att_state_manager self._key_dtype = torch.int64 self._value_dtype = torch.int64 if page_size < 1: raise ValueError(f"page_size must be >= 1, got {page_size}") - if mem_manager is not None and page_size != mem_manager.page_size: - raise ValueError( - f"RadixCache page_size {page_size} must match mem_manager page_size {mem_manager.page_size}" - ) + allocator = None if att_state_manager is None else att_state_manager.full_att_allocator + if allocator is not None and page_size != allocator.page_size: + raise ValueError(f"RadixCache page_size {page_size} must match allocator page_size {allocator.page_size}") self.page_size = page_size self.root_node = TreeNode(page_size=page_size) @@ -546,9 +543,9 @@ def _print_helper(self, node: TreeNode, indent): return def free_radix_cache_to_get_enough_token(self, need_token_num): - assert self.mem_manager is not None - if need_token_num > self.mem_manager.allocator.can_use_mem_size: - need_evict_token_num = need_token_num - self.mem_manager.allocator.can_use_mem_size + allocator = self.att_state_manager.full_att_allocator + if need_token_num > allocator.can_use_mem_size: + need_evict_token_num = need_token_num - allocator.can_use_mem_size release_mems = [] def release_mem(mem_index): @@ -557,7 +554,7 @@ def release_mem(mem_index): self.evict(need_evict_token_num, release_mem) mem_index = torch.concat(release_mems) - self.mem_manager.free(mem_index) + allocator.free(mem_index) return diff --git a/lightllm/server/router/manager.py b/lightllm/server/router/manager.py index 189d3b6b78..8513fc030a 100644 --- a/lightllm/server/router/manager.py +++ b/lightllm/server/router/manager.py @@ -27,7 +27,7 @@ from lightllm.utils.profiler import ProfilerCmd from lightllm.server.router.token_load import TokenLoad from lightllm.server.metrics.manager import MetricClient -from lightllm.common.kv_cache_mem_manager import ReadOnlyStaticsMemoryManager +from lightllm.common.attention_state.allocators import KVCapacityReader from lightllm.utils.graceful_utils import graceful_registry from lightllm.utils.process_check import start_parent_check_thread from lightllm.utils.envs_utils import get_unique_server_name @@ -67,7 +67,7 @@ def __init__(self, args: StartArgs): self.shm_max_total_token_num = SharedInt("shm_max_total_token_num") self.shm_req_manager = ShmReqManager() # 用共享内存进行共享，router 模块读取进行精确的调度估计 - self.read_only_statics_mem_manager = ReadOnlyStaticsMemoryManager() + self.kv_capacity_reader = KVCapacityReader() # 初始化 radix_cache_client 用于读取 prompt cache 的管理信息 self.radix_cache_client = None @@ -235,8 +235,7 @@ async def loop_for_fwd( for dp_index in range(self.dp_size_in_node): token_ratio1 = self.get_used_tokens(dp_index) / self.max_total_token_num token_ratio2 = ( - self.max_total_token_num - - self.read_only_statics_mem_manager.get_unrefed_token_num(dp_index) + self.max_total_token_num - self.kv_capacity_reader.get_available_tokens(dp_index) ) / self.max_total_token_num d_i = dp_index estimated_peak_token_count = self.shared_token_load.get_estimated_peak_token_count(d_i) @@ -387,11 +386,11 @@ def get_used_tokens(self, dp_index): if not self.args.disable_dynamic_prompt_cache: return ( self.max_total_token_num - - self.read_only_statics_mem_manager.get_unrefed_token_num(dp_index) + - self.kv_capacity_reader.get_available_tokens(dp_index) - self.radix_cache_client.get_unrefed_tokens_num(dp_index) ) else: - return self.max_total_token_num - self.read_only_statics_mem_manager.get_unrefed_token_num(dp_index) + return self.max_total_token_num - self.kv_capacity_reader.get_available_tokens(dp_index) def _add_req(self, group_req_indexes: GroupReqIndexes): req_group = [] diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index e135be36b6..056c36f596 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -1,6 +1,5 @@ import enum import torch -import torch.distributed as dist import numpy as np import collections import pickle @@ -8,7 +7,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.attention_state import AttentionStateManager 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 +39,7 @@ class InferenceContext: backend: Optional["ModeBackend"] = None req_manager: ReqManager = None # gpu 请求管理 + att_state_manager: AttentionStateManager = None radix_cache: Union[HybridAttPagedRadixCache, RadixCache] = None shm_req_manager: ShmReqManager = None # 共享内存请求对象管理 requests_mapping: Dict[int, "InferReq"] = None @@ -64,6 +65,7 @@ def register( self.args = get_env_start_args() self.backend: ModeBackend = backend self.req_manager = req_manager + self.att_state_manager = backend.model.att_state_manager 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 +76,7 @@ def register( self.vocab_size = vocab_size - self.is_hybrid_att_model = isinstance(self.req_manager, HybridAttentionReqManager) + self.is_hybrid_att_model = backend.is_hybrid_att_model self.output_constraint_state_factory = self._init_output_constraint_state_factory() return @@ -174,11 +176,11 @@ def _free_req_mem_without_radix_insert(self, free_token_index: List, req: "Infer 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]) + self.att_state_manager.bounded_att_small_page_allocator.free([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())) + self.att_state_manager.bounded_att_big_page_allocator.free(list(req.hybrid_len_to_big_page_id.values())) req.hybrid_len_to_big_page_id.clear() # 解除请求对已命中 GPU radix cache 前缀节点的引用。 @@ -292,7 +294,7 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): # 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())) + self.att_state_manager.bounded_att_big_page_allocator.free(list(req.hybrid_len_to_big_page_id.values())) req.hybrid_len_to_big_page_id.clear() req.cur_kv_len = shared_kv_len @@ -305,20 +307,6 @@ def _hybrid_att_free_req(self, free_token_index: List, req: "InferReq"): assert False, f"error state: cur_kv_len: {req.cur_kv_len}" - def _save_promptcache_kvbuffer(self): - """ - save prompt cache kv buffer - 这个接口是用于保存非量化的缓存prompt cache资源，是定制场景使用的接口，当前代码中不会有调用。 - 其保存的 kv 会配合量化推理模式, 加载到量化推理的prompt cache中, 提升量化推理的精度。 - like paper: - https://arxiv.org/abs/2403.01241 - """ - prompt_cache_token_id = list(self.radix_cache.root_node.children.values())[0].token_id_key - print(f"prompt_cache_token_id : {prompt_cache_token_id}") - index = range(len(prompt_cache_token_id)) - prompt_cache_kv_buffer = self.radix_cache.mem_manager.get_index_kv_buffer(index) - torch.save(prompt_cache_kv_buffer, f"prompt_cache_rank_{dist.get_rank()}.pt") - @torch.no_grad() def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool = True): if len(finished_request_ids) == 0: @@ -348,7 +336,8 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool if free_token_index: free_token_index = custom_cat(free_token_index) - self.req_manager.free(free_req_index, free_token_index) + self.att_state_manager.full_att_allocator.free(free_token_index) + self.req_manager.free(free_req_index) finished_req_ids_set = set(finished_request_ids) self.infer_req_ids = [_id for _id in self.infer_req_ids if _id not in finished_req_ids_set] @@ -358,8 +347,8 @@ def _filter(self, finished_request_ids: List[int], modify_shm_finish_state: bool f"free a batch state:\n" f"radix refed token num {self.radix_cache.get_refed_tokens_num()}\n" f"radix hold token num {self.radix_cache.get_tree_total_tokens_num()}\n" - f"mem manager can alloc token num {self.req_manager.mem_manager.allocator.can_use_mem_size}\n" - f"mem manager total size {self.req_manager.mem_manager.allocator.size}\n" + f"KV allocator available tokens {self.att_state_manager.full_att_allocator.can_use_mem_size}\n" + f"KV allocator total tokens {self.att_state_manager.full_att_allocator.size}\n" ) return @@ -388,7 +377,7 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool): if len(free_token_index) != 0: free_token_index = custom_cat(free_token_index) - self.req_manager.free_token(free_token_index) + self.att_state_manager.full_att_allocator.free(free_token_index) return self def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool): @@ -422,7 +411,7 @@ def get_can_alloc_token_num(self): radix_cache_unref_token_num = ( self.radix_cache.get_tree_total_tokens_num() - self.radix_cache.get_refed_tokens_num() ) - return self.req_manager.mem_manager.allocator.can_use_mem_size + radix_cache_unref_token_num + return self.att_state_manager.full_att_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.""" @@ -441,8 +430,9 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR # 完整图片可能跨过多个大页结束点，例如页长 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 + slots = self.att_state_manager.bounded_att_big_page_allocator.alloc(1) + assert slots is not None + big_page_id = slots[0] 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，匹配时也不能在此恢复。 @@ -452,10 +442,12 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR 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, - reqs=reqs, - buffer_indexes=big_page_buffer_ids, + selected = [(req, slot) for req, slot in zip(reqs, big_page_buffer_ids) if slot != -1] + self.att_state_manager.save_bounded_att_runtime_to_pages( + req_indices=[req.req_idx for req, _ in selected], + state_token_lens=[req.get_chuncked_input_token_len() for req, _ in selected], + page_kind="big", + page_slot_ids=[slot for _, slot in selected], ) assert not self.args.disable_chunked_prefill, "chunked prefill must be enabled for hybrid attention models" @@ -467,13 +459,16 @@ def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferR 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() + allocator = self.att_state_manager.bounded_att_small_page_allocator + slots = allocator.alloc(1) if allocator is not None else None + req.tail_small_page_buffer_id = slots[0] if slots else None 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=req, - buffer_idx=dst_buffer_idx, - state_cache_manager=self.radix_cache.small_page_buffers, + self.att_state_manager.save_bounded_att_runtime_to_pages( + req_indices=[req.req_idx], + state_token_lens=[req.hybrid_cache_len], + page_kind="small", + page_slot_ids=[dst_buffer_idx], ) return @@ -638,7 +633,7 @@ def _init_all_state(self): self.stop_sequences = self.sampling_param.shm_param.stop_sequences.to_token_ids_list() self.multimodal_params = self.multimodal_params.to_dict() self.mm_input_spans = ( - get_mm_input_spans(self.get_input_token_ids(), self.multimodal_params) + get_mm_input_spans(self.shm_req.shm_prompt_ids.arr[: self.shm_req.input_len], self.multimodal_params) if self.args.disable_chunked_mm_input else [] ) @@ -679,6 +674,16 @@ def _match_radix_cache(self): self.shm_req.shm_cur_kv_len = self.cur_kv_len return + def _restore_attention_page(self, node, page_kind): + resources = g_infer_context.att_state_manager + slot = node.big_page_buffer_idx if page_kind == "big" else node.small_page_buffer_idx + resources.restore_bounded_att_runtime_from_pages( + req_indices=[self.req_idx], + state_token_lens=[node.node_prefix_total_len], + page_kind=page_kind, + page_slot_ids=[slot], + ) + def _hybrid_match_radix_cache(self): assert ( g_infer_context.is_hybrid_att_model is True @@ -713,9 +718,7 @@ def _hybrid_match_radix_cache(self): 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._restore_attention_page(share_node, "big") else: # 小页匹配 if big_page_is_disable: @@ -728,9 +731,7 @@ def _hybrid_match_radix_cache(self): 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, - ) + self._restore_attention_page(share_node, "small") else: # 如果 大页本质是被启用的，则需要使用小页的匹配结果, 将小页的kv 复制到的新申请的kv位置，同时释放 # 对应的小页对应的节点，递归找到对应最近的大叶节点进行返回,然后赋值到req.shared_node 对象上 @@ -742,7 +743,7 @@ def _hybrid_match_radix_cache(self): # 有充足的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) + tail_mems = g_infer_context.att_state_manager.full_att_allocator.alloc(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] @@ -751,8 +752,9 @@ def _hybrid_match_radix_cache(self): ] = 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 + g_infer_context.att_state_manager.copy_full_att_tokens( + value_tensor[cur_big_page_tokens:shared_kv_len].to(device="cuda", non_blocking=True), + tail_mems.to(device="cuda", non_blocking=True), ) # 尾部 KV 换到新 mem 后，同步拷贝已捕获的 top-k prompt logprobs。 self.prompt_selected_logprobs.copy_capture_slots_if_needed( @@ -760,11 +762,7 @@ def _hybrid_match_radix_cache(self): 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 + self._restore_attention_page(share_node, "small") big_page_shared_node = radix_cache.deref_to_first_big_page_node(node=share_node) self.shared_kv_node = big_page_shared_node @@ -788,9 +786,7 @@ def _hybrid_match_radix_cache(self): 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._restore_attention_page(share_node, "big") self.shm_req.shm_cur_kv_len = self.cur_kv_len @@ -802,7 +798,7 @@ def _hybrid_match_radix_cache(self): if self.cur_kv_len == 0: # 说明没有任何命中 - g_infer_context.req_manager.init_hybrid_attention_state(req=self) + g_infer_context.att_state_manager.reset_bounded_att_state([self.req_idx]) return def is_master_req(self): diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 6c621b6adb..7997f49934 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -1,5 +1,3 @@ -import os - import numpy as np import torch import time @@ -15,7 +13,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,40 +144,61 @@ def init_model(self, kvargs): } self.model, self.is_multimodal = get_model(model_cfg, model_kvargs) self.model: TpPartBaseModel = self.model # for easy typing - set_random_seed(2147483647) - self.is_hybrid_att_model = isinstance(self.model.req_manager, HybridAttentionReqManager) - - if self.is_hybrid_att_model: - self.small_page_buffers = self.model.req_manager.create_small_page_cache_manager( - size=self.args.linear_att_cache_size, + self.draft_models = [] + if self.args.mtp_mode is not None: + self.init_mtp_draft_model(model_kvargs) + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.req_manager import ReqManager + + models = (self.model, *self.draft_models) + assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" + assert all( + item.max_req_num == self.model.max_req_num for item in models + ), "models must share the request domain" + # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 + max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) + with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + req_manager = ReqManager(max_request_num=self.model.max_req_num, max_sequence_length=max_seq_len) + for item in models: + item.bind_req_manager(req_manager=req_manager) + for item in models: + item.prepare_attention_states() + manager: AttentionStateManager = self.model.att_state_manager + assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" + with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.finish_state_registration() + manager.load_full_att_calibration(args=self.model.args) + manager.initialize_storage( + self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction ) - else: - self.small_page_buffers = None + for item in models: + item.init_runtime() + torch.cuda.empty_cache() + set_model_init_status(True) + set_random_seed(2147483647) + self.is_hybrid_att_model = bool(self.model.att_state_manager.get_bounded_attention_states()) if not self.use_dynamic_prompt_cache: self.radix_cache = None else: if self.is_hybrid_att_model: self.radix_cache = HybridAttPagedRadixCache( - total_token_num=self.model.mem_manager.size, + total_token_num=self.model.att_state_manager.full_att_allocator.size, rank_in_node=self.rank_in_node, hash_page_size=self.args.linear_att_hash_page_size, big_page_num=self.args.linear_att_page_block_num, - kv_cache_mem_manager=self.model.mem_manager, - small_page_buffers=self.small_page_buffers, + att_state_manager=self.model.att_state_manager, ) else: self.radix_cache = RadixCache( - total_token_num=self.model.mem_manager.size, + total_token_num=self.model.att_state_manager.full_att_allocator.size, rank_in_node=self.rank_in_node, - mem_manager=self.model.mem_manager, + att_state_manager=self.model.att_state_manager, page_size=self.args.page_size, ) - if "prompt_cache_kv_buffer" in model_cfg: - assert self.use_dynamic_prompt_cache - self.preload_prompt_cache_kv_buffer(model_cfg) - self.logger.info(f"loaded model class {self.model.__class__}") cache_placement_controller = create_cache_placement_controller( @@ -219,16 +237,30 @@ 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) + from lightllm.server.router.model_infer.shared_att_resources_utils import publish_att_resources + + # PD 和 DP 使用独立导出，避免 PD 读者收到大小页池或 DP 读者丢失页池。 + if self.args.run_mode in ("prefill", "decode"): + publish_att_resources( + self.model.att_state_manager, + self.model.req_manager.req_to_token_indexs, + kind="pd", + reader_count=self.node_world_size, + ) + if self.args.enable_dp_prompt_cache_fetch: + publish_att_resources( + self.model.att_state_manager, + self.model.req_manager.req_to_token_indexs, + kind="dp", + reader_count=self.dp_size_in_node - 1, + ) dist.barrier(group=self.node_nccl_group) # 同一 DP 组内只需主 rank 初始化真实的 capture buffer 并执行后续相关操作； # 非主 rank 不需要分配 buffer，避免重复占用内存。 if self.is_master_in_dp: # Capture 只保存 allocator 管理的真实 KV 槽位，HOLD 页对应的 padding 写入由 kernel 过滤。 - kv_cache_size = self.model.mem_manager.size + kv_cache_size = self.model.att_state_manager.full_att_allocator.size if self.args.enable_prompt_logprobs: mgr = PromptLogprobsCaptureManager.get_instance() if mgr is not None: @@ -248,7 +280,6 @@ def init_model(self, kvargs): self.shm_pd_trans_io_buffer = ShmObjsIOBuffer(tail_str="pd") if self.args.mtp_mode is not None: - self.init_mtp_draft_model(model_kvargs) self.init_spec_engine() if self.args.enable_cpu_cache: @@ -274,7 +305,7 @@ 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 + from lightllm.server.router.model_infer.shared_att_resources_utils import load_att_resources torch.cuda.set_device(get_current_device_id()) @@ -284,17 +315,25 @@ 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): + # 每个 DP 只保留对应 TP 分片，统一以节点内 DP rank 索引。 + self.dp_att_state_managers = {} + self.dp_req_to_token_indices = {} + # 取节点内的 TP 起点，兼容单个 TP 组跨节点、rank_in_dp 超过本节点 rank 范围的情况。 + for dp_rank, rank_idx in enumerate( + range(self.rank_in_node % self.dp_world_size, self.node_world_size, self.dp_world_size) + ): if rank_idx != self.rank_in_node: - self.mem_managers.append(MemoryManager.loads_from_shm(rank_idx)) + att_state_manager, req_to_token_indices = load_att_resources(rank_idx, kind="dp") + att_state_manager.get_dp_shared_operator().register_page_pools_pin_memory() else: - self.mem_managers.append(self.model.mem_manager) + att_state_manager = self.model.att_state_manager + req_to_token_indices = self.model.req_manager.req_to_token_indexs + self.dp_att_state_managers[dp_rank] = att_state_manager + self.dp_req_to_token_indices[dp_rank] = req_to_token_indices return def get_max_total_token_num(self): - return self.model.mem_manager.size + return self.model.att_state_manager.full_att_allocator.size def infer_loop(self): raise NotImplementedError() @@ -311,9 +350,10 @@ def init_mtp_draft_model(self, main_kvargs: dict): spec_mode = self.args.mtp_mode is_chained_draft = spec_mode in ("vanilla_with_att", "vanilla_no_att") - os.environ["DISABLE_CHECK_MAX_LEN_INFER"] = "1" - draft_model_count = self.max_draft_step if is_chained_draft else 1 + from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS + + assert 0 < draft_model_count <= MAX_DRAFT_MODELS, "at most 36 draft models are supported" draft_model_dirs = self.args.mtp_draft_model_dir assert draft_model_dirs is not None assert len(draft_model_dirs) >= draft_model_count @@ -322,7 +362,7 @@ def init_mtp_draft_model(self, main_kvargs: dict): draft_model_cfg, _ = PretrainedConfig.get_config_dict(draft_model_dirs[i]) draft_model_kvargs = { "weight_dir": draft_model_dirs[i], - "max_total_token_num": self.model.mem_manager.size, + "max_total_token_num": main_kvargs["max_total_token_num"], "load_way": main_kvargs["load_way"], "max_req_num": main_kvargs.get("max_req_num", 1000), "max_seq_length": main_kvargs.get("max_seq_length", 1024 * 5), @@ -632,7 +672,7 @@ def _alloc_req_kv_mem( old_hold_kv_len = req_obj.hold_kv_len new_hold_kv_len = old_hold_kv_len + alloc_token_num - mem_indexes = g_infer_context.req_manager.mem_manager.alloc(alloc_token_num) + mem_indexes = g_infer_context.att_state_manager.full_att_allocator.alloc(alloc_token_num) # 高频调度路径允许异步写入请求索引表，其他调用方默认保持原有的同步拷贝语义。 g_infer_context.req_manager.req_to_token_indexs[req_obj.req_idx, old_hold_kv_len:new_hold_kv_len].copy_( mem_indexes, non_blocking=no_blcoking_copy @@ -1009,27 +1049,6 @@ def _dp_all_reduce_decode_req_num(self, decode_reqs: List[InferReq]) -> int: max_decode_num = self.dp_reduce_tensor.item() return max_decode_num - def preload_prompt_cache_kv_buffer(self, model_cfg): - self.logger.info("Preload prompt cache kv buffer.") - cur_rank = dist.get_rank() - prompt_cache_kv_buffer_path = os.path.join( - self.weight_dir, model_cfg["prompt_cache_kv_buffer"][f"rank_{cur_rank}"] - ) - page_size = self.args.page_size - intact_kv_len = len(model_cfg["prompt_cache_token_ids"]) // page_size * page_size - if intact_kv_len == 0: - return - - prompt_cache_kv_buffer = torch.load(prompt_cache_kv_buffer_path, weights_only=True, map_location="cpu") - prompt_cache_kv_buffer = {name: buffer[:, :intact_kv_len] for name, buffer in prompt_cache_kv_buffer.items()} - intact_kv_index = self.radix_cache.mem_manager.alloc(intact_kv_len) - self.radix_cache.mem_manager.load_index_kv_buffer(intact_kv_index, prompt_cache_kv_buffer) - intact_token_ids = torch.tensor( - model_cfg["prompt_cache_token_ids"][:intact_kv_len], dtype=torch.int64, device="cpu" - ) - self.radix_cache.insert(intact_token_ids, intact_kv_index) - self.radix_cache.match_prefix(intact_token_ids, update_refs=True) - def init_rank_infos(self): self.node_world_size = get_node_world_size() self.rank_in_node = get_current_rank_in_node() diff --git a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py index 46586959f1..57d002049a 100644 --- a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py @@ -182,7 +182,7 @@ def _master_req_to_radix_cache(self, master_req: InferReq): prefix_len, new_shared_kv_node = self.radix_cache.insert(key, value) old_prefix_len = 0 if master_req.shared_kv_node is None else master_req.shared_kv_node.node_prefix_total_len assert old_prefix_len <= master_req.cur_kv_len - self.model.mem_manager.free( + self.model.att_state_manager.full_att_allocator.free( self.model.req_manager.req_to_token_indexs[master_req.req_idx][old_prefix_len:prefix_len] ) 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..03c69be293 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py @@ -1,15 +1,12 @@ # 该文件用于提供在数据dp并行的推理模式下，共享kv cache trans相关的功能函数模块 -import time import numpy as np import dataclasses import torch from typing import List -from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_dp_rank_in_node from lightllm.server.core.objs.shm_array import ShmArray from ...infer_batch import InferReq -from lightllm.utils.dist_utils import get_current_device_id from lightllm.server.router.model_infer.infer_batch import g_infer_context import torch.distributed as dist @@ -24,28 +21,50 @@ def __init__(self, max_req_num: int, dp_size_in_node: int, backend): self.backend: DPChunkedPrefillBackend = backend self.max_req_num = max_req_num - # 0 代表 kv_len, 1 代表 radix_cache_len + self.big_page_tokens = ( + backend.args.linear_att_hash_page_size * backend.args.linear_att_page_block_num + if backend.is_hybrid_att_model + else 0 + ) + max_big_pages = backend.args.max_req_total_len // self.big_page_tokens if self.big_page_tokens else 0 + # 原共享信息只保存 KV 长度和请求槽位，大页列表使用独立共享数组。 self.shared_req_infos = ShmArray( name="dp_shared_req_infos", shape=(self.max_req_num, dp_size_in_node, 2), dtype=np.int64, ) self.shared_req_infos.create_shm() + self.shared_req_big_page_ids: ShmArray | None = None + if max_big_pages: + # 第 i 项对应前缀长度 (i + 1) * big_page_tokens，未使用位置填 -1。 + self.shared_req_big_page_ids = ShmArray( + name="dp_shared_req_big_page_ids", + shape=(self.max_req_num, dp_size_in_node, max_big_pages), + dtype=np.int64, + ) + self.shared_req_big_page_ids.create_shm() self.dp_rank_in_node = get_dp_rank_in_node() assert get_env_start_args().diverse_mode is False - def fill_reqs_info(self, reqs: List[InferReq]): - """ - 填充请求的 kv 信息到共享内存中 - """ + def fill_reqs_info(self, reqs: List[InferReq]) -> None: + """发布每个请求的 KV 长度、请求槽位及完整有序大页列表。""" dist.barrier(group=self.backend.node_nccl_group) if self.backend.is_master_in_dp: - self.shared_req_infos.arr[0 : len(reqs), self.dp_rank_in_node, self._KV_LEN_INDEX] = [ - req.cur_kv_len for req in reqs - ] - self.shared_req_infos.arr[0 : len(reqs), self.dp_rank_in_node, self._REQ_IDX_INDEX] = [ - req.req_idx for req in reqs - ] + for req_index, req in enumerate(reqs): + row = self.shared_req_infos.arr[req_index, self.dp_rank_in_node] + row[self._KV_LEN_INDEX] = req.cur_kv_len + row[self._REQ_IDX_INDEX] = req.req_idx + if self.shared_req_big_page_ids is None: + continue + page_ids = self.shared_req_big_page_ids.arr[req_index, self.dp_rank_in_node] + page_ids.fill(-1) # 批次复用共享行时，清除上次较长请求留下的页 ID。 + # radix 页按路径从根到叶排列；请求持有的页按其前缀长度放到对应位置。 + if req.shared_kv_node is not None: + shared_pages = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) + page_ids[: len(shared_pages)] = shared_pages + for token_len, page_id in req.hybrid_len_to_big_page_id.items(): + if token_len <= req.cur_kv_len: + page_ids[token_len // self.big_page_tokens - 1] = page_id return def build_shared_kv_trans_tasks( @@ -59,83 +78,103 @@ def build_shared_kv_trans_tasks( dist.barrier(group=self.backend.node_nccl_group) trans_tasks: List[TransTask] = [] - rank_max_radix_cache_lens = np.max( - self.shared_req_infos.arr[0 : len(reqs), :, self._KV_LEN_INDEX], axis=1, keepdims=False - ) - # 如果发现自己是dp_rank 最小， radix_cache_len 最长的请求，则将数据写入到共享内存中。 - for req_index, req, max_req_radix_cache_len, req_dp_rank in zip( - list(range(len(reqs))), reqs, rank_max_radix_cache_lens, req_dp_ranks + available_lens = self.shared_req_infos.arr[0 : len(reqs), :, self._KV_LEN_INDEX] + max_available_kv_lens = np.max(available_lens, axis=1) + # 为本 DP 负责的请求选择可用前缀最长的来源。 + for req_index, req, max_available_kv_len, req_dp_rank in zip( + list(range(len(reqs))), + reqs, + max_available_kv_lens, + req_dp_ranks, ): # 当前请求是本 dp_rank 负责的 - is_current_dp_handle = req_dp_rank == self.dp_rank_in_node + is_local_dp_request = req_dp_rank == self.dp_rank_in_node # 计算需要传输的 kv 长度， 不能超过 req.get_cur_total_len() - 1 - trans_size = min(max_req_radix_cache_len, req.get_cur_total_len() - 1) - req.cur_kv_len + trans_size = min(max_available_kv_len, req.get_cur_total_len() - 1) - req.cur_kv_len target_kv_len = req.cur_kv_len + trans_size alloc_token_num = req._kv_cache_alloc_need(target_kv_len) if trans_size > 0 else 0 - if is_current_dp_handle and trans_size > 0 and alloc_token_num <= g_infer_context.get_can_alloc_token_num(): + if is_local_dp_request and trans_size > 0 and alloc_token_num <= g_infer_context.get_can_alloc_token_num(): assert req.hold_kv_len == req.cur_kv_len - mem_indexes = self.backend._alloc_req_kv_mem(req, alloc_token_num) - assert mem_indexes is not None - # mem_indexes 只描述需要复制的逻辑 KV；页尾预留槽位已经由 - # _alloc_req_kv_mem 写入请求表，但不参与本次跨 DP 传输。 - mem_indexes = mem_indexes[:trans_size] - 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 + src_dp_rank = int(available_lens[req_index].argmax()) + src_info = self.shared_req_infos.arr[req_index, src_dp_rank] + src_req_idx = int(src_info[self._REQ_IDX_INDEX]) + src_big_page_ids, dst_big_page_ids = [], [] + if self.backend.is_hybrid_att_model: + # 只复制新增前缀跨过的完整大页，已有前缀页继续使用本地引用。 + first_page = req.cur_kv_len // self.big_page_tokens + end_page = target_kv_len // self.big_page_tokens + if end_page > first_page: + src_big_page_ids = self.shared_req_big_page_ids.arr[ + req_index, src_dp_rank, first_page:end_page + ].tolist() + assert all( + page_id >= 0 for page_id in src_big_page_ids + ), "DP source is missing prefix big pages" + dst_token_indices = self.backend._alloc_req_kv_mem(req, alloc_token_num) + assert dst_token_indices is not None + if src_big_page_ids: + allocator = self.backend.model.att_state_manager.bounded_att_big_page_allocator + dst_big_page_ids = allocator.alloc(len(src_big_page_ids)) + assert dst_big_page_ids is not None, "DP fetch big page allocation must succeed" + # 只搬运有效 KV；页尾预留槽位已写入请求表，不参与本次传输。 + dst_token_indices = dst_token_indices[:trans_size] + src_token_indices = self.backend.dp_req_to_token_indices[src_dp_rank][ + src_req_idx, req.cur_kv_len : target_kv_len ] trans_tasks.append( TransTask( - req=req, - mem_indexes=mem_indexes, - 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, + dst_req=req, + src_dp_rank=src_dp_rank, + src_req_idx=src_req_idx, + src_token_indices=src_token_indices, + dst_token_indices=dst_token_indices, + src_big_page_ids=src_big_page_ids, + dst_big_page_ids=dst_big_page_ids, ) ) return trans_tasks - def kv_trans(self, trans_tasks: List["TransTask"]): - # kv 传输 - if len(trans_tasks) > 0: - max_kv_len_mem_indexes = [] - max_kv_len_dp_ranks = [] - mem_indexes = [] - - for i, trans_task in enumerate(trans_tasks): - max_kv_len_mem_indexes.append(trans_task.max_kv_len_mem_indexes) - max_kv_len_dp_ranks.extend([trans_task.max_kv_len_dp_rank] * len(trans_task.max_kv_len_mem_indexes)) - mem_indexes.append(trans_task.mem_indexes) - - max_kv_len_mem_indexes_tensor = torch.cat(max_kv_len_mem_indexes).to(dtype=torch.int64, device="cuda") - 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, - move_token_indexes=max_kv_len_mem_indexes_tensor, - token_dp_indexes=max_kv_len_dp_ranks_tensor, - mem_indexes=mem_indexes_tensor, - dp_size_in_node=self.backend.dp_size_in_node, - rank_in_dp=self.backend.rank_in_dp, + def kv_trans(self, trans_tasks: List["TransTask"]) -> None: + target = self.backend.model.att_state_manager + for task in trans_tasks: + source = self.backend.dp_att_state_managers[task.src_dp_rank] + # 源索引来自远端 CUDA 请求表；allocator 返回 CPU 索引，在此一次性准备目标 CUDA int64 索引。 + dst_token_indices = task.dst_token_indices.to(device="cuda", dtype=torch.long, non_blocking=True) + target.get_dp_shared_operator().copy_from_dp( + remote_att_manager=source, + src_token_indices=task.src_token_indices, + dst_token_indices=dst_token_indices, + src_big_page_ids=task.src_big_page_ids, + dst_big_page_ids=task.dst_big_page_ids, + src_req_idx=task.src_req_idx, + dst_req_idx=task.dst_req.req_idx, ) - # This copy runs during request initialization, not once per PD transfer page. - self.backend.logger.debug(f"dp_i {self.dp_rank_in_node} transfer kv tokens num: {len(mem_indexes_tensor)}") + # Source requests may be released on other ranks immediately afterwards. + # All ranks participate, including those without a local transfer task. + dist.barrier(group=self.backend.node_nccl_group) - for trans_task in trans_tasks: - trans_task.req.cur_kv_len += len(trans_task.mem_indexes) - assert trans_task.req.cur_kv_len <= trans_task.req.hold_kv_len + for task in trans_tasks: + req = task.dst_req + # 复制完成后，才把目标大页发布为请求持有的快照；后续 radix 插入或释放沿用此表。 + if task.dst_big_page_ids: + first_page_end = (req.cur_kv_len // self.big_page_tokens + 1) * self.big_page_tokens + for page_offset, page_id in enumerate(task.dst_big_page_ids): + token_len = first_page_end + page_offset * self.big_page_tokens + req.hybrid_len_to_big_page_id[token_len] = page_id + req.cur_kv_len += len(task.dst_token_indices) + assert req.cur_kv_len <= 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 + req.shm_req.shm_cur_kv_len = req.cur_kv_len @dataclasses.dataclass class TransTask: - req: InferReq - mem_indexes: torch.Tensor - max_kv_len_dp_rank: int - max_kv_len_mem_manager_index: int - max_kv_len_mem_indexes: torch.Tensor + dst_req: InferReq # 本 DP 接收前缀的请求，持有目标请求槽位和新分配的资源。 + src_dp_rank: int # 来源 DP 在当前节点内的编号。 + src_req_idx: int # 远端请求槽位，用于复制 prefill 基准运行态。 + src_token_indices: torch.Tensor # 来源 CUDA 请求表中的有效 KV 索引。 + dst_token_indices: torch.Tensor # allocator 返回的目标 CPU 索引，提交时统一转为 CUDA int64。 + src_big_page_ids: List[int] # 按前缀长度递增排列的源大页槽位。 + dst_big_page_ids: List[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 4b72c2ec4f..ed7e5bb167 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 @@ -98,8 +98,10 @@ def _init_reqs(self, reqs: List[Tuple]): infer_reqs = g_infer_context.add_reqs(reqs, init_prefix_cache=True) req_dp_ranks = [req[3] for req in reqs] self.dp_kv_shared_module.fill_reqs_info(reqs=infer_reqs) - trans_taskes = self.dp_kv_shared_module.build_shared_kv_trans_tasks(reqs=infer_reqs, req_dp_ranks=req_dp_ranks) - self.dp_kv_shared_module.kv_trans(trans_tasks=trans_taskes) + transfer_tasks = self.dp_kv_shared_module.build_shared_kv_trans_tasks( + reqs=infer_reqs, req_dp_ranks=req_dp_ranks + ) + self.dp_kv_shared_module.kv_trans(trans_tasks=transfer_tasks) # other_dp_reqs 只是为本 DP 做完 prefix cache / KV 拉取后的临时本地对象， # 真正推理仍在其归属 DP 上进行。这里仅清理本 DP 的 InferReq 与 KV 引用， 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..261f090b0b 100644 --- a/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py +++ b/lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py @@ -1,4 +1,3 @@ -import threading import torch.distributed as dist import torch import dataclasses @@ -8,11 +7,9 @@ from collections import deque from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient -from lightllm.utils.config_utils import is_hybrid_att_model from lightllm.utils.envs_utils import get_env_start_args from ..infer_batch import InferReq from lightllm.utils.dist_utils import create_new_group_for_current_dp -from lightllm.common.basemodel.triton_kernel.kv_cache_offload import offload_gpu_kv_to_cpu, load_cpu_kv_to_gpu from lightllm.server.router.model_infer.infer_batch import g_infer_context from lightllm.utils.log_utils import init_logger from lightllm.utils.multimodal_chunk_utils import clamp_mm_cache_len @@ -35,11 +32,13 @@ def __init__(self, backend): dist.barrier(group=self.offload_sync_group) self.offload_sync_tensor = torch.empty((1,), dtype=torch.int32, device="cuda") - self.page_index_buffer = torch.empty((1024 * 1024 * 4,), dtype=torch.int32, device="cuda") - self.page_ready_buffer = torch.empty((1024 * 1024 * 4,), dtype=torch.bool, device="cuda") - self.cpu_cache_handle_queue: Deque[TransTask] = deque() self.cpu_cache_client = CpuKvCacheClient(only_create_meta_data=False, init_shm_data=False) + resources = self.backend.model.att_state_manager + self.cpu_cache_operator = resources.get_cpu_cache_operator() + meta = self.cpu_cache_client.kv_cache_tensor_meta + if meta.page_bytes != self.cpu_cache_operator.get_page_bytes(): + raise ValueError("CPU cache service and model resource page sizes differ") @lru_cache() def need_sync_compute_stream(self) -> bool: @@ -77,7 +76,6 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): continue page_len_list = req.shm_req.token_hash_page_len_list.get_all() - page_len_start_list = [0] + page_len_list assert len(page_list) <= len(page_len_list) matched_page_num = len(page_list) @@ -113,43 +111,13 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): assert ready_page_num <= matched_page_num need_pages = page_list[ready_page_num:matched_page_num] - mem_indexes = self.backend._alloc_req_kv_mem(req, need_token_num) - assert mem_indexes is not None - - if self.need_sync_compute_stream(): - # TODO fa3 现在必须使用同步模式, 未来需要移除 - 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) - page_indexes_cuda = torch.tensor(need_pages, dtype=torch.int32, device="cpu").cuda( - non_blocking=True - ) - # hybrid 页面加载必须按完整 page 处理，否则可能缺失恢复运行态所需的 checkpoint，所以 - # 这里需要进行pad操作，使操作的页面是完整的。 - _start = page_len_start_list[ready_page_num] - - _end = req.cur_kv_len - assert 0 <= _start <= _end, f"invalid pad range [{_start}, {_end}]" - mem_indexes_cuda = torch.cat( - [req_manager.req_to_token_indexs[req.req_idx, _start:_end], mem_indexes_cuda] - ) - - assert len(mem_indexes_cuda) == match_tokens - page_len_start_list[ready_page_num] - - # 更新 req 状态。 - 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( - mem_indexes=mem_indexes_cuda, - page_indexes=page_indexes_cuda, - cpu_cache_client=self.cpu_cache_client, + self.read_req_from_cpu( req=req, + cpu_page_indices=need_pages, + page_lengths=page_len_list[ready_page_num:matched_page_num], + start_token_len=page_len_list[ready_page_num - 1] if ready_page_num else 0, ) + idle_token_num -= need_token_num torch.cuda.current_stream().synchronize() @@ -166,6 +134,66 @@ def load_cpu_cache_to_reqs(self, reqs: List[InferReq]): self.cpu_cache_client.lock.release() return + def read_req_from_cpu( + self, + req: InferReq, + cpu_page_indices: List[int], + page_lengths: List[int], + start_token_len: int, + ) -> None: + """恢复已选定的请求前缀页面，分配资源并登记完整大页引用。 + + start_token_len 是首个 CPU 页的逻辑起点，可能早于已有 GPU 命中长度。 + 调用方负责选择命中范围、确认容量以及搬运完成后的 CPU 页引用释放。 + """ + if not cpu_page_indices: + return + resources = self.backend.model.att_state_manager + cpu_cache_operator = resources.get_cpu_cache_operator() + req_manager = self.backend.model.req_manager + need_token_num = page_lengths[-1] - req.cur_kv_len + + # 阶段一：分配新增 KV，并补齐首个 CPU 页在 GPU 上已经命中的 token 索引。 + mem_indexes = self.backend._alloc_req_kv_mem(req, need_token_num) + assert mem_indexes is not None + if self.need_sync_compute_stream(): + # FA3 的资源搬运需等待已有计算流完成。 + torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream()) + assert 0 <= start_token_len <= req.cur_kv_len, "invalid CPU page start" + token_indices = torch.cat( + [ + req_manager.req_to_token_indexs[req.req_idx, start_token_len : req.cur_kv_len], + mem_indexes.cuda(non_blocking=True), + ] + ) + assert token_indices.numel() == page_lengths[-1] - start_token_len + + # 阶段二：只为完整 bounded 页分配普通大页，碎页 LOAD 槽位由资源层管理。 + big_page_slot_ids = None + if self.backend.is_hybrid_att_model: + page_tokens = self.args.cpu_cache_token_page_size + complete_page_lengths = [length for length in page_lengths if length % page_tokens == 0] + assert len(complete_page_lengths) == len(page_lengths) - int( + page_lengths[-1] % page_tokens != 0 + ), "only the final CPU page may be partial" + allocator = resources.bounded_att_big_page_allocator + big_page_slot_ids = allocator.alloc(len(complete_page_lengths)) + assert big_page_slot_ids is not None, "not enough bounded state slots for CPU cache load" + + # 阶段三：统一恢复各类 State；异常直接向上传播，不在此处回收或继续执行。 + cpu_cache_operator.read_req( + cpu_page_buffers=self.cpu_cache_client.cpu_kv_cache_tensor, + cpu_page_indices=cpu_page_indices, + token_indices=token_indices, + layouts=cpu_cache_operator.get_page_layouts(), + big_page_slot_ids=big_page_slot_ids, + req_idx=req.req_idx if self.backend.is_hybrid_att_model else None, + state_token_len=page_lengths[-1], + ) + if big_page_slot_ids is not None: + req.hybrid_len_to_big_page_id.update(zip(complete_page_lengths, big_page_slot_ids)) + req.cur_kv_len = page_lengths[-1] + def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> List[InferReq]: """ 将满足cpu kv cache 卸载条件的请求进行处理, 并返回真的满足退出条件的请求list。 @@ -208,7 +236,7 @@ def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> L g_infer_context.get_overlap_stream().synchronize() # 发起将请求的 kv cache 卸载到 cpu cache 中的任务 - trans_task = self._start_kv_cache_offload_task(req=req, cpu_kv_cache_stream=cpu_stream) + trans_task = self.write_req_to_cpu(req=req, cpu_kv_cache_stream=cpu_stream) # 根据是否成功创建了卸载任务，决定是否将请求加入到处理队列中 if trans_task is not None: @@ -222,17 +250,18 @@ def offload_finished_reqs_to_cpu_cache(self, finished_reqs: List[InferReq]) -> L return true_finished_reqs - def _start_kv_cache_offload_task( - self, req: InferReq, cpu_kv_cache_stream: torch.cuda.Stream - ) -> Optional["TransTask"]: + def write_req_to_cpu(self, req: InferReq, cpu_kv_cache_stream: torch.cuda.Stream) -> Optional["TransTask"]: + """选择请求可缓存的前缀，分配 CPU 页并提交写出，返回待确认完成的任务。""" assert CacheTier.CPU in req.cache_tiers disk_offload_enable = CacheTier.DISK in req.cache_tiers + # KV 和 bounded 快照均在搬运流上复制，由后续完成事件统一确认就绪。 with torch.cuda.stream(cpu_kv_cache_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) + # 阶段一：确定本次可写出范围，申请 CPU 页。 if self.backend.is_master_in_dp: find_index = bisect.bisect_right(page_len_list, req.cur_kv_len) @@ -257,6 +286,12 @@ def _start_kv_cache_offload_task( finally: self.cpu_cache_client.lock.release() + # 分配不足时返回 -1 并补齐剩余位置，只保留成功分配的连续前缀。 + # 第一个 -1 及后续页面（包括可能的尾部碎页）均不参与本次写出。 + if -1 in page_list: + allocated_page_count = page_list.index(-1) + page_list = page_list[:allocated_page_count] + ready_list = ready_list[:allocated_page_count] item_size = len(page_list) if item_size == 0: dist.broadcast_object_list([0], group=self.gloo_group, group_src=0) @@ -276,26 +311,42 @@ def _start_kv_cache_offload_task( page_list = broadcast_data["page_list"] ready_list = broadcast_data["ready_list"] - page_indexes = torch.tensor(page_list, dtype=torch.int32, device="cpu", pin_memory=True) - page_readies = torch.tensor(ready_list, dtype=torch.bool, device="cpu", pin_memory=True) - assert len(page_indexes) <= self.page_index_buffer.shape[0] - cuda_page_indexes = self.page_index_buffer[: len(page_indexes)] - cuda_page_readies = self.page_ready_buffer[: len(page_readies)] - cuda_page_indexes.copy_(page_indexes, non_blocking=True) - cuda_page_readies.copy_(page_readies, non_blocking=True) + page_indexes = torch.tensor(page_list, dtype=torch.int32) + page_readies = torch.tensor(ready_list, dtype=torch.bool) move_token_num = page_len_list[item_size - 1] assert req.cur_kv_len >= move_token_num token_indexes = self.backend.model.req_manager.req_to_token_indexs[req.req_idx, 0:move_token_num] - mem_manager = self.backend.model.mem_manager - - mem_manager.operator.offload_gpu_kv_to_cpu_cache( - mem_indexes=token_indexes, - page_indexes=cuda_page_indexes, - page_readies=cuda_page_readies, - cpu_cache_client=self.cpu_cache_client, - req=req, + # 阶段二：只查找实际分配的 CPU 页所需快照；完整大页必须存在，不做缺失回退。 + resources = self.backend.model.att_state_manager + big_page_slot_ids, small_page_slot_ids = None, None + if self.backend.is_hybrid_att_model: + page_tokens = self.args.cpu_cache_token_page_size + radix_cache = g_infer_context.radix_cache + shared_slots = ( + radix_cache.get_big_page_ids_by_node(req.shared_kv_node) if radix_cache is not None else [] + ) + big_slots = {(i + 1) * page_tokens: slot for i, slot in enumerate(shared_slots)} + big_slots.update(req.hybrid_len_to_big_page_id) + big_page_slot_ids = [ + big_slots[length] for length in page_len_list[:item_size] if length % page_tokens == 0 + ] + small_page_slot_ids = [] + # 只有实际分配的前缀包含原始尾页，才可能使用小页快照。 + # 若分配遇到 -1，item_size 已截短，尾页不会进入此分支。 + includes_tail_page = item_size == len(page_len_list) + if includes_tail_page and page_len_list[-1] % page_tokens != 0: + small_page_slot_ids.append(req.tail_small_page_buffer_id) + # 整批提交并保留 token/page 对齐，由资源层跳过 ready 页。 + resources.get_cpu_cache_operator().write_req( + token_indices=token_indexes, + cpu_page_buffers=self.cpu_cache_client.cpu_kv_cache_tensor, + cpu_page_indices=page_list, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + cpu_page_ready_flags=ready_list, + big_page_slot_ids=big_page_slot_ids, + small_page_slot_ids=small_page_slot_ids, ) # 这个操作只是为了在offload 对应的cuda stream中，同步标记下对应的kv cache offload 操作已经完成， diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py index 786c60b30a..6f1e90590e 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_impl.py @@ -111,7 +111,7 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: .detach() .cpu() ) - self.model.mem_manager.free(mem_indexes) + self.model.att_state_manager.full_att_allocator.free(mem_indexes) req_obj.cur_kv_len = old_prefix_len req_obj.hold_kv_len = old_prefix_len if self.is_master_in_dp: 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..ca4b51e9b5 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 @@ -8,7 +8,9 @@ import pickle from typing import List, Dict, Union, Optional from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.attention_state import AttentionStateManager +from lightllm.server.router.model_infer.shared_att_resources_utils import load_att_resources +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.server.pd_io_struct import ( PDChunckedTransTask, PDChunckedTransTaskGroup, @@ -83,20 +85,20 @@ def _init_env( task_out_queue.put("proc_start") - # 从共享内存读取所有rank的mem_manager + # 从共享内存读取所有 rank 的 attention 资源 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) + resource_managers: List[AttentionStateManager] = [ + load_att_resources(rank, kind="pd")[0] for rank in range(node_world_size) ] - task_out_queue.put("get_mem_managers_ok") + task_out_queue.put("get_resources_ok") manager = _DecodeTransModule( args=args, device_id=device_id, task_in_queue=task_in_queue, task_out_queue=task_out_queue, - mem_managers=mem_managers, + resource_managers=resource_managers, up_status_in_queue=up_status_in_queue, ) assert manager is not None @@ -119,7 +121,7 @@ def __init__( device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, - mem_managers: List[MemoryManager], + resource_managers: List[AttentionStateManager], up_status_in_queue: Optional[mp.SimpleQueue], ): self.args = args @@ -127,12 +129,12 @@ 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.resource_managers = resource_managers 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( - page_num=self.args.pd_kv_page_num, page_size=self.args.pd_kv_page_size + kv_move_buffer = allocate_transfer_buffer( + self.resource_managers[device_id], self.args.pd_kv_page_num, device=f"cuda:{device_id}" ) + self.kv_move_buffer = kv_move_buffer self.copy_cuda_stream = torch.cuda.Stream(priority=-1) self.transporter = create_kv_transporter( args=self.args, @@ -170,13 +172,16 @@ 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.read_page_kv_move_buffer_to_mem( - mem_indexes=[0], - page_index=0, - dp_index=dp_index, - mem_managers=self.mem_managers, - dp_world_size=self.dp_world_size, + cpu_token_indices = torch.tensor([0], dtype=torch.long, device="cpu") + token_indices = cpu_token_indices.cuda(device=self.kv_move_buffer.device, non_blocking=True) + copy_page( + node_att_managers=self.resource_managers, + pd_page_buffers=self.kv_move_buffer, + write_to_pd_page=False, + token_indices=token_indices, + pd_page_index=0, + dp_rank_in_node=dp_index, + tp_world_size=self.dp_world_size, ) torch.cuda.current_stream().synchronize() return @@ -392,19 +397,43 @@ def read_page_to_mems_loop(self): torch.cuda.set_device(self.device_id) while True: trans_task: PDChunckedTransTask = self.ready_page_task_queue.get() + if trans_task.xfer_handle is not None: + # 只有 NCCL 接收端会创建本地 Xfer 句柄；NIXL 的远端 WRITE 完成通知 + # 已代表其写入完成，不经过这里的本地 CUDA event 轮询。 + # P 的 done 不代表 D 本地接收已完成；统一轮询传输状态， + # 直到接收 event 完成后，独立搬运流才读取 PD 页。 + while True: + status = self.transporter.check_task_status(trans_task=trans_task) + if status == "PROC": + time.sleep(0.001) + continue + if status != "DONE": + trans_task.error_info = trans_task.xfer_handle.error_info + self.failed_queue.put(trans_task) + break + if status != "DONE": + continue copy_start_event = torch.cuda.Event(enable_timing=True) 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, + # 只有 KV 页使用 token 索引；att_state 直接按请求槽位搬运运行态。 + token_indices = None + if trans_task.page_kind == "kv": + cpu_token_indices = torch.tensor(trans_task.mem_indexes, dtype=torch.long, device="cpu") + # 索引上传和后续页拷贝使用同一流，保证执行顺序。 + token_indices = cpu_token_indices.cuda(device=self.kv_move_buffer.device, non_blocking=True) + copy_page( + node_att_managers=self.resource_managers, + pd_page_buffers=self.kv_move_buffer, + write_to_pd_page=False, + token_indices=token_indices, + pd_page_index=trans_task.dst_page_index, + dp_rank_in_node=trans_task.decode_dp_index, + tp_world_size=self.dp_world_size, + task_kind=trans_task.page_kind, req_idx=trans_task.req_idx, + state_token_len=trans_task.end_kv_index if trans_task.page_kind == "att_state" else None, ) 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/nccl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py index d3f2c0214d..efacb98531 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py @@ -56,7 +56,7 @@ def __init__( assert args.run_mode in ["prefill", "decode"], args.run_mode self.is_prefill_node = args.run_mode == "prefill" self.capture_telemetry = False - self.num_pages, self.page_size, self.num_layers, self.kv_head_num, self.head_dims = kv_move_buffer.shape + self.num_pages = kv_move_buffer.shape[0] self.host_ip = host_ip or get_hostname_ip() assert self.host_ip is not None, "can not get host ip for NcclKVTransporter" @@ -91,11 +91,8 @@ def local_page_mem_desc(self) -> bytes: return pickle.dumps( { "num_pages": self.num_pages, - "page_size": self.page_size, - "num_layers": self.num_layers, - "kv_head_num": self.kv_head_num, - "head_dims": self.head_dims, - "dtype": str(self.kv_move_buffer.dtype), + # PD buffer 固定为二维 uint8，页长度就是单页字节数。 + "page_bytes": self.kv_move_buffer.shape[1], } ) @@ -114,6 +111,8 @@ def connect_add_remote_agent(self, remote_agent: PDAgentMetadata): metadata.agent_name == remote_agent.agent_name ), f"Peer name {metadata.agent_name} does not match remote name {remote_agent.agent_name}" + page = pickle.loads(remote_agent.page_reg_desc) + assert page["page_bytes"] == self.kv_move_buffer.shape[1], "PD peers must use the same page byte capacity" self.remote_agents[remote_agent.agent_name] = remote_agent logger.info(f"Added NCCL remote agent {remote_agent.agent_name} at {metadata.host_ip}:{metadata.control_port}") return @@ -265,16 +264,37 @@ def _get_notify_source_agent_name(self, notify: bytes) -> str: @dataclass class _NcclXferHandle: peer_name: str - event: torch.cuda.Event + mode: str + finish_cuda_event: Optional[torch.cuda.Event] = None status: str = "PROC" error_info: Optional[str] = None + def __post_init__(self) -> None: + assert self.mode in ("send", "recv"), f"invalid NCCL handle mode: {self.mode}" + def check_status(self) -> str: if self.status != "PROC": return self.status + if self.mode == "send": + return self._check_send_status() + return self._check_recv_status() + def _check_send_status(self) -> str: + if self.finish_cuda_event is None: + return "PROC" try: - if self.event.query(): + if self.finish_cuda_event.query(): + self.status = "DONE" + except BaseException as e: + self.status = "ERR" + self.error_info = str(e) + return self.status + + def _check_recv_status(self) -> str: + if self.finish_cuda_event is None: + return "PROC" + try: + if self.finish_cuda_event.query(): self.status = "DONE" except BaseException as e: self.status = "ERR" @@ -298,17 +318,22 @@ def send_page(self, trans_task: PDChunckedTransTask) -> _NcclXferHandle: stream = self._get_stream() comm.send(page_tensor, dst=1, stream=stream) - event = torch.cuda.Event() - event.record(stream) + send_finish_cuda_event = torch.cuda.Event() + send_finish_cuda_event.record(stream) if self.transporter.args.pd_detail_log: logger.debug( f"NCCL send page posted request_id={trans_task.request_id} " f"src_page={trans_task.src_page_index} dst_agent={self.peer_name}" ) - return _NcclXferHandle(peer_name=self.peer_name, event=event) + return _NcclXferHandle( + peer_name=self.peer_name, + mode="send", + finish_cuda_event=send_finish_cuda_event, + ) def start_recv(self, trans_task: PDChunckedTransTask): + trans_task.xfer_handle = _NcclXferHandle(peer_name=self.peer_name, mode="recv") self._get_recv_queue().put(copy.copy(trans_task)) return @@ -356,11 +381,17 @@ def _recv_page(self, trans_task: PDChunckedTransTask): comm = self._ensure_comm(is_server=False) stream = self._get_stream() comm.recv(page_tensor, src=0, stream=stream) + finish_cuda_event = torch.cuda.Event() + finish_cuda_event.record(stream) + # 必须在 record 返回后才发布；CUDA event 自身不会等待未来的 record。 + trans_task.xfer_handle.finish_cuda_event = finish_cuda_event if self.transporter.args.pd_detail_log: logger.debug( f"NCCL recv page done request_id={trans_task.request_id} " f"dst_page={trans_task.dst_page_index}" ) except BaseException as e: + trans_task.xfer_handle.status = "ERR" + trans_task.xfer_handle.error_info = str(e) trans_task.error_info = str(e) logger.exception(str(e)) self._drop_comm() @@ -434,6 +465,7 @@ def __init__( def _start_server(self, host_ip: str, port_min: int, port_max: int) -> tuple[ThreadedServer, int]: last_error = None for cur_port in range(port_min, port_max + 1): + server = None try: server = ThreadedServer( _NcclControlService(self), @@ -446,14 +478,23 @@ def _start_server(self, host_ip: str, port_min: int, port_max: int) -> tuple[Thr "allow_setattr": True, }, ) + # RPyC 构造时只 bind；启用 SO_REUSEADDR 后，并发进程可能都绑定同一端口成功。 + # 在当前线程先 listen，让端口冲突进入下面的重试；否则后台线程启动失败时， + # 外层已经返回该端口，无法捕获异常并切换到下一个端口。 + server.listener.listen(server.backlog) threading.Thread(target=server.start, daemon=True).start() logger.info(f"NCCL RPyC control channel listen on {host_ip}:{cur_port}") return server, cur_port except OSError as e: + # 构造失败时 server 尚未赋值；已构造但 listen 失败时，关闭 socket 后再重试。 + if server is not None: + server.close() last_error = e if e.errno == errno.EADDRINUSE: + # 当前端口已被占用，或并发启动时被其他进程抢先监听，继续尝试下一个端口。 logger.info(f"NCCL RPyC control port {host_ip}:{cur_port} is in use, try next port") else: + # 其他 socket 错误保留告警；整个端口范围都失败后，携带最后一个异常向上抛出。 logger.warning(f"Create NCCL RPyC control channel on {host_ip}:{cur_port} failed: {e}") raise RuntimeError(f"can not allocate NCCL control port in [{port_min}, {port_max}]") from last_error diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py index b80d52a747..ae4b62ff2b 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py @@ -61,17 +61,18 @@ def get_new_notifs(self) -> Dict[str, list[bytes]]: return self.nixl_agent.get_new_notifs() def _register_kv_move_buffer(self, kv_move_buffer: Tensor): - self.num_pages, self.page_size, self.num_layers, self.kv_head_num, self.head_dims = kv_move_buffer.shape - self.dtype_byte_size = kv_move_buffer.element_size() - self.page_len = self.page_size * self.num_layers * self.kv_head_num * self.head_dims * self.dtype_byte_size + self.num_pages = kv_move_buffer.shape[0] + self.page_bytes = kv_move_buffer[0].numel() * kv_move_buffer.element_size() self.page_reg_desc = self.nixl_agent.register_memory(kv_move_buffer) self.page_local_xfer_handles = self._create_paged_xfer_handles(self.page_reg_desc, self.num_pages) def _create_paged_xfer_handles(self, reg_desc: "nixlBind.nixlRegDList", page_num: int, agent_name: str = ""): - base_addr, _, device_id, _ = reg_desc[0] + base_addr, registered_bytes, device_id, _ = reg_desc[0] + # NIXL 原生描述记录整块注册内存的字节数，按对端页数检查单页容量是否一致。 + assert registered_bytes == page_num * self.page_bytes, "PD peers must use the same page byte capacity" pages_data = [] for page_id in range(page_num): - pages_data.append((base_addr + page_id * self.page_len, self.page_len, device_id)) + pages_data.append((base_addr + page_id * self.page_bytes, self.page_bytes, device_id)) descs = self.nixl_agent.get_xfer_descs(pages_data, "VRAM") return self.nixl_agent.prep_xfer_dlist(agent_name, descs, "VRAM") 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..b3c1dcc609 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 @@ -8,7 +8,9 @@ import pickle from typing import List, Dict, Optional, Union from lightllm.utils.log_utils import init_logger -from lightllm.common.kv_cache_mem_manager import MemoryManager +from lightllm.common.attention_state import AttentionStateManager +from lightllm.server.router.model_infer.shared_att_resources_utils import load_att_resources +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page from lightllm.server.pd_io_struct import PDAbortReq, PDChunckedTransTask from lightllm.utils.graceful_utils import graceful_registry from lightllm.server.core.objs import StartArgs @@ -58,20 +60,20 @@ def _init_env( graceful_registry(inspect.currentframe().f_code.co_name) task_out_queue.put("proc_start") - # 从共享内存读取所有rank的mem_manager + # 从共享内存读取所有 rank 的 attention 资源 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) + resource_managers: List[AttentionStateManager] = [ + load_att_resources(rank, kind="pd")[0] for rank in range(node_world_size) ] - task_out_queue.put("get_mem_managers_ok") + task_out_queue.put("get_resources_ok") manager = _PrefillTransModule( args=args, device_id=device_id, task_in_queue=task_in_queue, task_out_queue=task_out_queue, - mem_managers=mem_managers, + resource_managers=resource_managers, ) assert manager is not None task_out_queue.put("module_ready") @@ -93,19 +95,19 @@ def __init__( device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, - mem_managers: List[MemoryManager], + resource_managers: List[AttentionStateManager], ) -> 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.resource_managers = resource_managers - cur_mem_manager: MemoryManager = self.mem_managers[device_id] - kv_move_buffer = cur_mem_manager.alloc_paged_kv_move_buffer( - page_num=self.args.pd_kv_page_num, page_size=self.args.pd_kv_page_size + kv_move_buffer = allocate_transfer_buffer( + self.resource_managers[device_id], self.args.pd_kv_page_num, device=f"cuda:{device_id}" ) + self.kv_move_buffer = kv_move_buffer self.copy_cuda_stream = torch.cuda.Stream(priority=-1) self.transporter = create_kv_transporter( args=self.args, @@ -145,13 +147,16 @@ 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.write_mem_to_page_kv_move_buffer( - mem_indexes=[0], - page_index=0, - dp_index=dp_index, - mem_managers=self.mem_managers, - dp_world_size=self.dp_world_size, + cpu_token_indices = torch.tensor([0], dtype=torch.long, device="cpu") + token_indices = cpu_token_indices.cuda(device=self.kv_move_buffer.device, non_blocking=True) + copy_page( + node_att_managers=self.resource_managers, + pd_page_buffers=self.kv_move_buffer, + write_to_pd_page=True, + token_indices=token_indices, + pd_page_index=0, + dp_rank_in_node=dp_index, + tp_world_size=self.dp_world_size, ) torch.cuda.current_stream().synchronize() return @@ -203,15 +208,23 @@ 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, + # 只有 KV 页使用 token 索引；att_state 直接按请求槽位搬运运行态。 + token_indices = None + if trans_task.page_kind == "kv": + cpu_token_indices = torch.tensor(trans_task.mem_indexes, dtype=torch.long, device="cpu") + # 索引上传和后续页拷贝使用同一流，保证执行顺序。 + token_indices = cpu_token_indices.cuda(device=self.kv_move_buffer.device, non_blocking=True) + copy_page( + node_att_managers=self.resource_managers, + pd_page_buffers=self.kv_move_buffer, + write_to_pd_page=True, + token_indices=token_indices, + pd_page_index=trans_task.src_page_index, + dp_rank_in_node=trans_task.prefill_dp_index, + tp_world_size=self.dp_world_size, + task_kind=trans_task.page_kind, req_idx=trans_task.req_idx, + state_token_len=trans_task.end_kv_index if trans_task.page_kind == "att_state" else None, ) sync_event = torch.cuda.Event() sync_event.record() @@ -397,14 +410,14 @@ def _update_task_status_once(self): try: self.transporter.send_write_done_task_to_decode_node(completed_task) except BaseException as e: - logger.exception(f"send WRITE done nixl notify failed: {completed_task.to_str()}") - completed_task.error_info = f"send WRITE done nixl notify failed: {str(e)}" + logger.exception(f"send WRITE done PD notify failed: {completed_task.to_str()}") + completed_task.error_info = f"send WRITE done PD notify failed: {str(e)}" self.failed_queue.put(completed_task) continue if self.args.pd_detail_log: logger.debug( - f"send WRITE done nixl notify " + f"send WRITE done PD notify " f"request_id={completed_task.request_id} " f"kv=[{completed_task.start_kv_index},{completed_task.end_kv_index}) " f"src_page={completed_task.src_page_index} dst_page={completed_task.dst_page_index}" diff --git a/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py new file mode 100644 index 0000000000..516744de25 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/pd/resource_io.py @@ -0,0 +1,99 @@ +"""Select the local TP group and submit its State-owned PD operations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +import torch + +if TYPE_CHECKING: + from collections.abc import Sequence + from lightllm.common.attention_state.manager import AttentionStateManager + from typing import Literal + + +def allocate_transfer_buffer(resources: AttentionStateManager, page_count: int, device: str | torch.device | int): + if page_count <= 0: + raise ValueError("PD buffer must contain pages") + return torch.empty((page_count, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device=device) + + +def copy_page( + node_att_managers: Sequence[AttentionStateManager], + pd_page_buffers: torch.Tensor, + pd_page_index: int, + dp_rank_in_node: int, + tp_world_size: int, + token_indices: torch.Tensor | None, + write_to_pd_page: bool, + task_kind: Literal["kv", "att_state"] = "kv", + req_idx: int | None = None, + state_token_len: int | None = None, +) -> None: + """选择节点内指定 DP 的 TP 组，检查整组布局后提交 PD 页面读写。 + + node_att_managers 按节点内 DP 编号、组内 TP rank 顺序平铺。 + pd_page_buffers 是 PD 页集合，pd_page_index 选择本次传输的页。 + token_indices 在 kv 任务中为 CUDA 索引；att_state 任务传 None，不进行索引上传。 + write_to_pd_page=True 将 State 写入发送页；False 从接收页恢复 State。 + """ + tp_att_managers = node_att_managers[dp_rank_in_node * tp_world_size : (dp_rank_in_node + 1) * tp_world_size] + assert ( + dp_rank_in_node >= 0 + and tp_world_size > 0 + and len(tp_att_managers) == tp_world_size + and 0 <= pd_page_index < len(pd_page_buffers) + ), "invalid PD group or page index" + tp_transfer_layouts = [] + for att_manager in tp_att_managers: + pd_operator = att_manager.get_pd_operator() + match task_kind: + case "kv": + state_layouts = pd_operator.get_pd_full_att_page_layouts() + case "att_state": + state_layouts = pd_operator.get_pd_bounded_att_page_layouts() + case _: + assert False, f"unknown PD task kind {task_kind}" + assert ( + state_layouts and pd_operator.get_page_bytes() == pd_page_buffers.shape[1] + ), "PD page layout does not match this group" + tp_transfer_layouts.append(state_layouts) + + for att_manager, state_layouts in zip(tp_att_managers, tp_transfer_layouts): + pd_operator = att_manager.get_pd_operator() + match task_kind: + case "kv": + if write_to_pd_page: + pd_operator.write_page( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kind=task_kind, + layouts=state_layouts, + token_indices=token_indices, + ) + else: + pd_operator.read_page( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kind=task_kind, + layouts=state_layouts, + token_indices=token_indices, + ) + case "att_state": + if write_to_pd_page: + pd_operator.write_page( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kind=task_kind, + layouts=state_layouts, + req_idx=req_idx, + state_token_len=state_token_len, + ) + else: + pd_operator.read_page( + pd_page_buffers=pd_page_buffers, + pd_page_index=pd_page_index, + kind=task_kind, + layouts=state_layouts, + req_idx=req_idx, + state_token_len=state_token_len, + ) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/trans_process_obj.py b/lightllm/server/router/model_infer/mode_backend/pd/trans_process_obj.py index 84173a44ba..f180071fe3 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/trans_process_obj.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/trans_process_obj.py @@ -39,7 +39,7 @@ def init_all( up_status_in_queue, ) assert self.task_out_queue.get(timeout=30) == "proc_start" - assert self.task_out_queue.get(timeout=60) == "get_mem_managers_ok" + assert self.task_out_queue.get(timeout=60) == "get_resources_ok" return True diff --git a/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py b/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py index a34adb8c52..839c3b10ad 100644 --- a/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py +++ b/lightllm/server/router/model_infer/mode_backend/rl_backend_ops.py @@ -73,7 +73,8 @@ def _resume_memory_tags(self, tags: Optional[List[MemoryTag]]): for tag in memory_tags: self.backend.model.torch_memory_saver.resume(tag=tag) if MemoryTag.KV_CACHE in memory_tags: - self.backend.model.req_manager.init_hold_request_indexs() + model = self.backend.model + model.req_manager.bind_hold_page(model.att_state_manager.full_att_allocator.hold_token_indices) def release_memory_occupation(self, tags: Optional[List[MemoryTag]]): try: diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index edde24602d..a32cedd17d 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -24,6 +24,7 @@ def verify_mtp_tokens( b_mtp_index: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" + # 保留原调用接口中的步序参数；当前状态提交直接使用验证返回的接受长度。 accept_lengths, accepted_index = mtp_verify( req_to_next_token_ids=backend.model.req_manager.req_sampling_params_manager.req_to_next_token_ids, @@ -31,14 +32,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( - 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=backend.max_draft_step + 1, - ) + backend.model.att_state_manager.commit_bounded_att_state( + b_req_idx[b_req_mtp_start_loc], accept_lengths, role="main" + ) return accept_lengths, accepted_index diff --git a/lightllm/server/router/model_infer/shared_att_resources_utils.py b/lightllm/server/router/model_infer/shared_att_resources_utils.py new file mode 100644 index 0000000000..46a07e46f9 --- /dev/null +++ b/lightllm/server/router/model_infer/shared_att_resources_utils.py @@ -0,0 +1,76 @@ +"""分别发布 PD 运行态和 DP shared 资源，使用独立共享内存及读取计数。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal +from multiprocessing.reduction import ForkingPickler + +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 +from lightllm.utils.shm_utils import create_or_link_shm, get_service_shm_name + +if TYPE_CHECKING: + import torch + from lightllm.common.attention_state.manager import AttentionStateManager + + +def publish_att_resources( + resources: AttentionStateManager, + request_mapping: torch.Tensor, + kind: Literal["pd", "dp"], + reader_count: int, +) -> None: + """按用途选择导出副本；每个读者独立消费对应类别的 IPC 引用。""" + assert kind in ("pd", "dp"), f"unknown IPC storage kind {kind}" + if reader_count < 0: + raise ValueError("storage reader count cannot be negative") + if reader_count == 0: + return + if kv_trans_use_p2p(): + from .mode_backend.pd.p2p_fix import reduce_tensor + + mp.reductions.reduce_tensor.__code__ = reduce_tensor.__code__ + count = reader_count + if kind == "pd": + ipc_resources = resources.get_pd_operator().create_ipc_copy() + else: + ipc_resources = resources.get_dp_shared_operator().create_ipc_copy() + with FileLock(f"/tmp/{get_service_shm_name(f'rank_storage_{kind}_lock')}"): + # 每个读者独立序列化 CUDA IPC 引用；计数偏移变化可能使各份 pickle 长度不同。 + payloads = [ForkingPickler.dumps((ipc_resources, request_mapping)).tobytes() for _ in range(count)] + size = max(len(payload) for payload in payloads) + shm = create_or_link_shm( + name=f"rank_storage_{kind}_{get_current_rank_in_node()}", + expected_size=8 + count * size, + force_mode="create", + ) + shm.buf[:4] = count.to_bytes(4, "little") + shm.buf[4:8] = size.to_bytes(4, "little") + for index, payload in enumerate(payloads): + # 每个读者占用等长槽位；pickle.loads 会忽略完整对象之后的补零。 + shm.buf[8 + index * size : 8 + (index + 1) * size] = payload.ljust(size, b"\0") + shm.close() + + +def load_att_resources( + rank_in_node: int, + kind: Literal["pd", "dp"], +) -> tuple[AttentionStateManager, torch.Tensor]: + """读取指定用途的资源；DP 接收方随后通过 manager 统一注册共享页池。""" + assert kind in ("pd", "dp"), f"unknown IPC storage kind {kind}" + with FileLock(f"/tmp/{get_service_shm_name(f'rank_storage_{kind}_lock')}"): + shm = create_or_link_shm(name=f"rank_storage_{kind}_{rank_in_node}", expected_size=-1, force_mode="link") + try: + count = int.from_bytes(shm.buf[:4], "little") + size = int.from_bytes(shm.buf[4:8], "little") + if count <= 0: + raise RuntimeError("rank storage IPC exports are exhausted") + payload = bytes(shm.buf[8 + (count - 1) * size : 8 + count * size]) + resources, request_mapping = ForkingPickler.loads(payload) + shm.buf[:4] = (count - 1).to_bytes(4, "little") + return resources, request_mapping + finally: + shm.close() diff --git a/lightllm/server/router/req_queue/base_queue.py b/lightllm/server/router/req_queue/base_queue.py index 8d4ac38085..b306a5e904 100644 --- a/lightllm/server/router/req_queue/base_queue.py +++ b/lightllm/server/router/req_queue/base_queue.py @@ -3,7 +3,7 @@ from lightllm.utils.infer_utils import calculate_time from ..batch import Batch, Req from lightllm.server.core.objs import FinishStatus -from lightllm.utils.config_utils import get_fixed_kv_len, get_running_max_req_size_per_dp +from lightllm.utils.config_utils import get_running_max_req_size_per_dp from lightllm.server.core.objs import StartArgs from lightllm.utils.log_utils import init_logger @@ -18,11 +18,7 @@ def __init__(self, args: StartArgs, router, dp_index, dp_size_in_node) -> None: from lightllm.server.router.manager import RouterManager self.router: RouterManager = router - # max_total_token_num - get_fixed_kv_len() 是为了减去被特定 - # 推理模式预先占用了部分token kv 资源，这会导致整体可用的kv 资源 - # 在极端情况下减少，在非特定模式下，get_fixed_kv_len() 返回的都是 - # 0， 不会有任何影响。 - self.max_total_tokens = args.max_total_token_num - get_fixed_kv_len() + self.max_total_tokens = args.max_total_token_num self.running_max_req_size = get_running_max_req_size_per_dp(args) self.waiting_req_list: List[Req] = [] # List of queued requests self.router_token_ratio = args.router_token_ratio # ratio to determine whether the router is busy diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index 4b5ae89530..12ddbb1c68 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -55,8 +55,17 @@ def get_tokenizer( model_cfg, _ = PretrainedConfig.get_config_dict(tokenizer_name) model_type = model_cfg.get("model_type", "") if model_type == "deepseek_v32": - # Register the custom config before AutoTokenizer tries to resolve it. + # Register the custom config for HuggingFace metadata loading. from ..models import deepseek3_2 # noqa: F401 + from ..models.deepseek3_2.tokenizer import DeepSeekV32Tokenizer + + if tokenizer_mode == "slow": + raise ValueError("DeepSeek-V3.2 requires a fast tokenizer loaded from tokenizer.json") + # checkpoint 声明 LlamaTokenizerFast，但实际使用 ByteLevel；部分 Transformers + # 版本会按 Llama 类重建 SentencePiece 规则，丢失空格和换行，故直接保留 JSON 中的规则。 + tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_name, *args, **kwargs) + logger.info("Using DeepSeek-V3.2 tokenizer mode with Python-based chat template encoding.") + return DeepSeekV32Tokenizer(tokenizer) if model_type == "deepseek_v4": from ..models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer @@ -81,17 +90,6 @@ def get_tokenizer( "slowdown. Consider using a fast tokenizer instead." ) - # DeepSeek-V3.2 custom tokenizer mode: wraps the HF tokenizer with - # a Python-based apply_chat_template that uses encoding_dsv32.py. - if model_type == "deepseek_v32": - from ..models.deepseek3_2.tokenizer import DeepSeekV32Tokenizer - - hf_tokenizer = AutoTokenizer.from_pretrained( - tokenizer_name, trust_remote_code=trust_remote_code, *args, **kwargs - ) - logger.info("Using DeepSeek-V3.2 tokenizer mode with Python-based chat template encoding.") - return DeepSeekV32Tokenizer(hf_tokenizer) - if model_cfg["architectures"][0] == "TarsierForConditionalGeneration": from ..models.tarsier2.model import Tarsier2Tokenizer from ..models.qwen2_vl.vision_process import Qwen2VLImageProcessor diff --git a/lightllm/utils/backend_validator.py b/lightllm/utils/backend_validator.py index 0903b9b33c..3458f10131 100644 --- a/lightllm/utils/backend_validator.py +++ b/lightllm/utils/backend_validator.py @@ -99,11 +99,12 @@ def _validate_flashqla(): from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( chunk_gated_delta_rule as fla_chunk_gated_delta_rule, ) - from lightllm.common.state_cache_manager import LinearAttCacheConfig + from lightllm.models.qwen3next.config import Qwen3NextResourceConfig - linear_config = LinearAttCacheConfig.load_from_args() - num_k_heads = linear_config.num_linear_k_heads - num_v_heads = linear_config.num_linear_v_heads + linear_config = Qwen3NextResourceConfig.load_from_args() + # 配置只保存全局 head 数，校验算子的输入按当前 TP 大小计算本地分片。 + num_k_heads = max(1, linear_config.global_linear_k_heads // linear_config.tp_world_size) + num_v_heads = max(1, linear_config.global_linear_v_heads // linear_config.tp_world_size) head_k_dim = linear_config.head_linear_k_dim head_v_dim = linear_config.head_linear_v_dim qkv_dtype = linear_config.conv_state_dtype diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 7b258959e9..9457e77391 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -255,42 +255,6 @@ def get_hidden_size(model_path: str) -> Optional[int]: return None -@lru_cache(maxsize=None) -def get_num_key_value_heads(model_path: str) -> int: - num_key_value_heads = _get_config_llm_keyvalue(model_path=model_path, key_name=["num_key_value_heads"]) - if isinstance(num_key_value_heads, int): - return num_key_value_heads - return None - - -@lru_cache(maxsize=None) -def get_num_attention_heads(model_path: str) -> int: - num_attention_heads = _get_config_llm_keyvalue(model_path=model_path, key_name=["num_attention_heads"]) - if isinstance(num_attention_heads, int): - return num_attention_heads - return None - - -@lru_cache(maxsize=None) -def get_head_dim(model_path: str) -> int: - head_dim = _get_config_llm_keyvalue(model_path=model_path, key_name=["head_dim"]) - if isinstance(head_dim, int): - return head_dim - - # calcu head_dim - head_dim = get_hidden_size(model_path=model_path) // get_num_attention_heads(model_path=model_path) - - return head_dim - - -@lru_cache(maxsize=None) -def get_layer_num(model_path: str) -> int: - num_hidden_layers = _get_config_llm_keyvalue(model_path=model_path, key_name=["num_hidden_layers"]) - if isinstance(num_hidden_layers, int): - return num_hidden_layers - return None - - def get_eos_token_ids(model_path: str) -> Optional[List[int]]: # gemma4 special eos_token_id try: @@ -408,19 +372,6 @@ def get_dtype(model_path: str): return torch_dtype -@lru_cache(maxsize=None) -def get_fixed_kv_len(): - start_args = get_env_start_args() - model_cfg = get_config_json(start_args.model_dir) - if "prompt_cache_token_ids" in model_cfg: - fixed_kv_len = len(model_cfg["prompt_cache_token_ids"]) - # 固定 KV 最终会插入 radix cache，只加载完整的模型 KV 页面；不足一页 - # 的尾部直接截断，因此 router 也只扣除实际常驻的页面容量。 - return fixed_kv_len // start_args.page_size * start_args.page_size - else: - return 0 - - @lru_cache(maxsize=None) def has_vision_module(model_path: str) -> bool: try: diff --git a/lightllm/utils/cpu_cache_layout.py b/lightllm/utils/cpu_cache_layout.py new file mode 100644 index 0000000000..05c90e17f6 --- /dev/null +++ b/lightllm/utils/cpu_cache_layout.py @@ -0,0 +1,63 @@ +"""CPU cache 服务布局构建；仅依赖 State 配置，不加载静态标定。""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from lightllm.common.attention_state import AttentionStateManager +from lightllm.utils.envs_utils import get_env_start_args + +if TYPE_CHECKING: + from collections.abc import Mapping + from lightllm.common.attention_state.types import StateTransferLayout + + +def build_cpu_cache_layouts() -> Mapping[str, StateTransferLayout]: + """使用启动流程相同的 meta State 定义，由管理器统一生成 CPU 服务页布局。""" + from transformers.configuration_utils import PretrainedConfig + from lightllm.models import get_draft_model_class + from lightllm.common.attention_state.utils import MAX_DRAFT_MODELS + from lightllm.utils.llm_utils import get_llm_model_class + + # 阶段一：读取 main 和启用 attention 的 draft 配置，不构造实际模型。 + args = get_env_start_args() + main_config, _ = PretrainedConfig.get_config_dict(args.model_dir) + models = [("main", get_llm_model_class(), main_config)] + if args.mtp_mode is not None and args.mtp_mode not in ("vanilla_no_att", "eagle_no_att"): + count = args.mtp_step if args.mtp_mode == "vanilla_with_att" else 1 + assert args.mtp_draft_model_dir is not None and len(args.mtp_draft_model_dir) >= count + assert 0 < count <= MAX_DRAFT_MODELS, "at most 36 draft models are supported" + for index, directory in enumerate(args.mtp_draft_model_dir[:count]): + config, _ = PretrainedConfig.get_config_dict(directory) + models.append((f"draft{index}", get_draft_model_class(config, args.mtp_mode), config)) + + # 阶段二：复用模型定义入口创建 meta State，并累计逻辑层号；CPU 页不包含静态标定参数。 + manager = AttentionStateManager(request_capacity=0) # 只描述 meta 布局，不分配请求或缓存存储。 + logical_layer_start = 0 + for role, model_class, config in models: + config = model_class.normalize_model_config(config) + count = config["num_hidden_layers"] + logical_layer_ids = tuple(range(logical_layer_start, logical_layer_start + count)) + full_states = model_class.create_full_att_states( + model_config=config, + args=args, + model_logical_layer_ids=logical_layer_ids, + ) + bounded_states = model_class.create_bounded_att_states( + model_config=config, + args=args, + model_logical_layer_ids=logical_layer_ids, + request_capacity=1, + device="meta", + ) + # 无重叠的池化窗口自行返回空 CPU 区域；这里只创建 meta State，不初始化 CUDA。 + for state in (*full_states, *bounded_states): + manager.register_state(role, state) + logical_layer_start += count + + # 阶段三：确认混合模型的页边界一致，由管理器聚合 CPU 联合页布局。 + manager.finish_state_registration() + if manager.get_bounded_attention_states(): + assert ( + args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num + ), "hybrid CPU pages must match the large state-page token boundary" + return manager.get_cpu_cache_operator().get_page_layouts() diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 4f93317d61..f7720e64fb 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -12,23 +12,12 @@ get_env_start_args, enable_huge_page, enable_cpu_cache_numa_interleave, - get_llm_data_type, - get_added_mtp_kv_layer_num, ) from lightllm.utils.log_utils import init_logger -from lightllm.utils.config_utils import get_num_key_value_heads, get_head_dim, get_layer_num, is_hybrid_att_model -from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.common.kv_cache_mem_manager import ( - MemoryManager, - PPLINT8KVMemoryManager, - PPLINT4KVMemoryManager, - Deepseek2MemoryManager, -) from typing import List, Tuple, Optional from tqdm import tqdm from lightllm.utils.dist_utils import get_current_device_id -from lightllm.common.state_cache_manager import get_hybrid_cache_config logger = init_logger(__name__) @@ -62,68 +51,20 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": args = get_env_start_args() assert args.enable_cpu_cache - is_hybrid_model = is_hybrid_att_model(args.model_dir) - mem_manager_class = None if is_hybrid_model else select_mem_manager_class() - if is_hybrid_model: - hybrid_config = get_hybrid_cache_config() - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=1, - layer_num=1, - num_heads=1, - head_dim=hybrid_config.get_cpu_cache_big_page_bytes(), - data_type=torch.uint8, - scale_head_dim=0, - scale_data_type=get_llm_data_type(), - ) - elif mem_manager_class is Deepseek2MemoryManager: - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=args.cpu_cache_token_page_size, - layer_num=get_layer_num(args.model_dir), - num_heads=1, - head_dim=512 + 64, - data_type=get_llm_data_type(), - scale_head_dim=0, - scale_data_type=get_llm_data_type(), - ) - elif mem_manager_class is MemoryManager: - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=args.cpu_cache_token_page_size, - layer_num=get_layer_num(args.model_dir), - num_heads=get_num_key_value_heads(args.model_dir) * 2, - head_dim=get_head_dim(args.model_dir), - data_type=get_llm_data_type(), - scale_head_dim=0, - scale_data_type=get_llm_data_type(), - ) - elif mem_manager_class is PPLINT8KVMemoryManager: - cpu_cache_meta = CpuKVCacheMeta( - page_num=0, - token_page_size=args.cpu_cache_token_page_size, - layer_num=get_layer_num(args.model_dir), - num_heads=get_num_key_value_heads(args.model_dir) * 2, - head_dim=get_head_dim(args.model_dir), - data_type=torch.int8, - scale_head_dim=get_head_dim(args.model_dir) // 8, - scale_data_type=get_llm_data_type(), - ) - else: - logger.error(f"not support mem manager: {mem_manager_class} for cpu kv cache") - raise Exception(f"not support mem manager: {mem_manager_class} for cpu kv cache") - - if args.mtp_mode is not None: - # TODO 可能会存在不同mtp模式的精度问题 - if not is_hybrid_model: - # 对于非 hybrid 模型，需要额外增加 mtp 的 kv 层数， - # 对于 hybrid 模型，如 qwen 3.5 mtp，已经将 kv 数据 - # 打包成一个块了，所以不需要额外增加，其 layer_num 一直都保持为 1 - cpu_cache_meta.layer_num += get_added_mtp_kv_layer_num() - - cpu_cache_page_num = int( - (args.cpu_cache_storage_size * 1024 * 1024 * 1024) / (cpu_cache_meta.calcu_one_page_size()) + from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts + + layouts = build_cpu_cache_layouts() + page_bytes = 0 + for layout in layouts.values(): + page_bytes = max(page_bytes, layout.byte_end) + # State 区间按 8 字节对齐，整个 CPU 物理页再按 16 字节对齐。 + page_bytes = (page_bytes + 15) // 16 * 16 + cpu_cache_meta = CpuKVCacheMeta( + page_num=0, + page_bytes=page_bytes, ) + + cpu_cache_page_num = int((args.cpu_cache_storage_size * 1024 * 1024 * 1024) / (cpu_cache_meta.page_bytes)) cpu_cache_meta.page_num = cpu_cache_page_num logger.info(f"cpu kv cache page num: {cpu_cache_meta.page_num}") @@ -133,36 +74,18 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": @dataclasses.dataclass class CpuKVCacheMeta: + """External byte-page storage, independent of model tensor dimensions.""" + page_num: int - token_page_size: int - layer_num: int - num_heads: int - head_dim: int - data_type: torch.dtype - scale_head_dim: int - scale_data_type: torch.dtype - - def calcu_size(self): - return self.page_num * self.calcu_one_page_size() - - def calcu_one_page_size(self): - return ( - self.token_page_size - * self.layer_num - * self.num_heads - * (self.head_dim * self.data_type.itemsize + self.scale_head_dim * self.scale_data_type.itemsize) - ) + page_bytes: int + + def __post_init__(self): + if self.page_num < 0 or self.page_bytes <= 0: + raise ValueError("CPU cache requires a nonnegative page count and positive page size") - def get_merged_head_dim(self): - """ - 返回将head_dim 和 scale_head_dim 看成融合成一个head_dim时候, head_dim的长度。 - """ - assert ( - self.head_dim * self.data_type.itemsize + self.scale_head_dim * self.scale_data_type.itemsize - ) % self.data_type.itemsize == 0 - return ( - self.head_dim * self.data_type.itemsize + self.scale_head_dim * self.scale_data_type.itemsize - ) // self.data_type.itemsize + @property + def size_bytes(self): + return self.page_num * self.page_bytes @lru_cache(maxsize=None) diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py new file mode 100644 index 0000000000..4e4317f3fa --- /dev/null +++ b/test/att_resources/validate_model_cache.py @@ -0,0 +1,366 @@ +"""Validate resource restoration against an uninterrupted, real model decode. + +Run from the repository root with local model weights, for example: + PYTHONPATH=. python test/att_resources/validate_model_cache.py \ + --model_dir /models/Qwen3.5-0.8B --report /tmp/hybrid-cache.json + +This checks GPU pages, CPU pages, DP copies and PD buffers on one GPU. Network +transport is tested separately. Each CPU page captures a real forward boundary. +Decode uses CUDA graphs by default; pass --disable_cudagraph to check eager mode. +Native MTP can be checked with --mtp_mode vanilla_with_att --mtp_step 1 and +--mtp_draft_model_dir pointing to the checkpoint containing its draft weights. +MTP uses identical candidate packets with varying accepted widths, not sampling. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +import json +from pathlib import Path +import socket +import sys + +import numpy as np +import torch +import torch.distributed as dist +from transformers import PretrainedConfig + +from lightllm.distributed import dist_group_manager +from lightllm.models import get_model +from lightllm.server.api_cli import make_argument_parser +from lightllm.utils.dist_utils import init_distributed_env +from lightllm.utils.envs_utils import get_unique_server_name, set_env_start_args, set_unique_server_name +from lightllm.utils.service_shm_cleanup import ServiceShmCleanup + +if TYPE_CHECKING: + from argparse import Namespace + from collections.abc import Sequence + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "benchmark" / "static_inference")) +from static_benchmark import ( + StaticBenchmarkExecutor, + TokenSource, + add_static_benchmark_args, + build_cases, + build_model_kvargs, + init_mtp_draft_models, + normalize_args, +) + + +@torch.inference_mode() +def validate(model: TpPartBaseModel, args: Namespace, draft_models: Sequence[TpPartBaseModel] = ()): + executor = StaticBenchmarkExecutor( + args=args, + model=model, + draft_models=draft_models, + token_source=TokenSource(args), + ) + resources = model.att_state_manager + (full_state,) = resources.get_full_attention_states("main") + for draft in draft_models: + assert draft.att_state_manager is resources + assert not hasattr(draft, "full_att_state") + (draft_state,) = resources.get_full_attention_states("draft0") + assert draft_state is not full_state + assert draft_state.capacity == full_state.capacity == resources.full_att_allocator.size + assert draft.req_manager is model.req_manager + prefix_len, page_tokens = args.input_len, args.cpu_cache_token_page_size + if prefix_len % page_tokens: + raise ValueError("input_len must contain complete CPU pages") + pages = list(range(prefix_len // page_tokens)) + lengths = [(page + 1) * page_tokens for page in pages] + hybrid = model.is_hybrid_att_model + big = resources.bounded_att_big_page_allocator + source_slots = big.alloc(len(pages)) if hybrid else None + allocated_slots = list(source_slots) if hybrid else [] + source_req = executor._alloc_req_indexes(1) + resources.reset_bounded_att_state(source_req.tolist()) + tokens = np.random.default_rng(42).integers( + low=100, + high=10000, + size=(1, prefix_len), + dtype=np.int64, + ) + for page, length in zip(pages, lengths): + batch = executor._make_prefill_input(tokens[:, page * page_tokens : length], source_req, page * page_tokens) + output = model.forward(batch) + if draft_models: + executor._fill_mtp_prefill_kv( + case=None, + main_prefill_input=batch, + main_output=output, + first_next_ids=executor._argmax_ids(output.logits), + ) + if hybrid: + resources.save_bounded_att_runtime_to_pages( + req_indices=source_req.tolist(), + state_token_lens=[length], + page_kind="big", + page_slot_ids=[source_slots[page]], + ) + torch.cuda.synchronize() + next_token = output.logits.argmax(dim=-1).cpu() + source_indices = model.req_manager.req_to_token_indexs[int(source_req[0]), :prefix_len].clone() + + cpu = torch.zeros( + (len(pages), resources.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) + resources.get_cpu_cache_operator().write_req( + token_indices=source_indices, + cpu_page_buffers=cpu, + cpu_page_indices=pages, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + big_page_slot_ids=source_slots if hybrid else None, + ) + # Host-side checkpoint repacking consumes the CPU page on the next call. + torch.cuda.synchronize() + full_att_layouts = resources.get_pd_operator().get_pd_full_att_page_layouts() + runtime_layouts = resources.get_pd_operator().get_pd_bounded_att_page_layouts() + pd = {"kv": torch.empty((1, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda")} + if runtime_layouts: + pd["att_state"] = torch.empty( + (1, resources.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda" + ) + resources.get_pd_operator().write_page( + pd_page_buffers=pd["kv"], + pd_page_index=0, + kind="kv", + layouts=full_att_layouts, + token_indices=source_indices, + ) + if "att_state" in pd: + resources.get_pd_operator().write_page( + pd_page_buffers=pd["att_state"], + pd_page_index=0, + kind="att_state", + layouts=runtime_layouts, + req_idx=int(source_req[0]), + state_token_len=prefix_len, + ) + + targets = {} + for kind in ("gpu_pages", "cpu_pages", "dp_copy", "pd_buffer"): + request = executor._alloc_req_indexes(1) + resources.reset_bounded_att_state(request.tolist()) + executor._ensure_req_kv_capacity(request, torch.tensor([prefix_len], dtype=torch.int32)) + indices = model.req_manager.req_to_token_indexs[int(request[0]), :prefix_len].clone() + assert not set(indices.tolist()) & set(source_indices.tolist()) + if kind == "gpu_pages": + resources.copy_full_att_tokens(source_indices, indices) + if hybrid: + resources.restore_bounded_att_runtime_from_pages( + req_indices=request.tolist(), + state_token_lens=[prefix_len], + page_kind="big", + page_slot_ids=[source_slots[-1]], + ) + elif kind == "cpu_pages": + slots = big.alloc(len(pages)) if hybrid else None + if hybrid: + allocated_slots.extend(slots) + resources.get_cpu_cache_operator().read_req( + cpu_page_buffers=cpu, + cpu_page_indices=pages, + token_indices=indices, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + big_page_slot_ids=slots, + req_idx=int(request[0]) if hybrid else None, + state_token_len=prefix_len if hybrid else None, + ) + elif kind == "dp_copy": + resources.get_dp_shared_operator().copy_from_dp( + remote_att_manager=resources, + src_token_indices=source_indices, + dst_token_indices=indices, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=int(source_req[0]), + dst_req_idx=int(request[0]), + ) + else: + resources.get_pd_operator().read_page( + pd_page_buffers=pd["kv"], + pd_page_index=0, + kind="kv", + layouts=full_att_layouts, + token_indices=indices, + ) + if "att_state" in pd: + resources.get_pd_operator().read_page( + pd_page_buffers=pd["att_state"], + pd_page_index=0, + kind="att_state", + layouts=runtime_layouts, + req_idx=int(request[0]), + state_token_len=prefix_len, + ) + targets[kind] = request + + errors = {kind: [] for kind in targets} + draft_errors = {kind: [] for kind in targets} if draft_models else None + generated = [] + accepted_widths = [] + sequence_length = prefix_len + step_width = args.mtp_step + 1 if draft_models else 1 + for step in range(args.output_len): + generated.append(int(next_token[0])) + accepted_width = step % step_width + 1 + accepted_widths.append(accepted_width) + + def decode(request: torch.Tensor): + request_rows, mtp_index = executor._build_mtp_decode_index_tensors(request, step_width) + lengths = executor._build_mtp_seq_len(torch.tensor([sequence_length], dtype=torch.int32), step_width) + batch = executor._make_decode_input( + batch_size=step_width, + req_idx=request_rows, + mtp_index=mtp_index, + seq_len=lengths, + input_ids=next_token.repeat(step_width), + max_kv_seq_len=int(lengths[-1]), + ) + output = model.forward(batch) + logits = [output.logits.clone()] + for draft in draft_models: + batch.input_ids = output.logits.argmax(dim=-1) + batch.mtp_draft_input_hiddens = output.mtp_collector.spec_hidden + output = draft.forward(batch) + logits.append(output.logits.clone()) + if draft_models and hybrid: + cuda_req_indices = request.cuda(non_blocking=True) + accepted_steps = torch.full_like(cuda_req_indices, accepted_width, dtype=torch.int32) + for item in (model, *draft_models): + resources.commit_bounded_att_state( + req_indices=cuda_req_indices, + accepted_steps=accepted_steps, + role=item.att_state_role, + ) + return logits + + reference = decode(source_req) + assert all(torch.isfinite(logits).all() for logits in reference) + for kind, request in targets.items(): + actual = decode(request) + for restored_logits, reference_logits in zip(actual, reference): + torch.testing.assert_close( + actual=restored_logits, + expected=reference_logits, + rtol=1e-4, + atol=1e-4, + ) + assert torch.equal(restored_logits.argmax(dim=-1), reference_logits.argmax(dim=-1)) + errors[kind].append(float((actual[0] - reference[0]).abs().max())) + if draft_models: + draft_errors[kind].append([float((a - b).abs().max()) for a, b in zip(actual[1:], reference[1:])]) + next_token = reference[0][accepted_width - 1 : accepted_width].argmax(dim=-1).cpu() + sequence_length += accepted_width + + torch.cuda.synchronize() + if hybrid: + big.free(allocated_slots) + executor._reset_model_cache() + assert resources.full_att_allocator.can_use_mem_size == resources.full_att_allocator.size + return { + "model": args.model_dir, + "model_class": type(model).__name__, + "transport": "local_buffers", + "prefix_tokens": prefix_len, + "cpu_pages": len(pages), + "decode_steps": args.output_len, + "generated_token_ids": generated, + "max_logit_errors": errors, + "draft_models": [type(draft).__name__ for draft in draft_models], + "max_draft_logit_errors": draft_errors, + "accepted_widths": accepted_widths, + "allocator_reclaimed": True, + "cuda_graphs": not args.disable_cudagraph, + } + + +def main(): + parser = make_argument_parser() + add_static_benchmark_args(parser) + parser.add_argument("--report", required=True) + parser.set_defaults( + tp=1, + dp=1, + data_type="bfloat16", + batch_size=1, + max_batch_size=1, + input_len=128, + output_len=4, + batch_max_tokens=256, + max_total_token_num=2048, + running_max_req_size=8, + graph_max_batch_size=2, + max_req_total_len=512, + page_size=16, + cpu_cache_token_page_size=64, + enable_cpu_cache=True, + linear_att_hash_page_size=64, + linear_att_page_block_num=1, + linear_att_cache_size=16, + llm_prefill_att_backend=["fa3", "triton"], + llm_decode_att_backend=["fa3", "triton"], + disable_cudagraph=False, + pd_kv_page_size=2048, + ) + parsed = parser.parse_args() + if parsed.tp != 1 or parsed.dp != 1 or parsed.mtp_mode not in (None, "vanilla_with_att"): + parser.error("this validation uses TP=1, DP=1 and optionally native vanilla_with_att MTP") + with socket.socket() as port: + port.bind(("127.0.0.1", 0)) + parsed.nccl_port = port.getsockname()[1] + args = normalize_args(parsed, build_cases(parsed)) + set_unique_server_name(args) + set_env_start_args(args) + torch.set_num_threads(4) + torch.manual_seed(42) + try: + kvargs = build_model_kvargs(args, 0) + init_distributed_env(kvargs) + dist_group_manager.create_groups(group_size=1) + config, _ = PretrainedConfig.get_config_dict(args.model_dir) + model, _ = get_model(config, kvargs) + draft_models = init_mtp_draft_models(args, kvargs, model) + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.req_manager import ReqManager + + models = (model, *draft_models) + assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" + assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" + # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 + max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) + with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + req_manager = ReqManager(max_request_num=model.max_req_num, max_sequence_length=max_seq_len) + for item in models: + item.bind_req_manager(req_manager=req_manager) + for item in models: + item.prepare_attention_states() + manager: AttentionStateManager = model.att_state_manager + assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" + with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.finish_state_registration() + manager.load_full_att_calibration(args=model.args) + manager.initialize_storage( + model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction + ) + for item in models: + item.init_runtime() + torch.cuda.empty_cache() + set_model_init_status(True) + report = validate(model, args, draft_models) + Path(args.report).write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2), flush=True) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + ServiceShmCleanup(get_unique_server_name()).cleanup_service_resources() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/static_inference/static_benchmark.py b/test/benchmark/static_inference/static_benchmark.py index 279840d5af..980c1f20e1 100644 --- a/test/benchmark/static_inference/static_benchmark.py +++ b/test/benchmark/static_inference/static_benchmark.py @@ -568,8 +568,8 @@ def _ensure_req_kv_capacity(self, req_idx: torch.Tensor, target_kv_lens: torch.T req_id = int(req_id) max_target_len_by_req[req_id] = max(max_target_len_by_req.get(req_id, 0), int(target_len)) - mem_manager = self.model.req_manager.mem_manager - page_size = mem_manager.page_size + allocator = self.model.att_state_manager.full_att_allocator + page_size = allocator.page_size for req_id, target_len in max_target_len_by_req.items(): new_hold_kv_len = align_up(target_len, page_size) old_hold_kv_len = self._hold_kv_len_by_req.get(req_id, 0) @@ -577,7 +577,7 @@ def _ensure_req_kv_capacity(self, req_idx: torch.Tensor, target_kv_lens: torch.T continue alloc_token_num = new_hold_kv_len - old_hold_kv_len - mem_indexes = mem_manager.alloc(alloc_token_num) + mem_indexes = allocator.alloc(alloc_token_num) self.model.req_manager.req_to_token_indexs[req_id, old_hold_kv_len:new_hold_kv_len].copy_(mem_indexes) self._hold_kv_len_by_req[req_id] = new_hold_kv_len @@ -729,7 +729,7 @@ def _alloc_req_indexes(self, batch_size: int) -> torch.Tensor: return torch.tensor(req_indexes, dtype=torch.int32, device="cpu") def _reset_model_cache(self): - self.model.mem_manager.free_all() + self.model.att_state_manager.full_att_allocator.free_all() self.model.req_manager.free_all() self._hold_kv_len_by_req.clear() torch.cuda.synchronize() @@ -1182,13 +1182,12 @@ def init_mtp_draft_models(args: SimpleNamespace, main_kvargs: Dict, main_model) if args.mtp_mode not in MTP_MODES: return [] - os.environ["DISABLE_CHECK_MAX_LEN_INFER"] = "1" draft_models = [] for draft_dir in args.mtp_draft_model_dir: mtp_cfg, _ = PretrainedConfig.get_config_dict(draft_dir) mtp_kvargs = { "weight_dir": draft_dir, - "max_total_token_num": main_model.mem_manager.size, + "max_total_token_num": main_kvargs["max_total_token_num"], "load_way": main_kvargs["load_way"], "max_req_num": main_kvargs["max_req_num"], "max_seq_length": main_kvargs["max_seq_length"], @@ -1271,11 +1270,49 @@ def run_worker(args_dict: Dict, case_dicts: List[Dict], rank_id: int, ans_queue) torch.cuda.empty_cache() model, _ = get_model(model_cfg, model_kvargs) - cases = resolve_batch_max_prefill_cases(args, cases, model.mem_manager.size) - cases = resolve_profile_decode_cases(args, cases, model.mem_manager.size) + draft_models = init_mtp_draft_models(args, model_kvargs, model) + from lightllm.utils.envs_utils import set_model_init_status + from lightllm.utils.torch_memory_saver_utils import MemoryTag + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.req_manager import ReqManager + + models = (model, *draft_models) + assert all(getattr(item, "att_state_manager", None) is None for item in models), "resources already initialized" + assert all(item.max_req_num == model.max_req_num for item in models), "models must share the request domain" + # main/draft 共用请求映射；按所有模型的最大需求分配，再逐模型绑定。 + max_seq_len = max(item.get_req_manager_max_seq_len() for item in models) + with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + req_manager = ReqManager(max_request_num=model.max_req_num, max_sequence_length=max_seq_len) + for item in models: + item.bind_req_manager(req_manager=req_manager) + for item in models: + item.prepare_attention_states() + manager: AttentionStateManager = model.att_state_manager + assert all(item.att_state_manager is manager for item in models), "models must share one resource manager" + with model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + manager.finish_state_registration() + manager.load_full_att_calibration(args=model.args) + manager.initialize_storage( + model.args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction + ) + for item in models: + item.init_runtime() + torch.cuda.empty_cache() + set_model_init_status(True) + cases = resolve_batch_max_prefill_cases(args, cases, model.att_state_manager.full_att_allocator.size) + cases = resolve_profile_decode_cases(args, cases, model.att_state_manager.full_att_allocator.size) if defer_cudagraph: init_deferred_cudagraph(args, cases, model_kvargs, model) - draft_models = init_mtp_draft_models(args, model_kvargs, model) + for draft in draft_models: + batch_size = args.graph_max_batch_size + if args.enable_decode_microbatch_overlap: + batch_size //= 2 + batch_size *= draft.mtp_manager.get_decode_tokens_per_request(True) + draft.graph_max_batch_size = draft._align_decode_batch_size(batch_size) + draft.disable_cudagraph = False + draft._init_att_backend() + draft._init_att_backend1() + draft._init_cudagraph() token_source = TokenSource(args) executor = StaticBenchmarkExecutor(args, model, draft_models, token_source) diff --git a/test/cpu_cache_kernel/test_speed.py b/test/cpu_cache_kernel/test_speed.py index d59142e09f..322d1b148d 100644 --- a/test/cpu_cache_kernel/test_speed.py +++ b/test/cpu_cache_kernel/test_speed.py @@ -1,276 +1,156 @@ -""" -Speed benchmark for copy_cpu_cache_to_kv_buffer in linear_att_cpu_cache_copy.py. +"""Measure the public CPU-page restore path for one TP shard of a hybrid model. -Test configuration (matching the user's LinearAttCacheConfig): - tp_world_size=8, full_att_all_num_kv_heads=2, full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, full_att_head_dim=256, - num_linear_k_heads=2, num_linear_v_heads=8, - head_linear_k_dim=128, head_linear_v_dim=128, - conv_kernel_size=4, 36 linear layers and 12 full-attention layers, - conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - layer_types=(linear_attention, linear_attention, linear_attention, full_attention) * 12 +Run from the repository root, for example: + PYTHONPATH=. python test/cpu_cache_kernel/test_speed.py --tokens-per-page 16384 --iterations 100 +The reported time includes validation, KV transfers, and CPU state-page repacking. """ -import os -import json +import argparse +import statistics import time -import triton -import torch -from easydict import EasyDict - -# --------------------------------------------------------------------------- -# Step 0 – set up environment args BEFORE any import that calls -# get_env_start_args() / LinearAttCacheConfig.load_from_args(). -# --------------------------------------------------------------------------- -_env_args = { - "cpu_cache_token_page_size": 2048 * 8, # big_page_token_num - "linear_att_hash_page_size": 2048, - "linear_att_page_block_num": 8, # 512 * 1 == 512 - "data_type": "bfloat16", - "linear_att_ssm_data_type": "float32", - "model_dir": "/tmp/fake_model", # dummy – not used when config is built directly - "tp": 8, - "dp": 1, - "running_max_req_size": 2048, - "enable_cpu_cache": True, -} -os.environ["LIGHTLLM_START_ARGS"] = json.dumps(_env_args) - -# --------------------------------------------------------------------------- -# Step 1 – build LinearAttCacheConfig directly (avoids needing a real model dir) -# --------------------------------------------------------------------------- -from lightllm.common.state_cache_manager import LinearAttCacheConfig - -linear_config = LinearAttCacheConfig( - tp_world_size=8, - full_att_all_num_kv_heads=2, - full_att_dtype=torch.bfloat16, - full_att_num_kv_heads=1, - full_att_head_dim=256, - global_linear_k_heads=16, - global_linear_v_heads=64, - num_linear_k_heads=2, - num_linear_v_heads=8, - head_linear_k_dim=128, - head_linear_v_dim=128, - conv_kernel_size=4, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=torch.float32, - layer_types=("linear_attention", "linear_attention", "linear_attention", "full_attention") * 12, -) -print(f"LinearAttCacheConfig:\n{linear_config}\n", flush=True) - -# --------------------------------------------------------------------------- -# Step 2 – derive sizes from the config -# --------------------------------------------------------------------------- -big_page_token_num = _env_args["cpu_cache_token_page_size"] # 512 -full_att_layer_num = linear_config.get_main_model_full_att_layer_num() # 12 - -full_att_bytes = linear_config.get_cpu_cache_full_att_bytes() # per big page -conv_bytes = linear_config.get_cpu_cache_conv_bytes() -ssm_bytes = linear_config.get_cpu_cache_ssm_bytes() -total_bytes = full_att_bytes + conv_bytes + ssm_bytes -print( - f"Per-page bytes full_att={full_att_bytes:,} conv={conv_bytes:,} " f"ssm={ssm_bytes:,} total={total_bytes:,}", - flush=True, -) -total_bytes = linear_config.get_cpu_cache_big_page_bytes() - -# --------------------------------------------------------------------------- -# Step 3 – allocate tensors -# --------------------------------------------------------------------------- -grid_num = 8 -PAGE_NUM = 1 # number of big pages to copy per call -SEQ_LEN = 2048 * 8 # total sequence length in gpu_full_att_kv_state dim-1 -BIG_PAGE_COUNT = PAGE_NUM # big_page_buffer_ids length == page_indexes length - -# --- GPU tensors --- -mem_indexes = torch.arange(0, big_page_token_num * PAGE_NUM, dtype=torch.int64, device="cpu") -big_page_buffer_ids = torch.arange(0, BIG_PAGE_COUNT, dtype=torch.int64, device="cpu") -page_indexes = torch.arange(0, PAGE_NUM, dtype=torch.int32, device="cpu") - -gpu_full_att_kv_state = torch.empty( - ( - full_att_layer_num, - SEQ_LEN, - 2 * max(1, linear_config.full_att_num_kv_heads // linear_config.tp_world_size), - linear_config.full_att_head_dim, - ), - dtype=linear_config.full_att_dtype, - device="cuda", -) - -# --- CPU tensors --- -buffer_count = triton.cdiv(SEQ_LEN, big_page_token_num) + 2 # matches Qwen3NextMemManager - - -conv_shape = linear_config.get_conv_state_shape() -cpu_kv_conv_state = torch.empty( - (buffer_count, linear_config.linear_layer_num, *conv_shape), - dtype=linear_config.conv_state_dtype, - device="cpu", - pin_memory=True, -) - -ssm_shape = linear_config.get_ssm_state_shape() # (num_linear_v_heads, head_linear_k_dim, head_linear_v_dim) -cpu_kv_ssm_state = torch.empty( - (buffer_count, linear_config.linear_layer_num, *ssm_shape), - dtype=linear_config.ssm_state_dtype, - device="cpu", - pin_memory=True, -) - -# conv_shape = linear_config.get_conv_state_shape() -# cpu_kv_conv_state = torch.empty( -# (buffer_count, linear_config.linear_layer_num, *conv_shape), -# dtype=linear_config.conv_state_dtype, device="cuda", -# ) - -# ssm_shape = linear_config.get_ssm_state_shape() # (num_linear_v_heads, head_linear_k_dim, head_linear_v_dim) -# cpu_kv_ssm_state = torch.empty( -# (buffer_count, linear_config.linear_layer_num, *ssm_shape), -# dtype=linear_config.ssm_state_dtype, device="cuda", -# ) - -# cpu_cache_tensor: [page_num, 1, 1, 1, total_bytes] -cpu_cache_tensor = torch.empty( - (PAGE_NUM, 1, 1, 1, total_bytes), - dtype=torch.uint8, - device="cpu", - pin_memory=True, -) - -# Move GPU tensors to CUDA -mem_indexes_cuda = mem_indexes.cuda(non_blocking=True) -big_page_buffer_ids_cuda = big_page_buffer_ids.cuda(non_blocking=True) -page_indexes_cuda = page_indexes.cuda(non_blocking=True) -gpu_full_att_kv_state = gpu_full_att_kv_state.cuda(non_blocking=True) - -torch.cuda.synchronize() -print("All tensors allocated and moved to GPU.\n", flush=True) - -# --------------------------------------------------------------------------- -# Step 4 – import and warm-up the triton kernel -# --------------------------------------------------------------------------- -from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_cpu_cache_to_kv_buffer, -) +import torch -print("Warming up …", flush=True) -copy_cpu_cache_to_kv_buffer( - mem_indexes=mem_indexes_cuda, - big_page_buffer_ids=big_page_buffer_ids_cuda, - page_indexes=page_indexes_cuda, - gpu_full_att_kv_state=gpu_full_att_kv_state, - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_cache_tensor=cpu_cache_tensor, - tp_rank=0, - tp_world_size=linear_config.tp_world_size, - big_page_token_num=big_page_token_num, - linear_config=linear_config, - grid_num=grid_num, +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.full_att import ( + GQAFullAttentionState, + GQAAttentionConfig, ) -torch.cuda.synchronize() -print("Warm-up done.\n", flush=True) - -# --------------------------------------------------------------------------- -# Step 5 – benchmark -# --------------------------------------------------------------------------- -WARMUP_ITERS = 10 -BENCH_ITERS = 100 - -print(f"Benchmarking ({BENCH_ITERS} iterations, {PAGE_NUM} pages / {big_page_token_num} tokens each) …", flush=True) - -# Warm-up -for _ in range(WARMUP_ITERS): - copy_cpu_cache_to_kv_buffer( - mem_indexes=mem_indexes_cuda, - big_page_buffer_ids=big_page_buffer_ids_cuda, - page_indexes=page_indexes_cuda, - gpu_full_att_kv_state=gpu_full_att_kv_state, - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_cache_tensor=cpu_cache_tensor, - tp_rank=0, - tp_world_size=linear_config.tp_world_size, - big_page_token_num=big_page_token_num, - linear_config=linear_config, - grid_num=grid_num, +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import StatePageAllocator +from lightllm.common.attention_state.bounded_att import create_gdn_state, ConvSSMStateConfig +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import set_env_start_args + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens-per-page", type=int, default=16384) + parser.add_argument("--pages", type=int, default=1) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--warmup", type=int, default=10) + args = parser.parse_args() + if min(args.tokens_per_page, args.pages, args.iterations) <= 0 or args.warmup < 0: + parser.error("page sizes and iterations must be positive; warmup must be nonnegative") + set_env_start_args(StartArgs(cpu_cache_token_page_size=args.tokens_per_page)) + cfg = Qwen3NextResourceConfig( + tp_world_size=8, + full_att_all_num_kv_heads=2, + full_att_head_dim=256, + global_linear_k_heads=16, + global_linear_v_heads=64, + head_linear_k_dim=128, + head_linear_v_dim=128, + conv_kernel_size=4, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4, + all_layer_num=48, + ) + tokens = args.tokens_per_page * args.pages + full = GQAFullAttentionState( + GQAAttentionConfig( + logical_layer_groups=tuple((layer,) for layer in tuple(range(12))), + global_head_num=1, + head_dim=256, + dtype=torch.bfloat16, + tp_rank=0, + tp_world_size=8, + ), + ) + full.allocate_storage(tokens) + resources = AttentionStateManager(request_capacity=1) + bounded_state = create_gdn_state( + config=ConvSSMStateConfig( + logical_layer_ids=tuple( + layer for layer in range(cfg.all_layer_num) if (layer + 1) % cfg.full_attention_interval + ), + global_k_heads=cfg.global_linear_k_heads, + global_v_heads=cfg.global_linear_v_heads, + key_dim=cfg.head_linear_k_dim, + value_dim=cfg.head_linear_v_dim, + conv_width=cfg.conv_kernel_size - 1, + conv_dtype=cfg.conv_state_dtype, + ssm_dtype=cfg.ssm_state_dtype, + tp_rank=0, + tp_world_size=cfg.tp_world_size, + ), + request_capacity=1, + mtp_step=0, + device="cuda", + ) + bounded_state.initialize_page_pools(big_page_capacity=args.pages + 2, small_page_capacity=0) + resources.register_state("main", full) + resources.register_state("main", bounded_state) + resources.full_att_allocator = FullAttKVAllocator(tokens) + resources.bounded_att_big_page_allocator = StatePageAllocator(args.pages + 2, reserved_slots=2) + resources.bounded_att_cpu_load_big_page_slot = args.pages + resources.bounded_att_cpu_offload_big_page_slot = args.pages + 1 + resources.finish_state_registration() + pages = list(range(args.pages)) + slots = resources.bounded_att_big_page_allocator.alloc(args.pages) + cpu_page_buffers = torch.zeros( + (args.pages, resources.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) + # Populate every TP region through the production API before measuring restores. + full.k_storage.zero_() + full.v_storage.zero_() + resources.save_bounded_att_runtime_to_pages( + [0] * args.pages, + state_token_lens=[(page + 1) * args.tokens_per_page for page in pages], + page_kind="big", + page_slot_ids=slots, ) -torch.cuda.synchronize() - -# Timed runs -times = [] -for _ in range(BENCH_ITERS): torch.cuda.synchronize() - t0 = time.perf_counter() - copy_cpu_cache_to_kv_buffer( - mem_indexes=mem_indexes_cuda, - big_page_buffer_ids=big_page_buffer_ids_cuda, - page_indexes=page_indexes_cuda, - gpu_full_att_kv_state=gpu_full_att_kv_state, - cpu_kv_conv_state=cpu_kv_conv_state, - cpu_kv_ssm_state=cpu_kv_ssm_state, - cpu_cache_tensor=cpu_cache_tensor, - tp_rank=0, - tp_world_size=linear_config.tp_world_size, - big_page_token_num=big_page_token_num, - linear_config=linear_config, - grid_num=grid_num, + token_indices = torch.arange(tokens, dtype=torch.int64, device="cuda") + resources.get_cpu_cache_operator().write_req( + token_indices=token_indices, + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=pages, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + big_page_slot_ids=slots, ) torch.cuda.synchronize() - t1 = time.perf_counter() - times.append(t1 - t0) - -# --------------------------------------------------------------------------- -# Step 6 – report -# --------------------------------------------------------------------------- -import statistics - -times_ms = [t * 1e3 for t in times] -total_tokens = PAGE_NUM * big_page_token_num -# Calculate head_scale_size (same logic as in copy_cpu_cache_to_kv_buffer) -if linear_config.full_att_all_num_kv_heads % linear_config.tp_world_size == 0: - head_scale_size = 1 -else: - head_scale_size = linear_config.tp_world_size // linear_config.full_att_all_num_kv_heads - -# Each TP rank copies: -# - full_att_bytes / head_scale_size (full attention is sharded by head_scale_size) -# - conv_bytes / tp_world_size (conv state is sharded by tp_rank) -# - ssm_bytes / tp_world_size (ssm state is sharded by tp_rank) -full_att_bytes = linear_config.get_cpu_cache_full_att_bytes() -conv_bytes = linear_config.get_cpu_cache_conv_bytes() -ssm_bytes = linear_config.get_cpu_cache_ssm_bytes() + def restore(): + resources.get_cpu_cache_operator().read_req( + cpu_page_buffers=cpu_page_buffers, + cpu_page_indices=pages, + token_indices=token_indices, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + big_page_slot_ids=slots, + req_idx=0, + state_token_len=tokens, + ) + + for _ in range(args.warmup): + restore() + torch.cuda.synchronize() + times = [] + for _ in range(args.iterations): + start = time.perf_counter() + restore() + torch.cuda.synchronize() + times.append(time.perf_counter() - start) + token_bytes = ( + full.k_storage[:, :tokens].numel() * full.k_storage.element_size() + + full.v_storage[:, :tokens].numel() * full.v_storage.element_size() + ) + pool = ( + resources.main_kda_conv_bf16_ssm_fp32_state + if hasattr(resources, "main_kda_conv_bf16_ssm_fp32_state") + else resources.main_gdn_conv_bf16_ssm_fp32_state + ).big_pages_pool + state_bytes = args.pages * sum(t[0].numel() * t.element_size() for t in (pool.conv_state, pool.ssm_state)) + median = statistics.median(times) + print( + f"Pages: {args.pages}; tokens/page: {args.tokens_per_page}; " + f"CPU bytes/page: {resources.get_cpu_cache_operator().get_page_bytes():,}" + ) + print(f"Mean: {statistics.mean(times) * 1000:.3f} ms; median: {median * 1000:.3f} ms") + print(f"KV + state payload throughput: {(token_bytes + state_bytes) / median / 1024**3:.3f} GiB/s") -bytes_per_page_per_tp = ( - full_att_bytes - * max(1, linear_config.full_att_all_num_kv_heads // linear_config.tp_world_size) - / linear_config.full_att_all_num_kv_heads - + conv_bytes // linear_config.tp_world_size - + ssm_bytes // linear_config.tp_world_size -) -total_bytes_copied = PAGE_NUM * bytes_per_page_per_tp -print() -print("=" * 60) -print(f" copy_cpu_cache_to_kv_buffer speed benchmark") -print("=" * 60) -print(f" Pages / call : {PAGE_NUM}") -print(f" Tokens / page : {big_page_token_num}") -print(f" Total tokens / call : {total_tokens}") -print(f" Bytes / page (total) : {total_bytes:,}") -print(f" Bytes / page (per TP) : {bytes_per_page_per_tp:,}") -print(f" Total bytes / call : {total_bytes_copied:,} ({total_bytes_copied / 1024**3:.3f} GB)") -print(f" Iterations : {BENCH_ITERS}") -print(f" Mean latency : {statistics.mean(times_ms):.3f} ms") -print(f" Median latency : {statistics.median(times_ms):.3f} ms") -print(f" Std latency : {statistics.stdev(times_ms):.3f} ms") -print(f" Min latency : {min(times_ms):.3f} ms") -print(f" Max latency : {max(times_ms):.3f} ms") -print(f" Throughput (tokens/s) : {total_tokens / statistics.mean(times_ms) * 1e3:,.0f}") -print(f" Throughput (GB/s) : {total_bytes_copied / 1024**3 / statistics.mean(times_ms) * 1e3:.3f}") -print("=" * 60) +if __name__ == "__main__": + main() diff --git a/unit_tests/common/att_resources/conftest.py b/unit_tests/common/att_resources/conftest.py new file mode 100644 index 0000000000..c074e861b3 --- /dev/null +++ b/unit_tests/common/att_resources/conftest.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import json +import os +from dataclasses import asdict + +import pytest + + +@pytest.fixture(autouse=True) +def allocator_environment(kv_allocator_environment: None, monkeypatch: pytest.MonkeyPatch): + """独立资源测试也需要启动参数，以选择页池的共享分配策略。""" + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + if "LIGHTLLM_START_ARGS" not in os.environ: + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs()))) + get_env_start_args.cache_clear() + yield + # 方法缓存以 State 实例为 key；测试会反复创建实例和替换参数，结束后释放缓存引用。 + from lightllm.common.attention_state.full_att import FullAttentionState + from lightllm.common.attention_state.bounded_att import BoundedAttentionState + + state_classes = [FullAttentionState, BoundedAttentionState] + for state_class in state_classes: + state_classes.extend(state_class.__subclasses__()) + for getter in (state_class.get_cpu_transfer_regions, state_class.get_pd_transfer_regions): + if hasattr(getter, "cache_clear"): + getter.cache_clear() + get_env_start_args.cache_clear() + + +@pytest.fixture +def pd_args(monkeypatch: pytest.MonkeyPatch): + """为独立资源测试提供与生产路径一致的全局 PD 配置。""" + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv( + "LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs(pd_kv_page_size=8, cpu_cache_token_page_size=8))) + ) + get_env_start_args.cache_clear() + yield get_env_start_args() + get_env_start_args.cache_clear() + + +@pytest.fixture +def cpu_args(monkeypatch: pytest.MonkeyPatch): + """为独立资源测试提供与生产路径一致的全局 CPU 页配置。""" + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs(cpu_cache_token_page_size=8)))) + get_env_start_args.cache_clear() + yield get_env_start_args() + get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py new file mode 100644 index 0000000000..ef44a54dff --- /dev/null +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -0,0 +1,45 @@ +"""Concrete main/draft composition used to exercise manager transfer operations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from dataclasses import replace +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator +from lightllm.common.attention_state.full_att import MLAConfig, MLAFullAttentionState +from lightllm.common.attention_state.bounded_att import GDNState +from unit_tests.common.att_resources.test_full_state import full_state +from unit_tests.common.att_resources.test_bounded_state import make_state + +if TYPE_CHECKING: + import torch + + +class MainDraftTestResourceManager(AttentionStateManager): + def __init__(self, device: str | torch.device | int = "cpu"): + super().__init__(request_capacity=3) + self.register_state("main", make_state(device=device)) + draft_state = GDNState( + config=replace(self.main_kda_conv_fp32_ssm_fp32_state.config, global_v_heads=2), + request_capacity=3, + mtp_step=self.main_kda_conv_fp32_ssm_fp32_state.mtp_step, + device=device, + ) + draft_state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") + self.register_state("draft0", draft_state) + self.register_state("main", full_state(device=device)) + prepared_full_state = MLAFullAttentionState( + layout_config=MLAConfig( + logical_layer_groups=self.main_gqa_fp32_state.config.logical_layer_groups, + latent_dim=self.main_gqa_fp32_state.config.head_dim, + latent_dtype=self.main_gqa_fp32_state.config.dtype, + ), + page_size=self.main_gqa_fp32_state.page_size, + device=device, + ) + prepared_full_state.allocate_storage(self.main_gqa_fp32_state.capacity) + self.register_state("draft0", prepared_full_state) + self.full_att_allocator = FullAttKVAllocator(12, page_size=4, pin_memory=False) + self.bounded_att_big_page_allocator = StatePageAllocator(3) + self.bounded_att_small_page_allocator = StatePageAllocator(2) + self.finish_state_registration() diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py new file mode 100644 index 0000000000..10c52c8cea --- /dev/null +++ b/unit_tests/common/att_resources/model_helpers.py @@ -0,0 +1,337 @@ +"""Select explicit model fields in tests parametrized over Qwen and GLM.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from contextlib import nullcontext +from dataclasses import replace +from types import SimpleNamespace + +import torch +from lightllm.common.attention_state.types import TensorField +from lightllm.common.attention_state.bounded_att import GDNState, KDAState, WindowCompressorState +from lightllm.common.attention_state.full_att import ( + FullAttentionState, + GQAAttentionConfig, + MLAConfig, + GroupQuantizedAttentionConfig, + StaticFP8AttentionConfig, + PackedMLAAttentionConfig, + CompressedAttentionConfig, + FP8E8M0PackedMLAConfig, + StaticFP8FullAttentionState, + GQAFullAttentionState, + GroupQuantizedFullAttentionState, + MLAFullAttentionState, + PackedMLAFullAttentionState, +) + + +if TYPE_CHECKING: + from lightllm.models.qwen3next.config import Qwen3NextResourceConfig + from lightllm.common.attention_state.allocators import FullAttKVAllocator + from collections.abc import Mapping, Sequence + from lightllm.common.attention_state.manager import AttentionStateManager + from lightllm.common.basemodel.basemodel import TpPartBaseModel + + +def bind_test_fp8_calibration( + state: StaticFP8FullAttentionState, + scales: torch.Tensor, + q_scales: torch.Tensor | None = None, +) -> None: + """将测试的本地期望值展开为标定文件格式，再通过生产绑定入口初始化。 + + 每个 TP 分片放置相同测试值；这样重建任意 rank 都能保留原测试的本地 scale。 + 真实的非重复全局 head 切片由 test_model_calibration 单独覆盖。 + """ + config = state.config + calibration = {"num_layers": len(config.logical_layer_groups), "quant_type": config.scale_mode} + if config.scale_mode == "per_head": + heads = config.head_num + k, v = scales.split(heads, dim=-1) + scales = torch.cat((k.repeat(1, config.tp_world_size), v.repeat(1, config.tp_world_size)), dim=-1) + assert q_scales is not None + q_scales = q_scales.repeat(1, config.tp_world_size) + calibration["num_head"] = heads * config.tp_world_size + calibration["q_calibration"] = { + "num_head": q_scales.shape[1], + "scales_shape": list(q_scales.shape), + "scales": q_scales.tolist(), + } + calibration["scales_shape"] = list(scales.shape) + calibration["scales"] = scales.tolist() + state.bind_calibration(calibration=calibration, storage_layer_start=0) + + +def rebuild_full_state( + template: FullAttentionState, + config: ( + GQAAttentionConfig + | MLAConfig + | GroupQuantizedAttentionConfig + | StaticFP8AttentionConfig + | PackedMLAAttentionConfig + | CompressedAttentionConfig + | FP8E8M0PackedMLAConfig + | None + ) = None, + token_capacity: int | None = None, + page_size: int | None = None, + device: str | torch.device | None = None, + scales: torch.Tensor | None = None, + q_scales: torch.Tensor | None = None, +) -> FullAttentionState: + """测试中按模板重建独立存储，显式传递运行时参数与静态标定值。""" + config = template.config if config is None else config + token_capacity = template.capacity if token_capacity is None else token_capacity + page_size = template.page_size if page_size is None else page_size + device = template.device if device is None else device + if isinstance(template, StaticFP8FullAttentionState): + prepared_full_state = type(template)( + config, + page_size=page_size, + device=device, + ) + bind_test_fp8_calibration( + state=prepared_full_state, + scales=template.scales if scales is None else scales, + q_scales=template.q_scales if q_scales is None else q_scales, + ) + prepared_full_state.allocate_storage(token_capacity) + return prepared_full_state + prepared_full_state = type(template)( + config, + page_size=page_size, + device=device, + ) + prepared_full_state.allocate_storage(token_capacity) + return prepared_full_state + + +def replace_logical_layer_groups(config, logical_layer_groups): + """替换测试配置的逻辑层分组，兼容 Packed MLA 的嵌套 logical_config。""" + if isinstance(config, PackedMLAAttentionConfig): + return replace( + config, + logical_config=replace( + config.logical_config, + logical_layer_groups=logical_layer_groups, + ), + ) + return replace(config, logical_layer_groups=logical_layer_groups) + + +def owned_states(attributes: Mapping[str, object], state_class: type | tuple[type, ...]): + """从显式成员选择指定类型的 State，按资源身份排序。""" + return tuple( + sorted( + (state for state in attributes.values() if isinstance(state, state_class)), + key=lambda state: state.state_name, + ) + ) + + +def main_linear_state(resources: AttentionStateManager): + (state,) = (state for state in resources.get_bounded_attention_states("main") if hasattr(state, "conv_state")) + return state + + +def prepare_model_states(model: TpPartBaseModel, logical_layer_ids: Sequence[int] | None = None): + config = model.config.get("text_config", model.config) + if logical_layer_ids is None: + logical_layer_ids = range(config.get("num_hidden_layers", config.get("n_layer"))) + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in logical_layer_ids] + if not hasattr(model, "torch_memory_saver"): + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + model.prepare_attention_states() + + +def create_full_attention_roles( + template: FullAttentionState, + main_layer_count: int, +) -> tuple[FullAttentionState, FullAttentionState | None]: + """仅用于测试：按配置层号创建各自拥有独立存储的 main 和单个 draft。""" + layers = template.config.logical_layer_groups + boundary = sum(all(layer < main_layer_count for layer in group) for group in layers) + assert boundary and all(layer < main_layer_count for group in layers[:boundary] for layer in group) + assert all( + layer >= main_layer_count for group in layers[boundary:] for layer in group + ), "draft layers must follow main layers and cannot share main storage" + states = [] + for start, end in ((0, boundary), (boundary, len(layers))): + if start == end: + states.append(None) + continue + if isinstance(template.config, PackedMLAAttentionConfig): + config = replace( + template.config, + logical_config=replace( + template.config.logical_config, + logical_layer_groups=layers[start:end], + ), + ) + else: + config = replace(template.config, logical_layer_groups=layers[start:end]) + if isinstance(template, StaticFP8FullAttentionState): + state = rebuild_full_state( + template=template, + config=config, + scales=template.scales[start:end].clone(), + q_scales=None if template.q_scales is None else template.q_scales[start:end].clone(), + ) + else: + state = rebuild_full_state(template, config=config) + states.append(state) + return states[0], states[1] + + +def read_kv_storage(state: FullAttentionState) -> torch.Tensor | None: + """测试只读比较用：按 K、V 顺序拼接；修改数据须直接操作 State 的真实 buffer。""" + if isinstance(state, (MLAFullAttentionState, PackedMLAFullAttentionState)): + return state.storage + if state.k_storage is None: + return None + return torch.cat((state.k_storage, state.v_storage), dim=2) + + +def full_state_test_buffers(state: FullAttentionState) -> dict[str, torch.Tensor | None]: + """仅供多方案参数化测试枚举实际存储，检查容量、独立分配及字节一致性。""" + if isinstance(state, GroupQuantizedFullAttentionState): + return { + "k": state.k_storage, + "v": state.v_storage, + "k_scale": state.k_scales, + "v_scale": state.v_scales, + } + if isinstance(state, (GQAFullAttentionState, StaticFP8FullAttentionState)): + return {"k": state.k_storage, "v": state.v_storage} + assert isinstance(state, (MLAFullAttentionState, PackedMLAFullAttentionState)) + return {"kv": state.storage} + + +def full_state_test_fields( + state: FullAttentionState, + buffers: dict[str, torch.Tensor] | None = None, +) -> dict[str, TensorField]: + """测试中独立描述实际数据，用于校验不同存储方案的传输内容。""" + buffers = full_state_test_buffers(state) if buffers is None else buffers + axes = ("layer", "token", "head", "dim") + if isinstance(state, PackedMLAFullAttentionState): + storage = buffers["kv"] + raw = storage.view(torch.uint8) + if state.config.dtype == torch.float8_e4m3fn: + fields = { + "kv_nope": TensorField(tensor=raw[..., :512], axes=axes), + "kv_scale": TensorField( + tensor=raw[..., 512:528].view(torch.float32), axes=("layer", "token", "head", "group") + ), + "kv_rope": TensorField(tensor=raw[..., 528:656].view(torch.bfloat16), axes=axes), + } + else: + fields = {"kv": TensorField(tensor=storage[..., : state.config.latent_dim], axes=axes)} + start = state.config.indexer_offset + end = start + state.config.indexer_dim + fields["indexer_k"] = TensorField(tensor=raw[..., start:end], axes=axes) + fields["indexer_scale"] = TensorField( + tensor=raw[..., end : end + 4].view(torch.float32), axes=("layer", "token", "head", "scale") + ) + return fields + return { + name: TensorField(tensor=tensor, axes=("layer", "token", "head", "group") if name.endswith("_scale") else axes) + for name, tensor in buffers.items() + } + + +def build_hybrid_resources( + config: Qwen3NextResourceConfig, + request_capacity: int, + mtp_step: int, + full_state: FullAttentionState, + full_allocator: FullAttKVAllocator, + big_page_capacity: int, + small_page_capacity: int, + draft_full_state: FullAttentionState | None = None, + tp_rank: int = 0, +) -> AttentionStateManager: + """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.allocators import StatePageAllocator + from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, create_gdn_state, create_kda_state + from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state + from lightllm.models.glm5_next.config import Glm5NextResourceConfig + + manager = AttentionStateManager(request_capacity=request_capacity) + glm = isinstance(config, Glm5NextResourceConfig) + create_state = create_kda_state if glm else create_gdn_state + bounded_state = create_state( + config=ConvSSMStateConfig( + logical_layer_ids=tuple( + layer for layer in range(config.all_layer_num) if (layer + 1) % config.full_attention_interval + ), + global_k_heads=config.global_linear_k_heads, + global_v_heads=config.global_linear_v_heads, + key_dim=config.head_linear_k_dim, + value_dim=config.head_linear_v_dim, + conv_width=config.conv_kernel_size - 1, + conv_dtype=config.conv_state_dtype, + ssm_dtype=config.ssm_state_dtype, + tp_rank=tp_rank, + tp_world_size=config.tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=mtp_step, + device="cuda", + ) + bounded_state.initialize_page_pools( + big_page_capacity=big_page_capacity, + small_page_capacity=small_page_capacity, + pin_memory=True, + ) + manager.register_state("main", bounded_state) + for role, state in (("main", full_state), ("draft0", draft_full_state)): + if state is None: + continue + manager.register_state(role, state) + if glm: + window = create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=tuple(layer for group in state.config.logical_layer_groups for layer in group), + window_size=config.index_kpool, + compress_stride=config.index_kpool, + candidate_dim=config.index_head_dim, + target="indexer", + dtype=config.conv_state_dtype, + tp_rank=tp_rank, + tp_world_size=config.tp_world_size, + ), + request_capacity=request_capacity, + mtp_step=mtp_step, + device="cuda", + ) + manager.register_state(role, window) + manager.full_att_allocator = full_allocator + manager.bounded_att_big_page_allocator = StatePageAllocator(big_page_capacity, reserved_slots=2) + manager.bounded_att_small_page_allocator = StatePageAllocator(small_page_capacity) if small_page_capacity else None + manager.bounded_att_cpu_load_big_page_slot = big_page_capacity - 2 + manager.bounded_att_cpu_offload_big_page_slot = big_page_capacity - 1 + manager.finish_state_registration() + return manager + + +def window_states( + manager: AttentionStateManager, + role: str | None = None, +) -> tuple[WindowCompressorState, ...]: + return tuple( + state for state in manager.get_bounded_attention_states(role) if isinstance(state, WindowCompressorState) + ) + + +def linear_states( + manager: AttentionStateManager, + role: str | None = None, +) -> tuple[GDNState | KDAState, ...]: + return tuple( + state for state in manager.get_bounded_attention_states(role) if isinstance(state, (GDNState, KDAState)) + ) diff --git a/unit_tests/common/att_resources/test_allocators.py b/unit_tests/common/att_resources/test_allocators.py new file mode 100644 index 0000000000..3b932d8de5 --- /dev/null +++ b/unit_tests/common/att_resources/test_allocators.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import pytest +import torch + +from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator + + +def test_slot_allocation_domains_and_reserved_storage(): + big = StatePageAllocator(5, reserved_slots=2) + small = StatePageAllocator(5) + assert big.alloc(3) == [0, 1, 2] + assert small.alloc(1) == [0] + assert big.alloc(1) is None + with pytest.raises(AssertionError): + big.free([0, 3]) + assert big.free_count == 0 + big.free([0]) + assert big.free_count == 1 and small.free_count == 4 + with pytest.raises(AssertionError): + big.free([0]) + with pytest.raises(AssertionError): + big.free([1, 1]) + assert big.alloc(1) == [0] + + +def test_kv_returns_complete_pages_and_rejects_partial_free_atomically(): + allocator = FullAttKVAllocator(16, 4, pin_memory=False) + tokens = allocator.alloc(8) + snapshot = tokens.clone() + assert tokens.tolist() == list(range(8)) + assert allocator.hold_token_indices == (16, 17, 18, 19) + assert ((tokens.view(-1, 4)[:, 0] % 4) == 0).all() + assert (tokens.view(-1, 4).diff(dim=1) == 1).all() + with pytest.raises(AssertionError): + allocator.free(tokens[:2]) + assert allocator.can_use_mem_size == 8 + more = allocator.alloc(8) + assert set(tokens.tolist()).isdisjoint(more.tolist()) + assert torch.equal(tokens, snapshot) + with pytest.raises(AssertionError, match="available"): + allocator.alloc(4) + # Pages may be returned in any order; tokens within each page remain ordered. + returned = tokens.view(-1, 4).flip(0).reshape(-1) + allocator.free(returned) + assert allocator.can_use_mem_size == 8 + assert more.tolist() == list(range(8, 16)) + reused = allocator.alloc(8) + assert torch.equal(reused, returned) + allocator.free(more) + allocator.free(reused) + assert allocator.can_use_mem_size == 16 + assert set(allocator.alloc(16).tolist()) == set(range(16)) + + +def test_kv_validation_and_empty_operations(): + with pytest.raises(AssertionError, match="HOLD"): + FullAttKVAllocator((1 << 31) - 2, 3, pin_memory=False) + allocator = FullAttKVAllocator(8, 4, pin_memory=False) + for n in [-4, 3]: + with pytest.raises(AssertionError): + allocator.alloc(n) + assert allocator.alloc(0).numel() == 0 + allocator.free([]) + assert allocator.can_use_mem_size == 8 + tokens = allocator.alloc(4) + with pytest.raises(AssertionError, match="complete KV pages"): + allocator.free(tokens[:3]) + with pytest.raises(AssertionError, match="exceeds allocated"): + allocator.free(list(range(8))) + assert allocator.can_use_mem_size == 4 + allocator.free_all() + assert allocator.can_use_mem_size == 8 + assert set(allocator.alloc(8).tolist()) == set(range(8)) + + +@pytest.mark.parametrize("page_size", [1, 4, 16]) +def test_kv_repeated_allocation_and_return_buffer_wrap(page_size: int): + allocator = FullAttKVAllocator(64, page_size, pin_memory=False) + # Keep half the domain live while cycling the other half through the return buffer. + retained = allocator.alloc(32).clone() + for _ in range(16): + current = allocator.alloc(32) + snapshot = current.clone() + assert set(current.tolist()) == set(range(32, 64)) + allocator.free(current.tolist()) + assert torch.equal(current, snapshot) + assert allocator.can_use_mem_size == 32 + allocator.free(retained) + assert set(allocator.alloc(64).tolist()) == set(range(64)) diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py new file mode 100644 index 0000000000..e3a09adb8a --- /dev/null +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -0,0 +1,372 @@ +"""Exercise shared-manager startup with independently prepared hybrid drafts.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from contextlib import nullcontext +from dataclasses import asdict +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.model_infer.mode_backend import base_backend +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend +from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from pathlib import Path + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def test_backend_rejects_more_than_36_drafts_before_loading(): + backend = ModeBackend.__new__(ModeBackend) + backend.args = StartArgs(mtp_mode="vanilla_with_att", mtp_step=37) + with pytest.raises(AssertionError, match="at most 36"): + backend.init_mtp_draft_model({}) + assert backend.draft_models == [] + + +@pytest.mark.parametrize("main_linear", [False, True]) +@pytest.mark.parametrize("draft_full", [False, True]) +@pytest.mark.parametrize("draft_count", [1, 2]) +@pytest.mark.parametrize("encoding", ["None", "fp8kv_sph", "fp8kv_spt"]) +@pytest.mark.parametrize("via_backend", [False, True]) +def test_backend_allocates_all_roles_before_warmup( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + main_linear: bool, + draft_full: bool, + draft_count: int, + encoding: str, + via_backend: bool, +): + args = StartArgs( + model_dir=str(tmp_path), + data_type="bfloat16", + llm_kv_type=encoding, + page_size=4, + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + linear_att_cache_size=2, + cpu_cache_token_page_size=8, + mtp_mode="vanilla_with_att", + mtp_step=draft_count, + ) + for key in ("CURRENT_RANK_IN_DP", "CURRENT_RANK_IN_NODE", "CURRENT_DEVICE_ID"): + monkeypatch.setenv("LIGHTLLM_" + key, "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + config = dict( + model_type="qwen3_5_text", + architectures=["Qwen3NextForCausalLM"], + num_hidden_layers=4, + n_layer=4, + num_attention_heads=4, + num_key_value_heads=2, + hidden_size=64, + head_dim=16, + full_attention_interval=2, + linear_num_key_heads=2, + linear_num_value_heads=4, + linear_key_head_dim=8, + linear_value_head_dim=8, + linear_conv_kernel_dim=3, + ) + draft_config = dict( + config, + num_hidden_layers=2 if draft_full else 1, + n_layer=2 if draft_full else 1, + num_key_value_heads=1, + head_dim=8, + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + draft_dir = tmp_path / "draft" + draft_dir.mkdir() + (draft_dir / "config.json").write_text(json.dumps(draft_config)) + args.mtp_draft_model_dir = [str(draft_dir)] * draft_count + main_full_layers = 2 if main_linear else 4 + if encoding != "None": + calibration_layers = main_full_layers + int(draft_full) * draft_count + scales = torch.arange(1, calibration_layers + 1, dtype=torch.float32) * 0.125 + calibration = { + "architectures": "Qwen3NextForCausalLM", + "num_head": 1, + "num_layers": calibration_layers, + "scales_shape": [calibration_layers, 2], + "scales": scales[:, None].expand(-1, 2).tolist(), + "q_calibration": { + "num_head": 1, + "scales_shape": [calibration_layers, 1], + "scales": (scales[:, None] * 2).tolist(), + }, + "quant_type": "per_head" if encoding == "fp8kv_sph" else "per_tensor", + } + calibration_path = tmp_path / "calibration.json" + calibration_path.write_text(json.dumps(calibration)) + args.kv_quant_calibration_config_path = str(calibration_path) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(args))) + get_env_start_args.cache_clear() + get_added_mtp_kv_layer_num.cache_clear() + main_cls = Qwen3NextTpPartModel if main_linear else LlamaTpPartModel + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: main_cls) + monkeypatch.setattr("lightllm.models.get_draft_model_class", lambda *_: Qwen3NextTpPartModel) + models, warmed, padded, captured, prefill_captured, ready = [], set(), set(), set(), set(), [] + checked = [] + mapping = SimpleNamespace(bind_hold_page=lambda indices: None) + request_allocations = [] + manager_creations = [] + from lightllm.common.attention_state import AttentionStateManager + + def create_empty_manager(request_capacity: int): + manager = AttentionStateManager(request_capacity=request_capacity) + assert not manager._state_registration_finished + assert not any(name.endswith("_state") for name in manager.__dict__) + manager_creations.append(request_capacity) + return manager + + monkeypatch.setattr("lightllm.common.basemodel.basemodel.AttentionStateManager", create_empty_manager) + + def create_req_manager(max_request_num: int, max_sequence_length: int): + assert len(models) == 1 + draft_count + assert all(not hasattr(item, "req_manager") for item in models) + request_allocations.append((max_request_num, max_sequence_length)) + return mapping + + monkeypatch.setattr("lightllm.common.req_manager.ReqManager", create_req_manager) + + def mark_backends(model: TpPartBaseModel): + # 第一个模型初始化 backend 前，全部 State 已创建完成，可通过缓存查询取得资源。 + for item in models: + for state in item.att_state_manager.get_full_attention_states(item.att_state_role): + logical_layer_ids = {layer_id for group in state.config.logical_layer_groups for layer_id in group} + for layer in item.layers_infer: + if layer.layer_num_ in logical_layer_ids: + assert ( + item.att_state_manager.get_state( + model_role=item.att_state_role, + state_type_filter="", + logical_layer_id=layer.layer_num_, + state_class=type(state), + ) + is state + ) + model.prefill_att_backend = model.decode_att_backend = object() + model.prefill_att_backend1 = model.decode_att_backend1 = None + model.backends_bound = True + + def warmup(model: TpPartBaseModel): + assert model.backends_bound + assert all(item.req_manager is mapping for item in models) + assert all(item.att_state_manager is models[0].att_state_manager for item in models) + assert model.att_state_manager.full_att_allocator.size == 32 + assert all(item.max_total_token_num == 32 for item in models[: models.index(model) + 1]) + assert all(item.max_total_token_num is None for item in models[models.index(model) + 1 :]) + assert captured == prefill_captured == {id(item) for item in models[: models.index(model)]} + assert not ready + warmed.add(id(model)) + + def capture(model: TpPartBaseModel): + assert warmed == padded == {id(item) for item in models[: models.index(model) + 1]} + assert not ready + captured.add(id(model)) + + def check_max_len(model: TpPartBaseModel): + assert model is models[0] + assert captured == prefill_captured == {id(model)} + assert not ready + checked.append(model) + + offset = 0 + for index in range(1 + draft_count): + cls = Qwen3NextTpPartModel if index else main_cls + model = cls.__new__(cls) + model.is_mtp_draft_model = bool(index) + model.args, model.config = args, draft_config.copy() if index else config.copy() + model.max_req_num, model.max_total_token_num, model.mem_fraction = 2, None, 0.9 + model.max_seq_length, model.batch_max_tokens = (32 if index else 24), 16 + model.tp_world_size_, model.tp_rank_ = 1, 0 + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(offset, offset + model.config["n_layer"])] + offset += len(model.layers_infer) + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + if index: + model.main_model, model.mtp_previous_draft_models = models[0], models[1:].copy() + model._init_att_backend = lambda model=model: mark_backends(model) + model._init_att_backend1 = lambda: None + model._init_hidden_collector = lambda: None + model._autotune_warmup = lambda model=model: warmup(model) + model._init_padded_req = lambda model=model: padded.add(id(model)) + model._init_cudagraph = lambda model=model: capture(model) + model._init_prefill_cuda_graph = lambda model=model: prefill_captured.add(id(model)) + model._check_max_len_infer = lambda model=model: check_max_len(model) + models.append(model) + + def memory_info(device: str | torch.device | int): + assert all(item.req_manager is mapping for item in models) + manager = models[0].att_state_manager + states = manager.get_full_attention_states() + assert all(read_kv_storage(state) is None for state in states) + assert manager.draft0_gdn_conv_bf16_ssm_fp32_state.conv_state.is_cuda + assert manager.draft0_gdn_conv_bf16_ssm_fp32_state.big_pages_pool is None + assert manager.full_att_allocator is None + assert not warmed + return manager.get_full_att_token_bytes() * (32 + args.page_size), 0 + + monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) + monkeypatch.setattr("lightllm.utils.envs_utils.set_model_init_status", ready.append) + assert all(not hasattr(item, "att_state_manager") for item in models) + + class RuntimeReady(Exception): + """Stop after model readiness, before unrelated queue and thread setup.""" + + if via_backend: + backend = ModeBackend.__new__(ModeBackend) + backend.init_rank_infos = lambda: None + backend.init_mtp_draft_model = lambda _: backend.draft_models.extend(models[1:]) + monkeypatch.setattr(base_backend, "init_distributed_env", lambda _: None) + monkeypatch.setattr(base_backend.dist_group_manager, "create_groups", lambda **_: None) + monkeypatch.setattr(base_backend, "TokenLoad", lambda *_: None) + monkeypatch.setattr(base_backend, "get_model", lambda *_: (models[0], False)) + + def mark_ready(value: bool): + ready.append(value) + raise RuntimeReady + + monkeypatch.setattr("lightllm.utils.envs_utils.set_model_init_status", mark_ready) + + def start_runtime(): + if via_backend: + with pytest.raises(RuntimeReady): + backend.init_model( + dict( + args=args, + world_size=1, + load_way="HF", + weight_dir=str(tmp_path), + max_total_token_num=None, + ) + ) + else: + assert all( + getattr(item, "att_state_manager", None) is None for item in models + ), "resources already initialized" + req_manager = create_req_manager( + max_request_num=models[0].max_req_num, + max_sequence_length=max(item.get_req_manager_max_seq_len() for item in models), + ) + for item in models: + item.bind_req_manager(req_manager=req_manager) + for index, item in enumerate(models): + item.prepare_attention_states() + assert all(other.att_state_manager is models[0].att_state_manager for other in models[: index + 1]) + assert all(not hasattr(other, "att_state_manager") for other in models[index + 1 :]) + main = models[0] + main.att_state_manager.finish_state_registration() + main.att_state_manager.load_full_att_calibration(args=args) + main.att_state_manager.initialize_storage( + args, token_capacity=main.max_total_token_num, mem_fraction=main.mem_fraction + ) + assert all(item.max_total_token_num is None for item in models) + for index, item in enumerate(models): + item.init_runtime() + assert all(not getattr(other, "backends_bound", False) for other in models[index + 1 :]) + assert warmed == padded == captured == prefill_captured == {id(other) for other in models[: index + 1]} + + start_runtime() + assert request_allocations == [(2, 32)] + assert manager_creations == [models[0].max_req_num + 1] + assert captured == {id(model) for model in models} + assert checked == [models[0]] + assert ready == ([True] if via_backend else []) + manager = models[0].att_state_manager + assert len(manager.get_bounded_attention_states()) == draft_count + int(main_linear) + assert all(not hasattr(model, "mtp_state_cursors") for model in models) + assert all(not hasattr(model, "full_att_state") for model in models) + draft_states = [manager.get_bounded_attention_states(f"draft{i}")[0] for i in range(draft_count)] + assert len({state.conv_state.data_ptr() for state in draft_states}) == draft_count + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([draft_count + 1], dtype=torch.int32, device="cuda"), + role="main", + ) + assert all(not state.mtp_state_index.any() for state in draft_states) + if main_linear: + assert manager.main_gdn_conv_bf16_ssm_fp32_state.mtp_state_index[0] == draft_count + for i, state in enumerate(draft_states): + assert len(state.config.logical_layer_ids) == 1 + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([i + 2], dtype=torch.int32, device="cuda"), + role=f"draft{i}", + ) + assert state.mtp_state_index[0] == i + 1 + assert all(not other.mtp_state_index.any() for other in draft_states[i + 1 :]) + if draft_full: + main_state = manager.get_full_attention_states("main")[0] + draft_state = manager.get_full_attention_states(f"draft{i}")[0] + assert main_state.config.head_dim == 16 + assert draft_state.config.head_dim == 8 + assert {buffer.data_ptr() for buffer in full_state_test_buffers(main_state).values()}.isdisjoint( + buffer.data_ptr() for buffer in full_state_test_buffers(draft_state).values() + ) + assert len(draft_state.config.logical_layer_groups) == 1 + else: + assert not manager.get_full_attention_states(f"draft{i}") + + if encoding != "None": + ranges = [("main", 0, main_full_layers)] + [ + (f"draft{i}", main_full_layers + i, main_full_layers + i + 1) for i in range(draft_count) + ] + for role, start, end in ranges: + for state in manager.get_full_attention_states(role): + assert torch.equal(state.scales[:, 0].cpu(), scales[start:end]) + if encoding == "fp8kv_sph": + assert torch.equal(state.q_scales[:, 0].cpu(), scales[start:end] * 2) + + layouts = build_cpu_cache_layouts() + assert ( + max(layout.byte_end for layout in layouts.values()) + 15 + ) // 16 * 16 == manager.get_cpu_cache_operator().get_page_bytes() + assert layouts.keys() == manager.get_cpu_cache_operator().get_page_layouts().keys() + for state_name, layout in layouts.items(): + runtime_layout = manager.get_cpu_cache_operator().get_page_layouts()[state_name] + assert [(r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in layout.regions] == [ + (r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in runtime_layout.regions + ] + slot = manager.bounded_att_big_page_allocator.alloc(1)[0] + for i, state in enumerate(manager.get_bounded_attention_states()): + state.conv_state.fill_(i + 3) + state.ssm_state.fill_(i + 7) + manager.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[4], + page_kind="big", + page_slot_ids=[slot], + ) + manager.reset_bounded_att_state([0]) + assert all(not state.conv_state[:, 0].any() for state in manager.get_bounded_attention_states()) + manager.restore_bounded_att_runtime_from_pages( + req_indices=[0], + state_token_lens=[4], + page_kind="big", + page_slot_ids=[slot], + ) + for i, state in enumerate(manager.get_bounded_attention_states()): + assert (state.conv_state[:, 0, :, : state.config.conv_width] == i + 3).all() + assert (state.ssm_state[:, 0] == i + 7).all() + with pytest.raises(AssertionError, match="resources already initialized"): + start_runtime() + assert all(item.att_state_manager is manager for item in models) + get_env_start_args.cache_clear() + get_added_mtp_kv_layer_num.cache_clear() diff --git a/unit_tests/common/att_resources/test_bounded_copy_kernel.py b/unit_tests/common/att_resources/test_bounded_copy_kernel.py new file mode 100644 index 0000000000..22d7ab470a --- /dev/null +++ b/unit_tests/common/att_resources/test_bounded_copy_kernel.py @@ -0,0 +1,287 @@ +"""聚合复制必须保持位模式、stride、流顺序和未参与区域。""" + +import pytest +import torch +import triton.language as tl + +from lightllm.common.attention_state.bounded_att.copy_kernel import copy_bounded_att_tensors +from lightllm.common.attention_state.bounded_att.utils import register_page_pool_tensor + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("direction", ["cuda_to_host", "host_to_cuda", "host_to_host", "cuda_to_cuda"]) +@pytest.mark.parametrize("grid_num", [1, 5, 16]) +@pytest.mark.parametrize("dtype", [torch.int8, torch.bfloat16, torch.float32, torch.float16, torch.int64]) +def test_aggregate_same_dtype_and_strides(direction: str, grid_num: int, dtype: torch.dtype) -> None: + src_device = "cuda" if direction.startswith("cuda") else "cpu" + dst_device = "cuda" if direction.endswith("cuda") else "cpu" + sources, targets, copies, expected_targets = [], [], [], [] + # Conv 窗口裁剪和 SSM 轴交换均保留间隙；使用非 2 次幂的边界覆盖最后一个 block。 + # 同批视图保持相同 shape，验证 grid 跳步和跨视图 block 分配。 + for _ in range(41): + shape = (3, 5, 7, 131) + source = torch.empty( + size=shape, + dtype=dtype, + device=src_device, + pin_memory=src_device == "cpu", + ) + source.copy_(torch.arange(source.numel()).reshape(shape) % 127) + target_shape = (shape[1], shape[0], shape[2] + 2, shape[3] + 4) + target = torch.full( + size=target_shape, + fill_value=-17, + dtype=dtype, + device=dst_device, + pin_memory=dst_device == "cpu", + ) + source_view = source[..., : shape[3] - 1] + target_view = target.permute(1, 0, 2, 3)[:, :, 1 : shape[2] + 1, 2 : shape[3] + 1] + expected = target.cpu().clone() + expected.permute(1, 0, 2, 3)[:, :, 1 : shape[2] + 1, 2 : shape[3] + 1].copy_(source_view.cpu()) + sources.append(source) + targets.append(target) + copies.append((source_view, target_view)) + expected_targets.append(expected) + copy_bounded_att_tensors(copies=copies, grid_num=grid_num) + torch.cuda.current_stream().synchronize() + for target, expected in zip(targets, expected_targets): + assert torch.equal(target.cpu(), expected) + + +def test_registered_shared_pages_and_stream_dependency() -> None: + source = torch.empty(size=(4099,), dtype=torch.float32).share_memory_() + target = torch.empty(size=(4099,), dtype=torch.float32).share_memory_() + register_page_pool_tensor(tensor=source) + register_page_pool_tensor(tensor=target) + producer = torch.cuda.Stream() + consumer = torch.cuda.Stream() + values = torch.arange(source.numel(), dtype=source.dtype, device="cuda") + with torch.cuda.stream(producer): + copy_bounded_att_tensors(copies=[(values, source)]) + with torch.cuda.stream(consumer): + consumer.wait_stream(producer) + # 这里不经过主机 synchronize；CPU→CPU 传输必须在 GPU stream 上等待上游。 + copy_bounded_att_tensors(copies=[(source, target)]) + restored = torch.empty_like(values) + copy_bounded_att_tensors(copies=[(target, restored)]) + consumer.synchronize() + assert torch.equal(restored, values) + assert torch.equal(target, values.cpu()) + + +@pytest.mark.parametrize( + "float_dtype,int_dtype,patterns", + [ + (torch.float16, torch.int16, [0x7E01, -32768, 0x7C00, 1]), + (torch.bfloat16, torch.int16, [0x7FC1, -32768, 0x7F80, 1]), + (torch.float32, torch.int32, [0x7FC00001, -2147483648, 0x7F800000, 1]), + (torch.float64, torch.int64, [0x7FF8000000000001, -9223372036854775808, 0x7FF0000000000000, 1]), + ], +) +def test_exact_float_bit_patterns(float_dtype: torch.dtype, int_dtype: torch.dtype, patterns: list[int]) -> None: + # 包含 NaN payload、负零和无穷，等宽整数复制不能触发精度转换或规范化。 + bits = torch.tensor(data=patterns, dtype=int_dtype, device="cuda") + target = torch.empty(size=(4,), dtype=float_dtype, pin_memory=True) + copy_bounded_att_tensors(copies=[(bits.view(float_dtype), target)]) + torch.cuda.synchronize() + assert torch.equal(target.view(int_dtype), bits.cpu()) + + +def test_one_kernel_for_same_dimension_batch(monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.attention_state.bounded_att import copy_kernel + + actual_kernel = copy_kernel._copy_bounded_att_1d + launches = [] + + class CountLaunches: + def __getitem__(self, grid: tuple[int]): + launches.append(grid) + return actual_kernel[grid] + + monkeypatch.setattr(copy_kernel, "_copy_bounded_att_1d", CountLaunches()) + source = torch.ones(size=(32, 67), device="cuda") + target = torch.zeros_like(source) + copy_bounded_att_tensors(copies=[(source[i], target[i]) for i in range(32)]) + assert launches == [(24,)] + assert torch.equal(source, target) + copy_bounded_att_tensors(copies=[]) + with pytest.raises(AssertionError, match="copy tensors must be nonempty"): + copy_bounded_att_tensors(copies=[(source[:0], target[:0])]) + assert len(launches) == 1 + + +@pytest.mark.parametrize("ndim", [1, 2, 3, 4, 5, 6]) +def test_up_to_six_dimensions(ndim: int) -> None: + # 非整 block 的大小，覆盖每个维度算子的跳步循环与尾部 mask。 + shape = (2,) * (ndim - 1) + (2051,) + source = torch.arange(2 ** (ndim - 1) * 2051, dtype=torch.float32, device="cuda") + source = source.reshape(shape) + if ndim >= 2: + source = source.transpose(0, ndim - 1) + target = torch.full( + size=source.shape, + fill_value=-1, + dtype=source.dtype, + pin_memory=True, + ) + copy_bounded_att_tensors(copies=[(source, target)]) + torch.cuda.synchronize() + assert torch.equal(source.cpu(), target) + + +@pytest.mark.parametrize("ndim", range(1, 7)) +@pytest.mark.parametrize("num_stages", [1, 2, 4, 8]) +def test_same_shape_batch_uses_one_kernel( + ndim: int, + num_stages: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from lightllm.common.attention_state.bounded_att import copy_kernel + + launches = [] + kernel_name = f"_copy_bounded_att_{ndim}d" + actual_kernel = getattr(copy_kernel, kernel_name) + + class CountLaunches: + def __getitem__(self, grid: tuple[int]): + launches.append(grid) + return actual_kernel[grid] + + monkeypatch.setattr(copy_kernel, kernel_name, CountLaunches()) + copies, targets, expected_targets = [], [], [] + for _ in range(3): + shape = (2,) * (ndim - 1) + (517,) + storage_shape = (4,) * (ndim - 1) + (1034,) + source = torch.arange(torch.Size(storage_shape).numel(), dtype=torch.float32, device="cuda").reshape( + storage_shape + ) + # 每个轴都保留间隙，保证降维后仍需调用对应的 1～6 维算子。 + source_view = source[(slice(None, None, 2),) * ndim] + target = torch.full(size=(*shape[:-1], 1034), fill_value=-17, dtype=source.dtype, pin_memory=True) + target_view = target[..., ::2] + copies.append((source_view, target_view)) + expected = target.clone() + expected[..., ::2].copy_(source_view.cpu()) + targets.append(target) + expected_targets.append(expected) + + copy_bounded_att_tensors(copies=copies, grid_num=5, num_stages=num_stages) + torch.cuda.synchronize() + assert launches == [(5,)] + for target, expected in zip(targets, expected_targets): + assert torch.equal(target, expected) + + +@pytest.mark.parametrize( + "layout", + [ + "ssm", + "conv_mtp", + "dense", + "width4", + "middle_unaligned", + "prefix", + "singletons", + "wide_singleton", + "wide_singleton_prefix", + "all_singletons", + ], +) +def test_transfer_views_preserve_gaps_and_alignment(layout: str, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.attention_state.bounded_att import copy_kernel + + launches = [] + + class CaptureLaunch: + def __init__(self, kernel, ndim: int) -> None: + self.kernel = kernel + self.ndim = ndim + + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs): + launches.append((self.ndim, kwargs["SHAPE"], kwargs["DTYPE"])) + return self.kernel[grid](**kwargs) + + return launch + + for ndim in range(1, 7): + name = f"_copy_bounded_att_{ndim}d" + monkeypatch.setattr(copy_kernel, name, CaptureLaunch(getattr(copy_kernel, name), ndim)) + + if layout == "ssm": + source = torch.arange(2 * 4 * 3 * 128, dtype=torch.float32, device="cuda").view(2, 4, 3, 128) + target = torch.full(size=(4, 2, 3, 128), fill_value=-1, dtype=source.dtype, pin_memory=True) + copies = [(source[:, index], target[index]) for index in (0, 2)] + expected = target.clone() + for index in (0, 2): + expected[index].copy_(source[:, index].cpu()) + expected_launch = (2, (2, 192), tl.uint64) + elif layout == "conv_mtp": + source = torch.arange(2 * 4 * 16 * 7, dtype=torch.float32, device="cuda").to(torch.bfloat16).view(2, 4, 16, 7) + target = torch.full(size=(4, 2, 16, 4), fill_value=-1, dtype=source.dtype, pin_memory=True) + copies = [(source[:, index, :, :4], target[index]) for index in (0, 2)] + expected = target.clone() + for index in (0, 2): + expected[index].copy_(source[:, index, :, :4].cpu()) + expected_launch = (3, (2, 16, 4), tl.uint16) + elif layout in ("prefix", "singletons", "wide_singleton", "wide_singleton_prefix", "all_singletons"): + if layout == "prefix": + shape = (2, 3, 4, 7) + target_shape = (2, 3, 4, 9) + width = 5 + expected_launch = (4, (2, 3, 4, 5), tl.uint32) + elif layout == "singletons": + shape = (1, 2, 1, 3, 1, 7) + target_shape = (1, 2, 1, 3, 1, 9) + width = 5 + # 尾部单元素轴可并入 width；前面的 padding 边界仍然保留。 + expected_launch = (5, (1, 2, 1, 3, 5), tl.uint32) + elif layout == "wide_singleton_prefix": + # dtype 扩宽只改变末轴长度，前后的单元素轴均保持不变。 + shape = (1, 3, 4) + target_shape = (1, 3, 6) + width = 2 + expected_launch = (3, (1, 3, 1), tl.uint64) + elif layout == "wide_singleton": + shape = (3, 4) + target_shape = (3, 6) + width = 2 + expected_launch = (2, (3, 1), tl.uint64) + else: + shape = target_shape = (1, 1, 1) + width = 1 + expected_launch = (1, (1,), tl.uint32) + source = torch.arange(torch.Size(shape).numel(), dtype=torch.float32, device="cuda").view(shape) + target = torch.full(size=target_shape, fill_value=-1, dtype=source.dtype, pin_memory=True) + copies = [(source[..., :width], target[..., :width])] + expected = target.clone() + expected[..., :width].copy_(source[..., :width].cpu()) + elif layout == "middle_unaligned": + source = torch.arange(64, dtype=torch.bfloat16, device="cuda") + target = torch.full(size=(64,), fill_value=-1, dtype=source.dtype, pin_memory=True) + expected = target.clone() + copies = [] + # 首尾地址都满足 8 字节对齐，中间请求偏移不对齐；整批必须保留 2 字节搬运。 + for src_offset, dst_offset in ((0, 0), (9, 16), (24, 32)): + copies.append((source[src_offset : src_offset + 8], target[dst_offset : dst_offset + 8])) + expected[dst_offset : dst_offset + 8].copy_(source[src_offset : src_offset + 8].cpu()) + expected_launch = (1, (8,), tl.uint16) + else: + shape = (4, 2, 3, 128) if layout == "dense" else (4, 6) + source = ( + torch.arange(torch.Size(shape).numel(), dtype=torch.float32, device="cuda").to(torch.bfloat16).view(shape) + ) + target = torch.full(size=shape, fill_value=-1, dtype=source.dtype, pin_memory=True) + copies = [(source[index], target[index]) for index in (0, 2)] + expected = target.clone() + for index in (0, 2): + expected[index].copy_(source[index].cpu()) + expected_launch = (1, (192,), tl.uint64) if layout == "dense" else (1, (3,), tl.uint32) + + copy_bounded_att_tensors(copies=copies) + torch.cuda.synchronize() + assert torch.equal(target, expected) + assert launches == [expected_launch] diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py new file mode 100644 index 0000000000..ffd59db4aa --- /dev/null +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -0,0 +1,436 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ( + GDNPagesPool, + GDNState, + KDAPagesPool, + KDAState, + ConvSSMStateConfig, +) +from lightllm.common.attention_state.types import ( + AxisMapping, + StateTransferLayout, + TransferRegion, +) + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for aggregated copies") + + +def make_state( + cls: type[KDAState] | type[GDNState] = KDAState, + device: str | torch.device | int = "cpu", + heads: int = 1, + mtp_step: int = 2, +): + config = ConvSSMStateConfig( + logical_layer_ids=(1, 3), + global_k_heads=heads, + global_v_heads=heads, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + state = cls( + config=config, + request_capacity=3, + mtp_step=mtp_step, + device=device, + ) + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", state) + state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device != "cpu") + return state + + +def page_layout(state: KDAState | GDNState): + cfg = state.config + conv_shape = (cfg.conv_channels, len(cfg.logical_layer_ids), cfg.conv_width) + conv_bytes = cfg.conv_channels * len(cfg.logical_layer_ids) * cfg.conv_width * 4 + ssm_shape = (len(cfg.logical_layer_ids), cfg.num_v_heads, cfg.key_dim, cfg.value_dim) + ssm_bytes = len(cfg.logical_layer_ids) * cfg.num_v_heads * cfg.key_dim * cfg.value_dim * 4 + end = 16 + conv_bytes + ssm_bytes + return StateTransferLayout( + state_name=state.state_name, + tokens_per_page=None, + byte_start=0, + byte_end=end, + regions=( + TransferRegion( + field="conv", + byte_start=16, + byte_end=16 + conv_bytes, + dtype=torch.float32, + axes=("channel", "layer", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=0, + external_start=0, + length=cfg.conv_channels, + ), + AxisMapping( + axis="window", + local_start=0, + external_start=0, + length=cfg.conv_width, + ), + ), + ), + TransferRegion( + field="ssm", + byte_start=16 + conv_bytes, + byte_end=end, + dtype=torch.float32, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=0, + length=cfg.num_v_heads, + ), + AxisMapping( + axis="key_dim", + local_start=0, + external_start=0, + length=cfg.key_dim, + ), + AxisMapping( + axis="value_dim", + local_start=0, + external_start=0, + length=cfg.value_dim, + ), + ), + ), + ), + ) + + +@pytest.mark.parametrize("cls,pool_cls", [(KDAState, KDAPagesPool), (GDNState, GDNPagesPool)]) +@pytest.mark.parametrize("device", ["cuda"]) +@pytest.mark.parametrize("page_kind", ["big", "small"]) +def test_prefill_state_is_saved_and_restored( + cls: type[KDAState] | type[GDNState], + pool_cls: type[KDAPagesPool] | type[GDNPagesPool], + device: str | torch.device | int, + page_kind: PageKind, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + config = ConvSSMStateConfig( + logical_layer_ids=(1, 3), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + state = cls( + config=config, + request_capacity=3, + mtp_step=2, + device=device, + ) + assert state.state_name is None and not hasattr(config, "state_name") + assert state.big_pages_pool is None and state.small_pages_pool is None + conv, ssm = state.get_layer_view(3) + assert conv.data_ptr() == state.conv_state[1].data_ptr() + assert ssm.data_ptr() == state.ssm_state[1].data_ptr() + for j in range(4): + state.conv_state[:, 0, :, j].fill_(j + 10) + for j in range(3): + state.ssm_state[:, j].fill_(j + 20) + state.mtp_state_index[2].fill_(1) + # 延后创建页池不能重建运行态或重置已接受的 MTP 游标。 + state.initialize_page_pools(big_page_capacity=3, small_page_capacity=2, pin_memory=device == "cuda") + assert conv.data_ptr() == state.conv_state[1].data_ptr() + assert ssm.data_ptr() == state.ssm_state[1].data_ptr() + assert type(state.big_pages_pool) is pool_cls + assert type(state.small_pages_pool) is pool_cls + assert state.big_pages_pool is not state.small_pages_pool + assert state.mtp_state_index.tolist() == [0, 0, 1] + state.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=page_kind, + page_slot_ids=[1], + ) + if device == "cuda": + torch.cuda.synchronize() + pool = state.big_pages_pool if page_kind == "big" else state.small_pages_pool + assert torch.equal(pool.conv_state[1, ..., 0], torch.full_like(pool.conv_state[1, ..., 0], 10)) + assert torch.equal(pool.conv_state[1, ..., 1], torch.full_like(pool.conv_state[1, ..., 1], 11)) + assert (pool.ssm_state[1] == 20).all() + state.mtp_state_index[1] = 2 + state.restore_bounded_att_runtime_from_pages( + req_indices=[1], + state_token_lens=[8], + page_kind=page_kind, + page_slot_ids=[1], + ) + assert state.mtp_state_index[1].item() == 0 + assert torch.equal(state.conv_state[:, 1, :, :2].cpu(), pool.conv_state[1]) + assert torch.equal(state.ssm_state[:, 3].cpu(), pool.ssm_state[1]) + + +def test_gdn_dimensions_and_page_restore(): + config = replace( + make_state().config, + global_k_heads=1, + global_v_heads=2, + key_dim=2, + value_dim=3, + ) + state = GDNState( + config=config, + request_capacity=3, + mtp_step=2, + device="cuda", + ) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2, pin_memory=True) + conv = torch.arange(40, dtype=config.conv_dtype).reshape(2, 10, 2) + ssm = torch.arange(24, dtype=config.ssm_dtype).reshape( + shape=(2, 2, 2, 3), + ) + state.conv_state[:, 0, :, :2].copy_(conv) + state.ssm_state[:, 0].copy_(ssm) + for page_kind in ("big", "small"): + state.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=page_kind, + page_slot_ids=[0], + ) + state.reset_bounded_att_state([0]) + for req, page_kind in enumerate(("big", "small"), start=1): + state.restore_bounded_att_runtime_from_pages( + req_indices=[req], + state_token_lens=[8], + page_kind=page_kind, + page_slot_ids=[0], + ) + assert torch.equal(state.conv_state[:, req, :, :2].cpu(), conv) + assert torch.equal(state.ssm_state[:, req * 3].cpu(), ssm) + with pytest.raises(AssertionError, match="equal Q/K/V"): + KDAPagesPool(2, config, pin_memory=False) + + +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): + state = make_state(cls, device="cuda") + state.conv_state[:, 0].fill_(7) + state.ssm_state[:, 0:3].fill_(11) + state.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[6], + page_kind="small", + page_slot_ids=[0], + ) + # 推进候选区和游标后，重置应清零全部运行态，同时保留已保存的页快照。 + state.conv_state[:, 0, :, 2:4].fill_(19) + state.ssm_state[:, 2].fill_(23) + state.mtp_state_index[0].fill_(2) + state.reset_bounded_att_state([0]) + assert not state.conv_state[:, 0].any() and not state.ssm_state[:, :3].any() + assert state.mtp_state_index[0] == 0 + assert (state.small_pages_pool.conv_state[0] == 7).all() + assert (state.small_pages_pool.ssm_state[0] == 11).all() + with pytest.raises(AssertionError, match="CUDA tensors"): + state.commit_bounded_att_state(torch.tensor([0]), torch.tensor([1], dtype=torch.int32)) + + +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +def test_cpu_pages_roundtrip_offsets(cls: type[KDAState] | type[GDNState]): + source = make_state(cls, device="cuda") + source.conv_state[:, 0].fill_(13) + source.ssm_state[:, 0].fill_(29) + source.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[0], + ) + layout = page_layout(source) + buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8, pin_memory=True) + source.write_state_pages_to_cpu( + page_kind="big", + page_slot_ids=[0], + cpu_page_buffers=buffer, + cpu_page_indices=[2], + layout=layout, + ) + torch.cuda.synchronize() + assert (buffer[:2] == 255).all() and (buffer[2, :16] == 255).all() + target = make_state(cls, device="cuda") + target.read_state_pages_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[2], + big_page_slot_ids=[1], + layout=layout, + ) + target.restore_bounded_att_runtime_from_pages( + req_indices=[2], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[1], + ) + assert (target.conv_state[:, 2, :, :2] == 13).all() + assert (target.ssm_state[:, 6] == 29).all() + # 页池只保存快照数据；前缀关联由调用方维护，读回的快照可以再次写出。 + target.read_state_pages_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[2], + big_page_slot_ids=[0], + layout=layout, + ) + target.write_state_pages_to_cpu( + page_kind="big", + page_slot_ids=[0], + cpu_page_buffers=buffer, + cpu_page_indices=[0], + layout=layout, + ) + torch.cuda.synchronize() + assert torch.equal(buffer[0], buffer[2]) + + +@pytest.mark.parametrize("device", ["cuda"]) +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +def test_pd_roundtrip_uses_baseline_state_and_page_relative_offset( + cls: type[KDAState] | type[GDNState], + device: str | torch.device | int, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source = make_state(cls, device=device) + source.conv_state.fill_(-1) + source.ssm_state.fill_(-1) + source.conv_state[:, 0, :, :2].fill_(23) + source.ssm_state[:, 0].fill_(47) + layout = page_layout(source) + # PD 只传 Conv/SSM，仍保留非零区域起点以覆盖组合页偏移。 + allocation = torch.full( + size=(3, layout.byte_end), + fill_value=254, + dtype=torch.uint8, + device=device, + ) + page = allocation[1] + source.write_to_pd_page( + req_idx=0, + pd_page_buffers=allocation, + pd_page_index=1, + state_token_len=8, + layout=layout, + ) + target = make_state(cls, device=device, mtp_step=1) + target.read_from_pd_page( + pd_page_buffers=allocation, + req_idx=2, + pd_page_index=1, + state_token_len=8, + layout=layout, + ) + assert (target.conv_state[:, 2, :, :2] == 23).all() + assert (target.ssm_state[:, 4] == 47).all() + assert (allocation[0] == 254).all() and (allocation[2] == 254).all() + assert (page[:16] == 254).all() + assert target.mtp_state_index[2] == 0 + + +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +@pytest.mark.parametrize("device", ["cuda"]) +@pytest.mark.parametrize("src_mtp_step,dst_mtp_step", [(0, 2), (2, 1)]) +def test_dp_runtime_copy_preserves_pages_and_other_requests(cls, device, src_mtp_step, dst_mtp_step): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source = make_state(cls, device=device, mtp_step=src_mtp_step) + target = make_state(cls, device=device, mtp_step=dst_mtp_step) + source.conv_state.fill_(-7) + source.ssm_state.fill_(-9) + target.conv_state.fill_(-1) + target.ssm_state.fill_(-2) + target.mtp_state_index.fill_(1) + target.big_pages_pool.conv_state.fill_(91) + target.big_pages_pool.ssm_state.fill_(92) + for req in (0, 2): + source.conv_state[:, req, :, : source.config.conv_width].fill_(10 + req) + source.ssm_state[:, req * (src_mtp_step + 1)].fill_(20 + req) + src_conv = source.conv_state.clone() + src_ssm = source.ssm_state.clone() + + target.copy_runtime_state_from_dp( + src_req_indices=[2, 0], + dst_req_indices=[0, 1], + remote_dp_bounded_state=source, + ) + + for src_req, dst_req in ((2, 0), (0, 1)): + assert (target.conv_state[:, dst_req, :, :2] == 10 + src_req).all() + assert (target.conv_state[:, dst_req, :, 2:] == -1).all() + dst_slot = dst_req * (dst_mtp_step + 1) + assert (target.ssm_state[:, dst_slot] == 20 + src_req).all() + assert (target.ssm_state[:, dst_slot + 1 : dst_slot + dst_mtp_step + 1] == -2).all() + assert target.mtp_state_index[dst_req] == 0 + assert (target.conv_state[:, 2] == -1).all() + assert (target.ssm_state[:, 2 * (dst_mtp_step + 1) :] == -2).all() + assert target.mtp_state_index[2] == 1 + assert (target.big_pages_pool.conv_state == 91).all() + assert (target.big_pages_pool.ssm_state == 92).all() + assert torch.equal(source.conv_state, src_conv) + assert torch.equal(source.ssm_state, src_ssm) + + +@pytest.mark.parametrize("cls", [KDAState, GDNState]) +def test_dp_big_page_copy_preserves_runtime_and_other_pages(cls): + source = make_state(cls, device="cuda") + target = make_state(cls, device="cuda") + for page in range(3): + source.big_pages_pool.conv_state[page].fill_(10 + page) + source.big_pages_pool.ssm_state[page].fill_(20 + page) + target.big_pages_pool.conv_state.fill_(-1) + target.big_pages_pool.ssm_state.fill_(-2) + target.small_pages_pool.conv_state.fill_(31) + target.small_pages_pool.ssm_state.fill_(32) + target.conv_state.fill_(41) + target.ssm_state.fill_(42) + target.mtp_state_index.fill_(1) + src_conv = source.big_pages_pool.conv_state.clone() + src_ssm = source.big_pages_pool.ssm_state.clone() + + target.copy_big_pages_from_dp( + src_page_ids=[2, 0], + dst_page_ids=[0, 1], + remote_dp_bounded_state=source, + ) + + torch.cuda.synchronize() + for src_page, dst_page in ((2, 0), (0, 1)): + assert torch.equal(target.big_pages_pool.conv_state[dst_page], src_conv[src_page]) + assert torch.equal(target.big_pages_pool.ssm_state[dst_page], src_ssm[src_page]) + assert (target.big_pages_pool.conv_state[2] == -1).all() + assert (target.big_pages_pool.ssm_state[2] == -2).all() + assert (target.small_pages_pool.conv_state == 31).all() + assert (target.small_pages_pool.ssm_state == 32).all() + assert (target.conv_state == 41).all() and (target.ssm_state == 42).all() + assert (target.mtp_state_index == 1).all() + assert torch.equal(source.big_pages_pool.conv_state, src_conv) + assert torch.equal(source.big_pages_pool.ssm_state, src_ssm) diff --git a/unit_tests/common/att_resources/test_capacity.py b/unit_tests/common/att_resources/test_capacity.py new file mode 100644 index 0000000000..1784d06cb8 --- /dev/null +++ b/unit_tests/common/att_resources/test_capacity.py @@ -0,0 +1,1113 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from typing import TYPE_CHECKING +from dataclasses import replace +from functools import lru_cache, partial +from math import prod + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.utils import deepcopy_result +from lightllm.common.attention_state.allocators import StatePageAllocator, FullAttKVAllocator +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState +from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.common.attention_state.full_att import ( + FullAttentionState, + GQAFullAttentionState, + GQAAttentionConfig, + MLAConfig, + create_gqa_state, + create_mla_state, + GroupQuantizedAttentionConfig, + GroupQuantizedFullAttentionState, + MLAFullAttentionState, + PackedMLAFullAttentionState, + StaticFP8FullAttentionState, +) + +if TYPE_CHECKING: + from easydict import EasyDict + + +def definition(kind: str): + if kind == "mla": + return create_mla_state( + config=MLAConfig(logical_layer_groups=((0,), (2,)), latent_dim=576, latent_dtype=torch.bfloat16), + page_size=4, + device="meta", + ) + if kind.startswith("packed"): + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state + + return create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0,), (2,)), + latent_dim=512 if kind == "packed_glm" else 576, + format="packed_group128_fp8" if kind == "packed_fp8" else "packed_bf16", + ), + page_size=4, + device="meta", + ) + kv_type = { + "int4": "int4kv", + "int8": "int8kv", + "static_head": "fp8kv_sph", + "static_tensor": "fp8kv_spt", + }.get(kind, "None") + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + ), + kv_type=kv_type, + page_size=4, + device="meta", + ) + if kind.startswith("static"): + heads = 2 if kind == "static_head" else 1 + bind_test_fp8_calibration( + state=state, + scales=torch.ones((2, heads * 2), dtype=torch.float32), + q_scales=torch.ones((2, 2), dtype=torch.float32) if heads == 2 else None, + ) + return state + + +def allocated_bytes(state: FullAttentionState, token_only: bool = False): + buffers = list(full_state_test_buffers(state).values()) + if not token_only and isinstance(state, StaticFP8FullAttentionState): + buffers.append(state.scales) + if state.q_scales is not None: + buffers.append(state.q_scales) + stores = {} + for buffer in buffers: + storage = buffer.untyped_storage() + stores[(buffer.device, storage.data_ptr())] = storage.nbytes() + return sum(stores.values()) + + +@pytest.mark.parametrize( + "kind", ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] +) +def test_budget_matches_token_stores_including_hold_and_padding( + monkeypatch: pytest.MonkeyPatch, + kind: str, + pd_args: EasyDict, +): + template = definition(kind) + state = rebuild_full_state( + template=template, + config=template.config, + token_capacity=20, + device="cpu", + ) + actual = allocated_bytes(state, token_only=True) + planning = AttentionStateManager(request_capacity=1) + planning.register_state("main", rebuild_full_state(template, config=template.config)) + planning.finish_state_registration() + token_bytes = planning.get_full_att_token_bytes() + pd_page_bytes = planning.get_pd_operator().get_page_bytes() + assert not hasattr(planning, "pd_page_bytes") + assert not hasattr(planning, "pd_transfer_layouts") + assert token_bytes * (20 + 4) == actual + if kind.startswith("static"): + assert allocated_bytes(state) > actual + assert state.get_token_bytes() == token_bytes + minimum = token_bytes * 4 + assert minimum > 0 + + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.full_att_allocator = FullAttKVAllocator(20, 4, pin_memory=False) + manager.finish_state_registration() + (state,) = manager.get_full_attention_states("main") + assert manager.get_pd_operator().get_page_bytes() == pd_page_bytes + args = StartArgs( + page_size=4, + run_mode="prefill", + pd_kv_page_size=8, + pd_kv_page_num=3, + ) + reserve = manager.get_pd_operator().get_page_bytes() * args.pd_kv_page_num + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 0) + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + for budget, expected in ((actual, 20), (actual - 1, 16), (minimum, 0), (minimum - 1, None)): + resources = AttentionStateManager(request_capacity=1) + resources.register_state("main", rebuild_full_state(template, config=template.config)) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda _: (budget + reserve, budget + reserve)) + resources.finish_state_registration() + if expected is None: + with pytest.raises(MemoryError, match="HOLD"): + resources.initialize_storage( + args=args, + token_capacity=None, + mem_fraction=1.0, + device="cpu", + ) + elif expected == 0: + with pytest.raises(AssertionError, match="at least one complete page"): + resources.initialize_storage( + args=args, + token_capacity=None, + mem_fraction=1.0, + device="cpu", + ) + else: + resources.initialize_storage( + args=args, + token_capacity=None, + mem_fraction=1.0, + device="cpu", + ) + assert resources.full_att_allocator.size == expected + assert resources.get_pd_operator().get_page_bytes() == manager.get_pd_operator().get_page_bytes() + + +def test_shared_allocator_accounts_for_distinct_state_stores(): + main, draft = definition("int4"), definition("static_head") + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", main) + manager.register_state("draft0", draft) + manager.finish_state_registration() + expected = sum( + allocated_bytes( + rebuild_full_state( + template=state, + config=state.config, + token_capacity=12, + page_size=4, + device="cpu", + ), + token_only=True, + ) + for state in (main, draft) + ) + assert manager.get_full_att_token_bytes() * (12 + 4) == expected + + +@pytest.mark.parametrize("state_cls", [KDAState, GDNState]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") +def test_pd_state_page_covers_one_snapshot_without_transfer_metadata( + monkeypatch: pytest.MonkeyPatch, + state_cls: type[KDAState] | type[GDNState], + pd_args: EasyDict, +): + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=2, + global_v_heads=2 if state_cls is KDAState else 4, + key_dim=3, + value_dim=3 if state_cls is KDAState else 5, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + tail = WindowCompressorState( + config=WindowCompressorStateConfig( + logical_layer_ids=(1,), + window_size=4, + compress_stride=4, + target="indexer", + candidate_dim=3, + dtype=torch.bfloat16, + ), + request_capacity=3, + mtp_step=3, + device="cuda", + ) + reference = state_cls( + config=config, + request_capacity=3, + mtp_step=3, + device="cuda", + ) + conv = reference.conv_state[:, 0, :, : config.conv_width] + ssm = reference.ssm_state[:, 0] + bounded_bytes = conv.numel() * conv.element_size() + ssm.numel() * ssm.element_size() + tail_bytes = ( + len(tail.config.logical_layer_ids) + * tail.config.window_size + * 2 + * tail.config.candidate_dim + * tail.config.dtype.itemsize + ) + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", reference) + manager.register_state("main", tail) + manager.finish_state_registration() + # PD 页只包含运行态数据，使用第二页验证 8 字节页跨度下的读写。 + assert manager.get_pd_operator().get_page_bytes() == (bounded_bytes + 7) // 8 * 8 + (tail_bytes + 7) // 8 * 8 + reference.conv_state.fill_(2) + reference.ssm_state.fill_(3) + tail.buffer.fill_(4) + pages = torch.zeros((2, manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + layouts = manager.get_pd_operator().get_pd_bounded_att_page_layouts() + assert {region.field for layout in layouts.values() for region in layout.regions} == {"conv", "ssm", "compressor"} + manager.get_pd_operator().write_page( + pd_page_buffers=pages, + pd_page_index=1, + kind="att_state", + layouts=layouts, + req_idx=0, + state_token_len=8, + ) + reference.conv_state.zero_() + reference.ssm_state.zero_() + tail.buffer.zero_() + manager.get_pd_operator().read_page( + pd_page_buffers=pages, + pd_page_index=1, + kind="att_state", + layouts=layouts, + req_idx=0, + state_token_len=8, + ) + assert (reference.conv_state[:, 0, :, : config.conv_width] == 2).all() + assert (reference.ssm_state[:, 0] == 3).all() + expected_tail = torch.zeros_like(tail.buffer[:, 0]) + for position in range(4, 8): + expected_tail[:, position % tail.buffer.shape[2]].fill_(4) + assert torch.equal(tail.buffer[:, 0], expected_tail) + assert not pages[0].any() + for state in (reference, tail): + monkeypatch.setattr( + state, + "get_pd_transfer_regions", + lambda *args, **kwargs: pytest.fail("cached page size must not rebuild regions"), + ) + assert manager.get_pd_operator().get_page_bytes() == pages.shape[1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_pd_layout_reserves_aligned_storage_for_mixed_token_formats( + monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, +): + pd_args.pd_kv_page_size = 3 + manager = AttentionStateManager(request_capacity=1) + for state_cls, role, dim, dtype in ( + (MLAFullAttentionState, "draft0", 1, torch.bfloat16), + (GQAFullAttentionState, "main", 3, torch.bfloat16), + (GQAFullAttentionState, "main", 1, torch.float32), + ): + if state_cls is MLAFullAttentionState: + config = MLAConfig( + logical_layer_groups=((0,),), + latent_dim=dim, + latent_dtype=dtype, + ) + else: + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=dim, + dtype=dtype, + ) + manager.register_state(role, state_cls(config, page_size=1, device="meta")) + manager.finish_state_registration() + # 无页头；三个 State 的 3-token 区域分别由 6/36/24 字节对齐到 8/40/24。 + assert manager.get_pd_operator().get_page_bytes() == 72 + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.allocate_storage( + token_capacity=4, + page_size=1, + big_page_tokens=4, + device="cuda", + ) + assert manager.get_pd_operator().get_page_bytes() == 72 + states = manager.get_full_attention_states() + for i, state in enumerate(states, 1): + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.fill_(i) + pages = torch.zeros((2, manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + page = pages[1] + layouts = manager.get_pd_operator().get_pd_full_att_page_layouts() + assert [min(region.byte_start for region in layouts[state.state_name].regions) for state in states] == [0, 8, 48] + manager.get_pd_operator().write_page( + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + kind="kv", + layouts=layouts, + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + ) + assert (page[:6].view(torch.bfloat16) == 1).all() + for state in states: + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + manager.get_pd_operator().read_page( + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + kind="kv", + layouts=layouts, + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + ) + for i, state in enumerate(states, 1): + assert (read_kv_storage(state)[:, :3] == i).all() + assert not pages[0].any() + + +@pytest.mark.parametrize("layout_first", [False, True]) +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +def test_capacity_uses_quantized_field_padding_and_preserves_page_boundaries( + monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, + layout_first: bool, + transfer: str, +): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + class AlignedFieldsState(GroupQuantizedFullAttentionState): + def _get_transfer_regions(self, page_tokens: int): + regions, offset = [], 0 + for region in super()._get_transfer_regions(page_tokens): + offset = (offset + 7) // 8 * 8 + regions.append( + replace(region, byte_start=offset, byte_end=offset + region.byte_end - region.byte_start) + ) + offset += prod(region.shape) * region.dtype.itemsize + return tuple(regions) + + pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 + config = GroupQuantizedAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=128, + scale_dtype=torch.float32, + bits=8, + group_size=128, + ) + manager = AttentionStateManager(request_capacity=1) + prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") + prepared_full_state.allocate_storage(8) + manager.register_state("main", prepared_full_state) + prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") + prepared_full_state.allocate_storage(8) + manager.register_state("draft0", prepared_full_state) + manager.full_att_allocator = FullAttKVAllocator(8, 1, pin_memory=False) + manager.finish_state_registration() + # 每个 State 的裸数据为 792 字节；scale 间 padding 使实际末尾为 796，页段需 800 字节。 + for state in manager.get_full_attention_states(): + assert state.get_token_bytes() * 3 == 792 + monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("capacity must use field layouts")) + get_layouts = ( + manager.get_cpu_cache_operator().get_page_layouts + if transfer == "cpu" + else manager.get_pd_operator().get_pd_full_att_page_layouts + ) + get_page_bytes = ( + manager.get_cpu_cache_operator().get_page_bytes + if transfer == "cpu" + else manager.get_pd_operator().get_page_bytes + ) + if layout_first: + with monkeypatch.context() as patch: + operator = manager.get_cpu_cache_operator() if transfer == "cpu" else manager.get_pd_operator() + patch.setattr(operator, "get_page_bytes", lambda: pytest.fail("layout must not query page bytes")) + layouts = get_layouts() + assert get_page_bytes() == 1600 + layouts = get_layouts() + assert get_layouts() is layouts + pages = torch.full( + size=(3, 1600), + fill_value=91, + dtype=torch.uint8, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", + ) + for index, state in enumerate(manager.get_full_attention_states()): + layout = layouts[state.state_name] + assert (layout.byte_start, layout.byte_end) == (index * 800, (index + 1) * 800) + regions = layout.regions + assert [region.byte_start for region in regions] == [index * 800 + n for n in (0, 384, 768, 784)] + for field_index, field in enumerate(full_state_test_fields(state, full_state_test_buffers(state)).values()): + field.tensor.zero_() + field.tensor[:, :3].fill_(index * 4 + field_index + 1) + if transfer == "cpu": + manager.get_cpu_cache_operator().write_req( + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=layouts, + ) + manager.get_cpu_cache_operator().read_req( + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), + layouts=layouts, + ) + else: + manager.get_pd_operator().write_page( + pd_page_buffers=pages, + pd_page_index=1, + kind="kv", + layouts=layouts, + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + ) + manager.get_pd_operator().read_page( + pd_page_buffers=pages, + pd_page_index=1, + kind="kv", + layouts=layouts, + token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), + ) + for state in manager.get_full_attention_states(): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + torch.testing.assert_close(field.tensor[:, 4:7], field.tensor[:, :3]) + assert (pages[[0, 2]] == 91).all() + for offset in (0, 800): + assert (pages[1, offset + 780 : offset + 784] == 91).all() + assert (pages[1, offset + 796 : offset + 800] == 91).all() + + +@pytest.mark.parametrize("full_page_tokens", [None, 8, 10, 16]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") +def test_pd_runtime_capacity_includes_field_padding( + monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, + full_page_tokens: int, +): + class PaddedGDNState(GDNState): + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self): + return tuple( + replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) + if region.field == "ssm" + else region + for region in super().get_pd_transfer_regions() + ) + + manager = AttentionStateManager(request_capacity=1) + if full_page_tokens is not None: + pd_args.pd_kv_page_size = full_page_tokens + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=1, + dtype=torch.float32, + ) + manager.register_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) + state = PaddedGDNState( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ), + request_capacity=1, + device="cuda", + ) + manager.register_state("main", state) + manager.finish_state_registration() + # Conv 为 6 字节，SSM 从偏移 72 开始，占 4 字节；完整运行态需要 80 字节。 + # 即使 KV 页不足，布局仍可独立查询；最终块大小查询选择两类占用的最大值。 + with monkeypatch.context() as patch: + patch.setattr( + manager.get_pd_operator(), + "get_pd_full_att_page_layouts", + lambda: pytest.fail("runtime layout must be independent"), + ) + patch.setattr( + manager.get_pd_operator(), + "get_page_bytes", + lambda: pytest.fail("layout must not query page bytes"), + ) + layouts = manager.get_pd_operator().get_pd_bounded_att_page_layouts() + layout = layouts[state.state_name] + assert layout.byte_end == 80 and layout.tokens_per_page is None + assert layout.regions[-1].byte_start == 72 + warnings = [] + monkeypatch.setattr("lightllm.common.attention_state.pd_operator.logger.warning", warnings.append) + expected = max(80, (full_page_tokens or 0) * 8) + assert manager.get_pd_operator().get_page_bytes() == expected + assert manager.get_pd_operator().get_pd_bounded_att_page_layouts() is layouts + # 实际传输页可大于运行态布局，选中第二页后仍从该页起点搬运。 + pages = torch.full((3, expected), 91, dtype=torch.uint8, device="cuda") + state.conv_state.fill_(3) + state.ssm_state.fill_(5) + manager.get_pd_operator().write_page( + pd_page_buffers=pages, + pd_page_index=1, + kind="att_state", + layouts=layouts, + req_idx=0, + state_token_len=8, + ) + state.conv_state.zero_() + state.ssm_state.zero_() + manager.get_pd_operator().read_page( + pd_page_buffers=pages, + pd_page_index=1, + kind="att_state", + layouts=layouts, + req_idx=0, + state_token_len=8, + ) + assert (state.conv_state == 3).all() and (state.ssm_state == 5).all() + assert (pages[[0, 2]] == 91).all() and (pages[1, 80:] == 91).all() + assert len(warnings) == int(full_page_tokens == 8) + if warnings: + assert "64 bytes" in warnings[0] and "using 80 bytes per transfer page" in warnings[0] + # 相同 manager 的容量查询命中缓存，不重复打印 warning。 + assert manager.get_pd_operator().get_page_bytes() == expected + assert len(warnings) == int(full_page_tokens == 8) + + +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +def test_page_bytes_refresh_with_storage_dependent_layout( + monkeypatch: pytest.MonkeyPatch, + pd_args: EasyDict, + transfer: str, +): + class PaddedState(GQAFullAttentionState): + def _get_transfer_regions(self, page_tokens: int): + return tuple( + replace( + region, + byte_start=region.byte_start + self.page_size * 8, + byte_end=region.byte_end + self.page_size * 8, + ) + for region in super()._get_transfer_regions(page_tokens) + ) + + pd_args.pd_kv_page_size = pd_args.cpu_cache_token_page_size = 3 + manager = AttentionStateManager(request_capacity=1) + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=3, + dtype=torch.bfloat16, + ) + manager.register_state("main", PaddedState(config, page_size=1, device="meta")) + manager.finish_state_registration() + get_page_bytes = ( + manager.get_cpu_cache_operator().get_page_bytes + if transfer == "cpu" + else manager.get_pd_operator().get_page_bytes + ) + assert get_page_bytes() == 48 + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.allocate_storage( + token_capacity=8, + page_size=4, + big_page_tokens=8, + device="cpu", + ) + assert get_page_bytes() == (80 if transfer == "cpu" else 72) + layouts = ( + manager.get_cpu_cache_operator().get_page_layouts() + if transfer == "cpu" + else manager.get_pd_operator().get_pd_full_att_page_layouts() + ) + assert layouts["main_gqa_bf16_state"].byte_end == 72 + + +@pytest.mark.parametrize("with_tail", [False, True]) +@pytest.mark.parametrize("extra_tokens", [-2, 0, 2]) +def test_hybrid_pd_capacity_covers_all_runtime_states_before_storage_allocation( + with_tail: bool, + extra_tokens: int, + pd_args: EasyDict, +): + manager = AttentionStateManager(request_capacity=1) + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=1, + dtype=torch.float32, + ) + full_state = GQAFullAttentionState(config, page_size=1, device="meta") + manager.register_state("main", full_state) + for cls, role in ((KDAState, "main"), (GDNState, "draft0")): + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + state = cls(config, request_capacity=1, device="meta") + manager.register_state(role, state) + if with_tail: + state = WindowCompressorState( + config=WindowCompressorStateConfig( + logical_layer_ids=(0,), + window_size=4, + compress_stride=4, + target="indexer", + candidate_dim=1, + dtype=torch.float32, + ), + request_capacity=1, + mtp_step=0, + device="meta", + ) + manager.register_state("main", state) + # 两份运行态各占 16 字节数据，tail 另占 32 字节，无页头及长度、版本字段。 + runtime_page_bytes = 32 + 32 * with_tail + page_tokens = runtime_page_bytes // 8 + extra_tokens + pd_args.pd_kv_page_size = page_tokens + manager.finish_state_registration() + assert manager.get_pd_operator().get_page_bytes() == max(runtime_page_bytes, runtime_page_bytes + extra_tokens * 8) + assert manager.full_att_allocator is None + assert read_kv_storage(manager.main_gqa_fp32_state) is None + assert not hasattr(manager, "pd_transfer_layouts") + + +@pytest.mark.parametrize( + "kind", ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] +) +@pytest.mark.parametrize("tp_world_size", [1, 4]) +@pytest.mark.parametrize("page_tokens", [1, 3]) +def test_cpu_page_bytes_derive_from_cached_transfer_layouts( + monkeypatch: pytest.MonkeyPatch, + cpu_args: EasyDict, + kind: str, + tp_world_size: int, + page_tokens: int, +): + cpu_args.cpu_cache_token_page_size = page_tokens + template = definition(kind) + if hasattr(template.config, "logical_config"): + config = replace( + template.config, + logical_config=replace( + template.config.logical_config, + tp_world_size=tp_world_size, + tp_rank=tp_world_size - 1, + ), + ) + else: + config = replace(template.config, tp_world_size=tp_world_size, tp_rank=tp_world_size - 1) + scales = torch.ones((len(config.logical_layer_groups), 2 * config.head_num)) if kind == "static_head" else None + q_scales = torch.ones((len(config.logical_layer_groups), config.head_num)) if kind == "static_head" else None + manager = AttentionStateManager(request_capacity=3) + for role in ("main", "draft0"): + manager.register_state( + role, + rebuild_full_state( + template=template, + config=config, + scales=scales, + q_scales=q_scales, + ), + ) + for state_cls, v_heads, value_dim in ((KDAState, 1, 1), (GDNState, 2, 3)): + # 奇数长度的 BF16 Conv 与 FP32 SSM 之间需要补齐，MTP 候选区不应计入页大小。 + bounded_config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=v_heads, + key_dim=1, + value_dim=value_dim, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + tp_rank=tp_world_size - 1, + tp_world_size=tp_world_size, + ) + bounded = state_cls( + config=bounded_config, + request_capacity=3, + mtp_step=2, + device="meta", + ) + manager.register_state("main", bounded) + manager.finish_state_registration() + # CPU 容量和布局都只依赖 State 的 CPU 字段描述，不读取 GPU 存储开销或 PD 布局。 + for state in (*manager.get_full_attention_states(), *manager.get_bounded_attention_states()): + if isinstance(state, FullAttentionState): + monkeypatch.setattr(state, "get_token_bytes", lambda: pytest.fail("CPU capacity must use field layouts")) + monkeypatch.setattr(state, "get_pd_transfer_regions", lambda: pytest.fail("CPU layout must be independent")) + layouts = manager.get_cpu_cache_operator().get_page_layouts() + assert min(region.byte_start for layout in layouts.values() for region in layout.regions) == 0 + assert not {"state_token_len", "format_version"} & { + region.field for layout in layouts.values() for region in layout.regions + } + expected_page_bytes = (max(layout.byte_end for layout in layouts.values()) + 15) // 16 * 16 + assert all(layout.tokens_per_page == page_tokens for layout in layouts.values()) + for state in (*manager.get_full_attention_states(), *manager.get_bounded_attention_states()): + monkeypatch.setattr(state, "get_cpu_transfer_regions", lambda: pytest.fail("cached layout must be reused")) + assert manager.get_cpu_cache_operator().get_page_bytes() == expected_page_bytes + assert manager.get_cpu_cache_operator().get_page_layouts() is layouts + + +def test_cpu_bounded_capacity_includes_state_field_padding(cpu_args: EasyDict): + class PaddedGDNState(GDNState): + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self): + return tuple( + replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) + if region.field == "ssm" + else region + for region in super().get_cpu_transfer_regions() + ) + + manager = AttentionStateManager(request_capacity=1) + state = PaddedGDNState( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ), + request_capacity=1, + device="meta", + ) + manager.register_state("main", state) + manager.finish_state_registration() + # Conv 占 6 字节，SSM 位于 [72, 76)，State 与物理页终点均对齐到 80。 + assert manager.get_cpu_cache_operator().get_page_bytes() == 80 + layout = manager.get_cpu_cache_operator().get_page_layouts()[state.state_name] + assert (layout.byte_start, layout.byte_end) == (0, 80) + assert layout.regions[-1].byte_end == 76 + + +@pytest.mark.parametrize("kind", ["packed", "packed_fp8", "packed_glm"]) +def test_cpu_packed_page_excludes_unused_gpu_row_padding(kind: str, cpu_args: EasyDict): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + cpu_args.cpu_cache_token_page_size = 3 + states = [] + for encoding in ("gqa", kind): + template = definition(encoding) + state = rebuild_full_state( + template=template, + config=template.config, + token_capacity=8, + device="cuda", + ) + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.view(torch.uint8).fill_(211) + states.append(state) + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", states[0]) + manager.register_state("draft0", states[1]) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + manager.finish_state_registration() + layouts = manager.get_cpu_cache_operator().get_page_layouts() + packed, gqa = manager.get_full_attention_states() + packed_layout, gqa_layout = layouts[packed.state_name], layouts[gqa.state_name] + data_end = max(r.byte_end for r in packed_layout.regions) + gqa_start = gqa_layout.byte_start + # 后一个 State 紧跟实际字段终点，仅保留对齐空间，不为 GPU packed 行 padding 预留位置。 + assert gqa_start == (data_end + 7) // 8 * 8 + assert gqa_start < read_kv_storage(packed)[:, :3].numel() * read_kv_storage(packed).element_size() + for value, state in enumerate((packed, gqa), 3): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + field.tensor[:, :3].fill_(value) + field.tensor[:, 4:7].zero_() + pages = torch.full( + size=(3, manager.get_cpu_cache_operator().get_page_bytes()), + fill_value=165, + dtype=torch.uint8, + pin_memory=True, + ) + manager.get_cpu_cache_operator().write_req( + token_indices=torch.tensor([0, 1, 2], dtype=torch.long, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=layouts, + ) + torch.cuda.synchronize() + assert (pages[1, data_end:gqa_start] == 165).all() + assert (pages[[0, 2]] == 165).all() + manager.get_cpu_cache_operator().read_req( + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=torch.tensor([4, 5, 6], dtype=torch.long, device="cuda"), + layouts=layouts, + ) + for value, state in enumerate((packed, gqa), 3): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + assert (field.tensor[:, 4:7] == value).all() + + +@pytest.mark.parametrize("page_tokens", [1, 9]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") +def test_cpu_bounded_snapshots_reserve_independent_aligned_regions(cpu_args: EasyDict, page_tokens: int): + cpu_args.cpu_cache_token_page_size = page_tokens + manager = AttentionStateManager(request_capacity=1) + manager.bounded_att_big_page_allocator = StatePageAllocator(2) + manager.bounded_att_small_page_allocator = StatePageAllocator(1) + for state_cls, role, value_dim in ((KDAState, "draft0", 1), (GDNState, "main", 3)): + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=value_dim, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + state = state_cls( + config=config, + request_capacity=1, + mtp_step=2, + device="cuda", + ) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1, pin_memory=True) + manager.register_state(role, state) + manager.finish_state_registration() + states = manager.get_bounded_attention_states() + layouts = manager.get_cpu_cache_operator().get_page_layouts() + # 两份快照分别含 10/22 字节数据，预留 16/24 字节；页 token 数不放大快照。 + assert [min(r.byte_start for r in layouts[state.state_name].regions) for state in states] == [0, 16] + assert manager.get_cpu_cache_operator().get_page_bytes() == 48 + assert [(layout.byte_start, layout.byte_end) for layout in layouts.values()] == [(0, 16), (16, 40)] + for i, state in enumerate(states): + state.conv_state.fill_(i + 3) + state.ssm_state.fill_(i + 7) + source_slots = manager.bounded_att_big_page_allocator.alloc(1) + manager.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=source_slots, + ) + pages = torch.full((2, manager.get_cpu_cache_operator().get_page_bytes()), 87, dtype=torch.uint8, pin_memory=True) + manager.get_cpu_cache_operator().write_req( + token_indices=torch.tensor([], dtype=torch.long), + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=layouts, + big_page_slot_ids=source_slots, + ) + torch.cuda.synchronize() + assert (pages[0] == 87).all() + assert (pages[1, 12:16] == 87).all() + manager.reset_bounded_att_state([0]) + target_slots = manager.bounded_att_big_page_allocator.alloc(1) + manager.get_cpu_cache_operator().read_req( + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=torch.tensor([], dtype=torch.long), + state_token_len=8, + layouts=layouts, + big_page_slot_ids=target_slots, + req_idx=0, + ) + for i, state in enumerate(states): + assert (state.conv_state[:, 0, :, :1] == i + 3).all() + assert (state.ssm_state[:, 0] == i + 7).all() + + +@pytest.mark.parametrize("transfer", ["cpu", "pd_token", "pd_runtime"]) +def test_transfer_layouts_are_lazy_cached_and_refreshed_after_storage_allocation( + monkeypatch: pytest.MonkeyPatch, + transfer: str, + pd_args: EasyDict, +): + pd_args.pd_kv_page_size = 3 + pd_args.cpu_cache_token_page_size = 3 + manager = AttentionStateManager(request_capacity=1) + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=3, + dtype=torch.bfloat16, + ) + manager.register_state("main", GQAFullAttentionState(config, page_size=1, device="meta")) + manager.register_state( + "main", + GDNState( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=1, + value_dim=1, + conv_width=1, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.bfloat16, + ), + request_capacity=1, + device="cpu", + ), + ) + manager.finish_state_registration() + if transfer == "cpu": + operator = manager.get_cpu_cache_operator() + getter = operator.get_page_layouts + elif transfer == "pd_token": + operator = manager.get_pd_operator() + getter = operator.get_pd_full_att_page_layouts + else: + operator = manager.get_pd_operator() + getter = operator.get_pd_bounded_att_page_layouts + before = getter() + assert getter() is before + state = manager.main_gdn_conv_bf16_ssm_bf16_state if transfer == "pd_runtime" else manager.main_gqa_bf16_state + before_layout = before[state.state_name] + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.allocate_storage( + token_capacity=8, + page_size=4, + big_page_tokens=8, + device="cpu", + ) + after = getter() + after_layout = after[state.state_name] + assert after is not before and getter() is after + assert after_layout is not before_layout + full_state = manager.main_gqa_bf16_state + assert full_state.page_size == 4 and not read_kv_storage(full_state).is_meta + page_bytes = ( + manager.get_cpu_cache_operator().get_page_bytes() + if transfer == "cpu" + else manager.get_pd_operator().get_page_bytes() + ) + if transfer == "pd_runtime": + assert after_layout.byte_end == 8 and page_bytes == 40 + elif transfer == "cpu": + layout_end = max(layout.byte_end for layout in after.values()) + assert (layout_end + 15) // 16 * 16 == page_bytes + else: + assert after_layout.byte_end == page_bytes + assert ( + not {"cpu_page_tokens", "cpu_page_bytes", "cpu_transfer_layouts", "pd_page_bytes", "pd_transfer_layouts"} + & manager.__dict__.keys() + ) + + +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +def test_manager_uses_state_regions_and_adds_each_resource_offset(transfer: str, pd_args: EasyDict): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + class ReversedCpuRegionsState(GQAFullAttentionState): + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self): + regions = super().get_cpu_transfer_regions() + # 本实现的 CPU 页先保存 V 再保存 K，PD 仍沿用 K/V 顺序。 + return tuple( + replace(region, byte_start=other.byte_start, byte_end=other.byte_end) + for region, other in zip(regions, reversed(regions)) + ) + + pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=2, + dtype=torch.float32, + ) + manager = AttentionStateManager(request_capacity=1) + prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") + prepared_full_state.allocate_storage(8) + manager.register_state("main", prepared_full_state) + prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") + prepared_full_state.allocate_storage(8) + manager.register_state("draft0", prepared_full_state) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + manager.finish_state_registration() + layouts = ( + manager.get_cpu_cache_operator().get_page_layouts() + if transfer == "cpu" + else manager.get_pd_operator().get_pd_full_att_page_layouts() + ) + for index, state in enumerate(manager.get_full_attention_states()): + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + state.k_storage[:, :4].fill_(index + 3) + state.v_storage[:, :4].fill_(index + 7) + # 1 层、4 tokens，K/V 各 1 head、每 head 2 个 FP32 元素。 + state_bytes = 4 * 2 * 2 * torch.float32.itemsize + k_offset = state_bytes // 2 if transfer == "cpu" else 0 + k_region, v_region = layouts[state.state_name].regions + assert k_region.field == "k" and v_region.field == "v" + assert k_region.byte_start == index * state_bytes + k_offset + assert v_region.byte_start == index * state_bytes + state_bytes // 2 - k_offset + + page_bytes = ( + manager.get_cpu_cache_operator().get_page_bytes() + if transfer == "cpu" + else manager.get_pd_operator().get_page_bytes() + ) + pages = torch.zeros( + size=(1, page_bytes), + dtype=torch.uint8, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", + ) + if transfer == "cpu": + manager.get_cpu_cache_operator().write_req( + token_indices=torch.arange(4, dtype=torch.long, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[0], + layouts=layouts, + ) + manager.get_cpu_cache_operator().read_req( + cpu_page_buffers=pages, + cpu_page_indices=[0], + token_indices=torch.arange( + start=4, + end=8, + dtype=torch.long, + device="cuda", + ), + layouts=layouts, + ) + else: + manager.get_pd_operator().write_page( + pd_page_buffers=pages, + pd_page_index=0, + kind="kv", + layouts=layouts, + token_indices=torch.arange(4, dtype=torch.long, device="cuda"), + ) + manager.get_pd_operator().read_page( + pd_page_buffers=pages, + pd_page_index=0, + kind="kv", + layouts=layouts, + token_indices=torch.arange( + start=4, + end=8, + dtype=torch.long, + device="cuda", + ), + ) + for state in manager.get_full_attention_states(): + torch.testing.assert_close(read_kv_storage(state)[:, 4:8], read_kv_storage(state)[:, :4]) + assert not read_kv_storage(state)[:, 8:].any() diff --git a/unit_tests/common/att_resources/test_compute_binding.py b/unit_tests/common/att_resources/test_compute_binding.py new file mode 100644 index 0000000000..85988d5edd --- /dev/null +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_template import TransformerLayerInferTpl +from lightllm.common.basemodel.triton_kernel.kv_copy.prefill_compact_kv_flashmla_fp8 import ( + get_prefill_kv_cache_and_remap_indices_triton, +) +from unit_tests.common.att_resources.test_capacity import definition + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kind", ["gqa", "static_head", "packed_glm"]) +@pytest.mark.parametrize("role", ["main", "draft0"]) +@pytest.mark.parametrize("capture", [False, True]) +def test_transformer_writes_logical_layer_through_manager(kind: str, role: str, capture: bool): + template = definition(kind) + state = rebuild_full_state( + template=template, + config=replace_logical_layer_groups(template.config, ((3,), (8,))), + token_capacity=8, + device="cuda", + ) + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + manager = AttentionStateManager(request_capacity=2) + manager.register_state(role, state) + manager.full_att_allocator = FullAttKVAllocator(state.capacity, state.page_size, pin_memory=False) + manager.finish_state_registration() + state = getattr(manager, f"{role}_{template.get_state_type()}_state") + infer = InferStateInfo() + infer.att_state_manager = manager + infer.is_draft_model = role == "draft0" + infer.att_state_role = role + infer.mem_index = torch.tensor([1, 5], dtype=torch.int32, device="cuda") + heads = 1 if kind == "packed_glm" else 2 * state.config.head_num + dim = state.config.latent_dim if kind == "packed_glm" else state.config.head_dim + kv = torch.full( + size=(2, heads, dim), + fill_value=0.5, + dtype=torch.bfloat16, + device="cuda", + ) + layer = SimpleNamespace(layer_num_=8) + TransformerLayerInferTpl._post_cache_kv( + self=layer, + cache_kv=kv, + infer_state=infer, + layer_weight=None, + ) + if capture: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + TransformerLayerInferTpl._post_cache_kv( + self=layer, + cache_kv=kv, + infer_state=infer, + layer_weight=None, + ) + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + graph.replay() + assert not hasattr(infer, "mem_manager") + assert not hasattr(infer, "full_att_state") + assert not read_kv_storage(state)[0].view(torch.uint8).any() + view = state.get_layer_view(8) + if kind == "packed_glm": + assert torch.equal(view.kv[infer.mem_index], kv) + assert not view.indexer_k.any() + else: + k, v = view + actual = torch.cat((k[infer.mem_index], v[infer.mem_index]), dim=1) + if kind == "static_head": + actual = actual.view(torch.float8_e4m3fn).float() * state.scales[1][None, :, None] + # Backends index fixed scales by physical row, even when model layer + # IDs are sparse or include appended draft layers. + assert state.get_storage_layer_index(logical_layer_id=8) == 1 + assert state.cpu_scales.shape == (2, 4) + assert torch.equal(actual, kv) + + +def test_packed_fp8_prefill_combines_history_and_fresh_kv_with_sparse_remapping(): + template = definition("packed_fp8") + state = rebuild_full_state( + template=template, + config=template.config, + token_capacity=8, + device="cuda", + ) + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + history_ids = torch.tensor([3, 7], dtype=torch.int32, device="cuda") + history = torch.randn( + size=(2, 1, 576), + dtype=torch.bfloat16, + device="cuda", + ) + state.write_full_att_kv(0, history_ids, history) + fresh = torch.full( + size=(1, 1, 576), + fill_value=123, + dtype=torch.bfloat16, + device="cuda", + ) + topk = torch.tensor([[7, 5, 3, 7, -1], [5, -1, -1, -1, -1]], dtype=torch.int32, device="cuda") + compact, mapped = get_prefill_kv_cache_and_remap_indices_triton( + packed_kv=state.get_layer_view(0).kv, + topk_mem_indices=topk, + prefill_mem_index=torch.tensor([5], dtype=torch.int32, device="cuda"), + prefill_cache_kv=fresh, + prefill_dtype=fresh.dtype, + ) + assert compact.shape == (3, 1, 576) + assert torch.equal(mapped == -1, topk == -1) + assert mapped[0, 0] == mapped[0, 3] + assert torch.equal(compact[mapped[0, 1]], fresh[0]) + assert torch.equal(compact[mapped[1, 0]], fresh[0]) + torch.testing.assert_close( + actual=compact[mapped[0, 0], :, :512], + expected=history[1, :, :512], + atol=0.15, + rtol=0.08, + ) + torch.testing.assert_close( + actual=compact[mapped[0, 2], :, :512], + expected=history[0, :, :512], + atol=0.15, + rtol=0.08, + ) + assert torch.equal(compact[mapped[0, 0], :, 512:], history[1, :, 512:]) + assert torch.equal(compact[mapped[0, 2], :, 512:], history[0, :, 512:]) diff --git a/unit_tests/common/att_resources/test_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py new file mode 100644 index 0000000000..badc2adf0c --- /dev/null +++ b/unit_tests/common/att_resources/test_dsv4_integration.py @@ -0,0 +1,287 @@ +"""通过模型工厂和生产 fused kernel 验证 V4 State 接入。""" + +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel +from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo +from lightllm.utils.envs_utils import get_env_start_args + + +@pytest.fixture +def cache(monkeypatch): + for key in ("LIGHTLLM_DP_WORLD_SIZE", "LIGHTLLM_CURRENT_RANK_IN_DP", "LIGHTLLM_CURRENT_DEVICE_ID"): + monkeypatch.setenv(key, "1" if key.endswith("WORLD_SIZE") else "0") + monkeypatch.setenv( + "LIGHTLLM_START_ARGS", + json.dumps( + dict( + page_size=256, + pd_kv_page_size=256, + run_mode="normal", + enable_dp_prompt_cache_fetch=False, + linear_att_hash_page_size=256, + linear_att_page_block_num=8, + cpu_cache_token_page_size=2048, + ) + ), + ) + get_env_start_args.cache_clear() + config = dict( + num_hidden_layers=3, + compress_ratios=[4, 128, 0], + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + sliding_window=128, + num_attention_heads=64, + ) + args = SimpleNamespace(llm_kv_type="None", page_size=256, mtp_step=3) + manager = AttentionStateManager(request_capacity=3) + for state in DeepseekV4TpPartModel.create_full_att_states( + model_config=config, + args=args, + model_logical_layer_ids=(0, 1, 2), + ): + state.device = torch.device("cuda") + manager.register_state("main", state) + bounded = DeepseekV4TpPartModel.create_bounded_att_states( + model_config=config, + args=args, + model_logical_layer_ids=(0, 1, 2), + request_capacity=3, + device="cuda", + ) + for state in bounded: + manager.register_state("main", state) + window = next(state for state in bounded if isinstance(state, FP8E8M0SlidingWindowState)) + manager.finish_state_registration() + manager.allocate_storage(token_capacity=8192, page_size=256, big_page_tokens=2048, small_page_capacity=2) + + class TestModel(SimpleNamespace): + __hash__ = object.__hash__ + + model = TestModel( + config=config, + tp_rank_=0, + tp_world_size_=1, + max_req_num=2, + batch_max_tokens=2048, + layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], + att_state_role="main", + args=SimpleNamespace(run_mode="normal", enable_prefill_microbatch_overlap=True), + ) + from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend + + backend = DeepseekV4FlashMlaFp8SparseAttBackend(model=model) + yield manager, window, backend + torch.cuda.synchronize() + get_env_start_args.cache_clear() + + +def make_infer(manager, window): + infer = DeepseekV4InferStateInfo() + infer.att_state_manager = manager + infer.att_state_role = "main" + infer.first_logical_layer_id = window.config.logical_layer_ids[0] + return infer + + +def read_packed_rows(buffer, slots, page_tokens): + """独立参考读取 packed 数据/scale 分区，不复用 State 搬运算子。""" + page_ids, rows = slots // page_tokens, slots % page_tokens + data = rows[:, None] * 576 + torch.arange(576, device=slots.device) + scale = page_tokens * 576 + rows[:, None] * 8 + torch.arange(8, device=slots.device) + return torch.cat((buffer[page_ids[:, None], data], buffer[page_ids[:, None], scale]), dim=-1) + + +def test_backend_reads_owned_state_views(cache): + manager, window, backend = cache + infer = make_infer(manager, window) + assert infer.swa_state is window + infer.is_prefill = False + for layer in range(3): + swa, compressed = backend.get_att_input_params(infer, layer) + assert swa.shape[1] == 74880 + assert (compressed is None) == (layer == 2) + if layer != 2: + assert compressed.data_ptr() == infer.get_compressed_kv_view(layer).kv.data_ptr() + + +@pytest.mark.parametrize("ratio,layer,is_indexer", [(4, 0, False), (4, 0, True), (128, 1, False)]) +def test_real_compressor_continues_after_checkpoint_restore(cache, ratio, layer, is_indexer): + from lightllm.models.deepseek_v4.layer_infer.compressor import fused_compress + + manager, window, _ = cache + length, boundary = 1024, 512 + head_dim = 128 if is_indexer else 512 + scores = torch.randn((length, head_dim * (4 if ratio == 4 else 2)), device="cuda") + cos = torch.ones((length, 32), device="cuda") + sin, weight = torch.zeros_like(cos), torch.ones(head_dim, device="cuda") + slots = [torch.arange(offset, offset + length, dtype=torch.int32, device="cuda") for offset in (0, 2048)] + + def compress(req_idx, start, end, decode=False): + infer = make_infer(manager, window) + positions = torch.arange(start, end, dtype=torch.int32, device="cuda") + width = end - start + infer.mem_index, infer.position_ids, infer.is_prefill = slots[req_idx][start:end], positions, not decode + infer._dsv4_token_to_batch_idx = None if decode else torch.zeros(width, dtype=torch.int32, device="cuda") + infer.b_req_idx = torch.full((width if decode else 1,), req_idx, dtype=torch.int32, device="cuda") + infer.b_mtp_index = ( + torch.arange(width, dtype=torch.int32, device="cuda") + if decode + else torch.zeros(1, dtype=torch.int32, device="cuda") + ) + infer.b_seq_len = positions + 1 if decode else torch.tensor([end], dtype=torch.int32, device="cuda") + infer.b_ready_cache_len = None if decode else torch.tensor([start], dtype=torch.int32, device="cuda") + infer.b_q_start_loc = None if decode else torch.zeros(1, dtype=torch.int32, device="cuda") + return fused_compress( + kv_score=scores[start:end], + infer_state=infer, + layer_idx=layer, + norm_weight=weight, + eps=1e-6, + head_dim=head_dim, + qk_rope_head_dim=64, + compress_ratio=ratio, + cos_table=cos, + sin_table=sin, + is_in_indexer=is_indexer, + ) + + reference = compress(0, 0, length) + reference_kv = None if is_indexer else make_infer(manager, window).get_compressed_kv_view(layer).kv.clone() + manager.reset_bounded_att_state([0, 1]) + compress(0, 0, boundary - 2) + compress(0, boundary - 2, boundary) + manager.save_bounded_att_runtime_to_pages([0], [boundary], "small", [0]) + manager.restore_bounded_att_runtime_from_pages([1], [boundary], "small", [0]) + for start, end in ((boundary, 768), (768, length - 4), (length - 4, length)): + output = compress(1, start, end, decode=end == length) + if is_indexer: + closed = (torch.arange(start, end, device="cuda") + 1) % ratio == 0 + torch.testing.assert_close(output[closed], reference[start:end][closed], rtol=0, atol=0) + if not is_indexer: + actual = make_infer(manager, window).get_compressed_kv_view(layer).kv + source = slots[0][boundary + ratio - 1 : length : ratio].long() // ratio + target = slots[1][boundary + ratio - 1 : length : ratio].long() // ratio + torch.testing.assert_close( + read_packed_rows(reference_kv, source, 256 // ratio), + read_packed_rows(actual, target, 256 // ratio), + rtol=0, + atol=0, + ) + + +def test_long_prefill_keeps_swa_prefix_and_runtime_tail(cache): + from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer + from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_kv_norm_rope_to_swa + from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + + manager, window, backend = cache + reference = torch.zeros((17, 74880), dtype=torch.uint8, device="cuda") + kv = torch.randn((2048, 512), dtype=torch.bfloat16, device="cuda") + weight, angles = torch.randn(512, dtype=torch.bfloat16, device="cuda"), torch.randn((2048, 32), device="cuda") + freqs = torch.complex(angles.cos(), angles.sin()) + layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + layer.layer_num_, layer.eps_, layer.freqs_cis = 2, 1e-6, freqs + for start, end in ((0, 255), (255, 641), (641, 1100), (1100, 2048)): + infer = make_infer(manager, window) + positions = torch.arange(start, end, device="cuda", dtype=torch.int32) + infer.is_prefill, infer.position_ids, infer.microbatch_index = True, positions, 0 + infer.b_req_idx = torch.zeros(1, dtype=torch.int32, device="cuda") + infer.b_ready_cache_len, infer.b_seq_len = ( + torch.tensor([value], dtype=torch.int32, device="cuda") for value in (start, end) + ) + prefill = backend.create_att_prefill_state(infer) + prefill.prepare_swa_workspace(infer) + infer.prefill_att_state = prefill + indices = torch.empty((positions.numel(), 128), dtype=torch.int32, device="cuda") + lengths, writes, runtime = (torch.empty_like(positions) for _ in range(3)) + build_swa_indices( + token_req_idx=torch.zeros_like(positions), + position_ids=positions, + swa_state=window, + swa_indices=indices, + swa_lengths=lengths, + swa_write_slots=writes, + prefill_state=prefill, + swa_runtime_write_slots=runtime, + token_to_batch_idx=torch.zeros_like(positions), + b_seq_len=infer.b_seq_len, + ) + infer.dsv4_swa_write_slots, infer.dsv4_swa_runtime_write_slots = writes, runtime + layer._norm_rope_and_cache_kv(kv[start:end], infer, SimpleNamespace(kv_norm_=SimpleNamespace(weight=weight))) + fused_kv_norm_rope_to_swa(kv[start:end], weight, 1e-6, freqs, positions, positions, reference, 128) + visible = torch.arange(max(0, start - 127), end, device="cuda") + offset = prefill.swa_prefill_workspace_offsets[0] + visible - max(0, start - 127) // 128 * 128 + actual = read_packed_rows(prefill.swa_buffer[2].view(-1, 74880), offset.long(), 128) + torch.testing.assert_close(actual[:, :-1], read_packed_rows(reference, visible, 128)[:, :-1], rtol=0, atol=0) + tail = torch.arange(max(0, end - window.config.window_size), end, device="cuda") + actual = read_packed_rows(window.swa_buffer[2, 0].view(-1, 74880), tail % window.swa_capacity_tokens, 128) + torch.testing.assert_close(actual[:, :-1], read_packed_rows(reference, tail, 128)[:, :-1], rtol=0, atol=0) + + +@pytest.mark.parametrize("draft_kind", ["mtp", "dspark"]) +def test_models_prepare_owned_states_in_one_manager(cache, draft_kind): + from contextlib import nullcontext + from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel + from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel + + _, _, backend = cache + config = dict( + backend.model.config, + compress_ratios=[4, 128, 0, 0, 0], + dspark_target_layer_ids=[0, 1], + dspark_block_size=4, + dspark_noise_token_id=1, + dspark_markov_rank=4, + ) + args = SimpleNamespace( + llm_kv_type="None", + page_size=256, + mtp_step=3, + mtp_mode="vanilla_with_att", + run_mode="decode", + enable_prefill_microbatch_overlap=False, + ) + main = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) + main.config, main.args, main.max_req_num, main.tp_rank_, main.tp_world_size_ = config, args, 2, 0, 1 + main.batch_max_tokens = 2048 + main.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(3)] + main.torch_memory_saver = SimpleNamespace(region=lambda **kwargs: nullcontext()) + main.prepare_attention_states() + cls = DeepseekV4MTPModel if draft_kind == "mtp" else DeepseekV4DSparkModel + drafts = [] + for index in range(2 if draft_kind == "mtp" else 1): + draft = cls.__new__(cls) + draft.main_model, draft.mtp_previous_draft_models = main, drafts.copy() + draft.config = cls.normalize_model_config(config) + assert draft.config == cls.normalize_model_config(draft.config) + draft.args, draft.max_req_num, draft.tp_rank_, draft.tp_world_size_ = args, 2, 0, 1 + draft.batch_max_tokens = main.batch_max_tokens + draft.torch_memory_saver = main.torch_memory_saver + first_layer = 3 + sum(len(model.layers_infer) for model in drafts) + draft.layers_infer = [ + SimpleNamespace(layer_num_=first_layer + i) for i in range(draft.config["num_hidden_layers"]) + ] + draft.prepare_attention_states() + drafts.append(draft) + manager = main.att_state_manager + assert all(model.att_state_manager is manager for model in drafts) + manager.finish_state_registration() + manager.allocate_storage(token_capacity=8192, page_size=256, big_page_tokens=2048, small_page_capacity=2) + assert len(manager.get_full_attention_states()) == 2 + assert len(manager.get_bounded_attention_states()) == len(drafts) + 4 + for model in (main, *drafts): + states = manager.get_bounded_attention_states(model.att_state_role) + state = next(item for item in states if isinstance(item, FP8E8M0SlidingWindowState)) + assert state.config.logical_layer_ids == tuple(layer.layer_num_ for layer in model.layers_infer) + assert state.swa_buffer.is_cuda + if model is not main: + assert len(states) == 1 diff --git a/unit_tests/common/att_resources/test_dsv4_state.py b/unit_tests/common/att_resources/test_dsv4_state.py new file mode 100644 index 0000000000..ae2ba27179 --- /dev/null +++ b/unit_tests/common/att_resources/test_dsv4_state.py @@ -0,0 +1,505 @@ +"""V4 State 独立验证：packed 字段、逻辑快照及不同 MTP ring 的 PD 恢复。""" + +from dataclasses import replace + +import pytest +import torch + +from lightllm.models.deepseek_v4.config import DSV4AttentionConfig +from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import ( + CompressedAttentionConfig, + CompressedAttentionState, +) +from lightllm.common.attention_state.bounded_att import ( + BoundedAttentionState, + FP8E8M0SlidingWindowConfig, + FP8E8M0SlidingWindowState, + WindowCompressorStateConfig, + WindowCompressorState, +) +from lightllm.common.attention_state.types import StateTransferLayout + + +def make_config() -> DSV4AttentionConfig: + return DSV4AttentionConfig( + logical_layer_ids=(2, 5, 8), + compress_ratios=(0, 4, 128), + head_dim=512, + rope_head_dim=64, + indexer_head_dim=128, + sliding_window=128, + ) + + +def make_full(ratio: int, device: str = "meta") -> CompressedAttentionState: + config = make_config() + layout = CompressedAttentionConfig( + logical_layer_groups=((5 if ratio == 4 else 8,),), + compress_ratio=ratio, + kv_dim=config.head_dim, + rope_dim=config.rope_head_dim, + block_token_size=config.compressed_kv_token_alignment, + indexer_dim=config.indexer_head_dim if ratio == 4 else 0, + ) + return CompressedAttentionState(layout_config=layout, page_size=256, device=device) + + +def pd_layout(state: BoundedAttentionState | CompressedAttentionState) -> StateTransferLayout: + regions = state.get_pd_transfer_regions() + return StateTransferLayout( + state_name=state.state_name, + tokens_per_page=None if isinstance(state, BoundedAttentionState) else 256, + byte_start=0, + byte_end=regions[-1].byte_end, + regions=regions, + ) + + +def gather_swa_reference(state: FP8E8M0SlidingWindowState, req_idx: int, length: int) -> torch.Tensor: + """按格式的独立参考计算读回窗口，覆盖数据/scale 分区和环绕。""" + config = state.config + packed = state.swa_buffer[:, req_idx].cpu().reshape(len(config.logical_layer_ids), -1) + result = torch.zeros((len(config.logical_layer_ids), config.window_size, config.kv_bytes), dtype=torch.uint8) + for window, position in enumerate(range(length - config.window_size, length)): + if position < 0: + continue + slot = position % state.swa_capacity_tokens + page = slot // 128 * config.get_packed_kv_bytes(128) + token = slot % 128 + data = page + token * config.kv_data_bytes + scale = page + 128 * config.kv_data_bytes + token * (config.kv_bytes - config.kv_data_bytes) + result[:, window, : config.kv_data_bytes] = packed[:, data : data + config.kv_data_bytes] + result[:, window, config.kv_data_bytes :] = packed[:, scale : scale + config.kv_bytes - config.kv_data_bytes] + return result + + +def test_compressed_meta_geometry() -> None: + config = make_config() + swa_config = config.create_swa_config() + assert swa_config.kv_bytes == 584 and swa_config.get_packed_kv_bytes(128) // 128 == 585 + for ratio, slot_bytes in ((4, 180), (128, 9)): + state = make_full(ratio=ratio) + assert state.config.token_slot_bytes == slot_bytes + assert state.storage is None + state.allocate_storage(token_capacity=512) + assert state.storage.shape == (1, 768, 1, slot_bytes) + view = state.get_layer_view(logical_layer_id=5 if ratio == 4 else 8) + assert view.kv.shape[0] == 3 + assert (view.indexer_k is not None) == (ratio == 4) + attention_view = state.config.get_attention_view(packed_kv=view.kv) + assert attention_view.shape == (3, 256 // ratio, 1, 584) + assert attention_view.stride(0) == view.kv.stride(0) + with pytest.raises(NotImplementedError, match="fused compressor"): + state.write_full_att_kv( + logical_layer_id=5 if ratio == 4 else 8, + token_indices=torch.empty(0, dtype=torch.int64, device="meta"), + kv=torch.empty(0, device="meta"), + ) + + +def test_compressed_factory_uses_independent_config() -> None: + from lightllm.common.attention_state.full_att import create_compressed_att_state + + definition = make_full(ratio=4) + assert type(definition.config).__bases__ == (object,) + state = create_compressed_att_state(config=definition.config, page_size=256, device="meta") + assert isinstance(state, CompressedAttentionState) + assert state.storage is None and state.get_token_bytes() == 180 + assert state.get_state_type() == "compressed_c4_fp8_group64_e8m0" + + +@pytest.mark.parametrize("page_tokens", [0, 128, 384]) +def test_compressed_transfer_pages_require_complete_blocks(page_tokens: int, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.utils.envs_utils import get_env_start_args + + args = get_env_start_args() + monkeypatch.setattr(args, "cpu_cache_token_page_size", page_tokens) + monkeypatch.setattr(args, "pd_kv_page_size", page_tokens) + state = make_full(ratio=4) + with pytest.raises(AssertionError, match="page token capacity"): + state.get_cpu_transfer_regions() + with pytest.raises(AssertionError, match="page token capacity"): + state.get_pd_transfer_regions() + + +@pytest.mark.parametrize("write", [True, False]) +@pytest.mark.parametrize("invalid", ["page_tokens", "page_bytes", "byte_start", "page_index"]) +def test_compressed_cpu_page_boundaries( + write: bool, + invalid: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """错误的布局或物理页在提交 GPU kernel 前失败，避免跨页写入其他缓存数据。""" + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setattr(get_env_start_args(), "cpu_cache_token_page_size", 256) + state = make_full(ratio=4) + regions = state.get_cpu_transfer_regions() + layout = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=256, + byte_start=0, + byte_end=regions[-1].byte_end, + regions=regions, + ) + pages = torch.empty((2, layout.byte_end), dtype=torch.uint8) + page_indices = [0] + match invalid: + case "page_tokens": + layout = replace(layout, tokens_per_page=128) + case "page_bytes": + pages = pages[:, :-1] + case "byte_start": + layout = replace(layout, byte_start=-1) + case "page_index": + page_indices = [2] + # 只需索引数量，不创建 CUDA Tensor；所有非法输入必须在读取实际存储前被拦截。 + indices = torch.empty(256, dtype=torch.int32, device="meta") + with pytest.raises(AssertionError, match="CPU page|CPU pages"): + if write: + state.write_to_cpu( + token_indices=indices, + cpu_page_buffers=pages, + cpu_page_indices=page_indices, + layout=layout, + ) + else: + state.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=page_indices, + token_indices=indices, + layout=layout, + ) + + +@pytest.mark.parametrize("ratio,block_tokens,indexer_dim,slot_bytes", [(4, 256, 0, 153), (8, 512, 32, 81)]) +def test_atomic_compressed_config_controls_layout( + ratio: int, + block_tokens: int, + indexer_dim: int, + slot_bytes: int, +) -> None: + """无需模型配置即可创建 State；indexer 是否存在独立于压缩比例。""" + from lightllm.common.attention_state.full_att import create_compressed_att_state + + config = CompressedAttentionConfig( + logical_layer_groups=((0, 4), (8,)), + compress_ratio=ratio, + kv_dim=512, + rope_dim=64, + block_token_size=block_tokens, + indexer_dim=indexer_dim, + tp_rank=1, + tp_world_size=2, + ) + state = create_compressed_att_state(config=config, page_size=block_tokens, device="meta") + assert not hasattr(config, "cache_config") + assert config.token_slot_bytes == slot_bytes + assert state.get_token_bytes() == 2 * slot_bytes + state.allocate_storage(token_capacity=2 * block_tokens) + assert state.storage.shape == (2, 3 * block_tokens, 1, slot_bytes) + assert state.get_storage_layer_index(0) == state.get_storage_layer_index(4) + view = state.get_layer_view(logical_layer_id=4) + assert view.kv.shape == (3, config.kv_block_bytes) + if indexer_dim: + assert view.indexer_k.shape == (3, block_tokens // ratio * (indexer_dim + 4)) + else: + assert view.indexer_k is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("ratio", [4, 128]) +def test_pd_compressed_kv_preserves_partial_block(ratio: int, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 256) + source, target = make_full(ratio=ratio, device="cuda"), make_full(ratio=ratio, device="cuda") + source.allocate_storage(token_capacity=512) + target.allocate_storage(token_capacity=512) + source.storage.random_(0, 256) + target.storage.zero_() + layout = pd_layout(source) + pages = torch.zeros((1, layout.byte_end), dtype=torch.uint8, device="cuda") + indices = torch.arange(17, dtype=torch.int32, device="cuda") + source.write_to_pd_page(token_indices=indices, pd_page_buffers=pages, pd_page_index=0, layout=layout) + target.read_from_pd_page(pd_page_buffers=pages, token_indices=indices, pd_page_index=0, layout=layout) + torch.cuda.synchronize() + assert torch.equal(source.storage[:, :256], target.storage[:, :256]) + assert not target.storage[:, 256:].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("ratio", [4, 128]) +def test_compressed_cpu_and_token_copy(ratio: int, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setattr(get_env_start_args(), "cpu_cache_token_page_size", 256) + source, target = make_full(ratio=ratio, device="cuda"), make_full(ratio=ratio, device="cuda") + source.allocate_storage(token_capacity=512) + target.allocate_storage(token_capacity=768) + source.storage.random_(0, 256) + regions = source.get_cpu_transfer_regions() + layout = StateTransferLayout( + state_name=source.state_name, + tokens_per_page=256, + byte_start=0, + byte_end=regions[-1].byte_end, + regions=regions, + ) + cpu_pages = torch.zeros((3, layout.byte_end), dtype=torch.uint8, pin_memory=True) + src_indices = torch.arange(256, 512, dtype=torch.int32, device="cuda") + dst_indices = torch.arange(256, dtype=torch.int64, device="cuda") + source.write_to_cpu( + token_indices=src_indices, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[2], + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=cpu_pages, + cpu_page_indices=[2], + token_indices=dst_indices, + layout=layout, + ) + target.copy_tokens_from_dp( + src_token_indices=src_indices, + dst_token_indices=dst_indices + 512, + remote_dp_full_att_state=source, + ) + torch.cuda.synchronize() + assert torch.equal(source.storage[:, 256:512], target.storage[:, :256]) + assert torch.equal(source.storage[:, 256:512], target.storage[:, 512:768]) + assert not target.storage[:, 256:512].any() + assert not cpu_pages[:2].any() + target.copy_full_att_tokens(src_indices=dst_indices, dst_indices=dst_indices + 256) + torch.cuda.synchronize() + assert torch.equal(source.storage[:, 256:512], target.storage[:, 256:512]) + assert not target.storage[:, 768:].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("ratio", [4, 128]) +@pytest.mark.parametrize("token_count", [128, 257]) +@pytest.mark.parametrize("from_dp", [False, True]) +def test_compressed_token_copy_rejects_partial_blocks(ratio: int, token_count: int, from_dp: bool) -> None: + """即使长度已按压缩比例对齐，也不能拆开含 KV/scale/indexer 的物理块。""" + source, target = make_full(ratio=ratio, device="cuda"), make_full(ratio=ratio, device="cuda") + source.allocate_storage(token_capacity=512) + target.allocate_storage(token_capacity=512) + indices = torch.arange(token_count, dtype=torch.int64, device="cuda") + with pytest.raises(AssertionError, match="complete compressed blocks"): + if from_dp: + target.copy_tokens_from_dp( + src_token_indices=indices, + dst_token_indices=indices, + remote_dp_full_att_state=source, + ) + else: + target.copy_full_att_tokens(src_indices=indices, dst_indices=indices) + + +def make_atomic_state(kind: str, mtp_step: int = 0, device: str = "meta") -> BoundedAttentionState: + if kind == "swa": + return FP8E8M0SlidingWindowState( + config=FP8E8M0SlidingWindowConfig(logical_layer_ids=(2, 5), window_size=128, kv_dim=512, rope_dim=64), + request_capacity=3, + mtp_step=mtp_step, + device=device, + ) + window, compress_stride, candidate_dim, target = { + "kv_w8": (8, 4, 1024, "kv"), + "indexer_w8": (8, 4, 256, "indexer"), + "kv_w128": (128, 128, 512, "kv"), + "generic_w12": (12, 4, 24, "kv"), + }[kind] + return WindowCompressorState( + config=WindowCompressorStateConfig( + logical_layer_ids=(2, 5), + window_size=window, + compress_stride=compress_stride, + candidate_dim=candidate_dim, + target=target, + ), + request_capacity=3, + mtp_step=mtp_step, + device=device, + ) + + +def runtime_buffer(state: BoundedAttentionState) -> torch.Tensor: + return state.swa_buffer if isinstance(state, FP8E8M0SlidingWindowState) else state.buffer + + +@pytest.mark.parametrize("kind", ["swa", "kv_w8", "indexer_w8", "kv_w128", "generic_w12"]) +def test_atomic_window_definitions(kind: str) -> None: + source, target = make_atomic_state(kind), make_atomic_state(kind, mtp_step=5) + assert source.get_pd_transfer_regions() == target.get_pd_transfer_regions() + cpu_regions = source.get_cpu_transfer_regions() + copied_regions = source.get_cpu_transfer_regions() + assert copied_regions == cpu_regions + if cpu_regions: + assert copied_regions is not cpu_regions and copied_regions[0] is not cpu_regions[0] + assert source.get_layer_view(2).shape[0] == 3 + assert not source.contains_logical_layer(8) + assert not hasattr(source.config, "compress_ratios") + if isinstance(source, WindowCompressorState): + assert source.config.overlap_tokens == source.config.window_size - source.config.compress_stride + assert (source.get_cpu_transfer_regions() == ()) == (source.config.overlap_tokens == 0) + source.initialize_page_pools(big_page_capacity=3, device="meta", pin_memory=False) + assert (source.big_pages_pool is None) == (source.config.overlap_tokens == 0) + else: + assert source.swa_capacity_tokens == 128 and target.swa_capacity_tokens == 256 + for state in (source, target): + attention_view = state.config.get_attention_view(packed_kv=state.get_layer_view(2)) + assert attention_view.shape == (3 * state.swa_capacity_tokens // 128, 128, 1, 584) + assert attention_view.stride(0) == state.config.get_packed_kv_bytes(128) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", ["swa", "kv_w8", "indexer_w8", "kv_w128", "generic_w12"]) +@pytest.mark.parametrize("length", [0, 1, 5, 127, 128, 129, 255, 260]) +@pytest.mark.parametrize("source_step,target_step", [(0, 5), (5, 0)]) +def test_atomic_pd_windows_ignore_mtp_capacity(kind: str, length: int, source_step: int, target_step: int) -> None: + source = make_atomic_state(kind, mtp_step=source_step, device="cuda") + target = make_atomic_state(kind, mtp_step=target_step, device="cuda") + runtime_buffer(source).random_(0, 127) + target.reset_bounded_att_state([1]) + layout = pd_layout(source) + pages = torch.full((2, layout.byte_end + 32), 53, dtype=torch.uint8, device="cuda") + source.write_to_pd_page(req_idx=0, pd_page_buffers=pages, pd_page_index=1, state_token_len=length, layout=layout) + target.read_from_pd_page(pd_page_buffers=pages, req_idx=1, pd_page_index=1, state_token_len=length, layout=layout) + torch.cuda.synchronize() + if isinstance(source, FP8E8M0SlidingWindowState): + assert torch.equal(gather_swa_reference(source, 0, length), gather_swa_reference(target, 1, length)) + else: + src, dst = runtime_buffer(source), runtime_buffer(target) + for position in range(max(0, length - source.config.window_size), length): + assert torch.equal(src[:, 0, position % src.shape[2]], dst[:, 1, position % dst.shape[2]]) + assert (pages[0] == 53).all() and (pages[1, layout.byte_end :] == 53).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", ["swa", "kv_w8", "indexer_w8", "kv_w128", "generic_w12"]) +def test_atomic_page_cpu_and_dp_roundtrip(kind: str) -> None: + source = make_atomic_state(kind, mtp_step=3, device="cuda") + target = make_atomic_state(kind, mtp_step=3, device="cuda") + source.initialize_page_pools(big_page_capacity=3) + target.initialize_page_pools(big_page_capacity=3) + runtime_buffer(source).random_(0, 127) + target.reset_bounded_att_state([0, 1]) + source.save_bounded_att_runtime_to_pages( + req_indices=[0], state_token_lens=[512], page_kind="big", page_slot_ids=[0] + ) + target.copy_runtime_state_from_dp(src_req_indices=[0], dst_req_indices=[1], remote_dp_bounded_state=source) + regions = tuple( + replace(region, byte_start=region.byte_start + 64, byte_end=region.byte_end + 64) + for region in source.get_cpu_transfer_regions() + ) + if not regions: + assert source.big_pages_pool is None + with pytest.raises(AssertionError, match="stride"): + source.restore_bounded_att_runtime_from_pages([0], [513], "big", [0]) + return + layout = StateTransferLayout( + state_name=None, tokens_per_page=256, byte_start=64, byte_end=regions[-1].byte_end, regions=regions + ) + pages = torch.full((3, layout.byte_end + 32), 53, dtype=torch.uint8, pin_memory=True) + source.write_state_pages_to_cpu( + page_kind="big", + page_slot_ids=[0, 0], + cpu_page_buffers=pages, + cpu_page_indices=[1, 2], + layout=layout, + cpu_page_ready_flags=[True, False], + ) + target.read_state_pages_from_cpu(cpu_page_buffers=pages, cpu_page_indices=[2], big_page_slot_ids=[0], layout=layout) + target.copy_big_pages_from_dp(src_page_ids=[0], dst_page_ids=[2], remote_dp_bounded_state=source) + target.restore_bounded_att_runtime_from_pages( + req_indices=[0], state_token_lens=[512], page_kind="big", page_slot_ids=[0] + ) + torch.cuda.synchronize() + assert torch.equal(source.big_pages_pool.buffer[0], target.big_pages_pool.buffer[0]) + assert torch.equal(source.big_pages_pool.buffer[0], target.big_pages_pool.buffer[2]) + assert torch.equal(runtime_buffer(source)[:, 0], runtime_buffer(target)[:, 1]) + assert (pages[:2] == 53).all() and (pages[2, :64] == 53).all() and (pages[2, layout.byte_end :] == 53).all() + if isinstance(source, FP8E8M0SlidingWindowState): + assert torch.equal(gather_swa_reference(source, 0, 512), gather_swa_reference(target, 0, 512)) + else: + for position in range(512 - source.config.overlap_tokens, 512): + ring = source.buffer.shape[2] + assert torch.equal(source.buffer[:, 0, position % ring], target.buffer[:, 0, position % ring]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_manager_joint_cpu_page_roundtrip(monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.attention_state import AttentionStateManager + from lightllm.utils.envs_utils import get_env_start_args + + args = get_env_start_args() + monkeypatch.setattr(args, "cpu_cache_token_page_size", 256) + monkeypatch.setattr(args, "pd_kv_page_size", 256) + manager = AttentionStateManager(request_capacity=3) + full_states = [make_full(ratio=ratio) for ratio in (4, 128)] + bounded = [make_atomic_state(kind, mtp_step=2, device="cuda") for kind in ("swa", "kv_w8", "indexer_w8", "kv_w128")] + for state in (*full_states, *bounded): + manager.register_state(model_role="main", state=state) + manager.finish_state_registration() + manager.allocate_storage(token_capacity=512, page_size=256, big_page_tokens=256, device="cuda") + for full in full_states: + full.storage[:, :256].random_(0, 256) + for state in bounded: + runtime_buffer(state).random_(0, 127) + manager.save_bounded_att_runtime_to_pages([0], [256], "big", [0]) + expected = gather_swa_reference(bounded[0], 0, 256) + layouts = manager.get_cpu_cache_operator().get_page_layouts() + assert bounded[-1].state_name not in layouts + pages = torch.zeros((2, manager.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + manager.get_cpu_cache_operator().write_req( + token_indices=torch.arange(256, dtype=torch.int64, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=layouts, + big_page_slot_ids=[0], + ) + manager.get_cpu_cache_operator().read_req( + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=torch.arange(256, 512, dtype=torch.int64, device="cuda"), + layouts=layouts, + big_page_slot_ids=[1], + req_idx=1, + state_token_len=256, + ) + torch.cuda.synchronize() + for full in full_states: + assert torch.equal(full.storage[:, :256], full.storage[:, 256:512]) + assert torch.equal(expected, gather_swa_reference(bounded[0], 1, 256)) + assert not pages[0].any() + ipc = manager.get_pd_operator().create_ipc_copy() + assert all(state.big_pages_pool is None for state in ipc.get_bounded_attention_states()) + assert bounded[0].big_pages_pool is not None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_atomic_dp_ipc_preserves_shared_pools(monkeypatch: pytest.MonkeyPatch) -> None: + import json + from lightllm.common.attention_state import AttentionStateManager + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps({"enable_dp_prompt_cache_fetch": True})) + get_env_start_args.cache_clear() + try: + manager = AttentionStateManager(request_capacity=3) + for kind in ("swa", "kv_w8", "kv_w128"): + state = make_atomic_state(kind, device="cuda") + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + exported = manager.get_dp_shared_operator().create_ipc_copy() + for state in exported.get_bounded_attention_states(): + for pool in (state.big_pages_pool, state.small_pages_pool): + if pool is not None: + assert pool.buffer.is_shared() + exported.get_dp_shared_operator().register_page_pools_pin_memory() + finally: + get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_fp8_compute_binding.py b/unit_tests/common/att_resources/test_fp8_compute_binding.py new file mode 100644 index 0000000000..7d2a8019de --- /dev/null +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -0,0 +1,194 @@ +"""FP8 backends obtain calibration from the role's matching storage scheme.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups + +from typing import TYPE_CHECKING +from dataclasses import replace +from types import SimpleNamespace +import weakref + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.basemodel.attention.base_att import AttControl, BaseAttBackend +from lightllm.common.basemodel.attention.fa3 import fp8 as fa3 +from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer +from lightllm.common.basemodel.infer_struct import InferStateInfo +from unit_tests.common.att_resources.test_capacity import definition + +if TYPE_CHECKING: + from collections.abc import Sequence + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kind", ["static_head", "static_tensor"]) +@pytest.mark.parametrize("role", ["main", "draft0"]) +@pytest.mark.parametrize("prefill", [False, True]) +def test_fp8_kernel_receives_calibration_for_role_and_physical_layer( + monkeypatch: pytest.MonkeyPatch, + kind: str, + role: str, + prefill: bool, +): + template = definition(kind) + states = [] + for value in (2, 5): + scales = template.scales + value + scales[1] += 10 + q_scales = None if kind == "static_tensor" else template.q_scales + value + states.append( + rebuild_full_state( + template=template, + config=replace_logical_layer_groups(template.config, ((3,), (8,))), + token_capacity=8, + device="cuda", + scales=scales, + q_scales=q_scales, + ) + ) + manager = AttentionStateManager(request_capacity=2) + manager.register_state("main", states[0]) + manager.register_state("draft0", states[1]) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + # Another encoding in the same role must not supply the FP8 calibration. + plain = definition("gqa") + plain = rebuild_full_state( + template=plain, + config=plain.config, + token_capacity=8, + device="cuda", + ) + manager.register_state(role, plain) + manager.finish_state_registration() + selected = getattr(manager, f"{role}_{template.get_state_type()}_state") + k, v = selected.get_layer_view(8) + infer = InferStateInfo() + infer.att_state_manager = manager + infer.is_draft_model = role == "draft0" + infer.att_state_role = role + infer.batch_size, infer.max_q_seq_len = 2, 1 + infer.b_seq_len = torch.ones(2, dtype=torch.int32, device="cuda") + backend = object.__new__(BaseAttBackend) + observed = {} + + if kind == "static_head": + state_cls = fa3.Fp8Fa3PrefillAttState if prefill else fa3.Fp8Fa3DecodeAttState + parent_cls = fa3.Fa3PrefillAttState if prefill else fa3.Fa3DecodeAttState + monkeypatch.setattr(parent_cls, "init_state", lambda self: None) + att = state_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + att.b_att_seq_len = infer.b_seq_len + att.init_state() + + def quantize(q: torch.Tensor, scales: torch.Tensor): + observed["q_scales"] = scales + return q + + def attention(**kwargs: object): + observed.update(kwargs) + return kwargs["q"] + + monkeypatch.setattr(fa3, "q_per_head_static_fp8_quant", quantize) + monkeypatch.setattr(fa3, "flash_attn_with_kvcache", attention) + else: + state_cls = flashinfer.Fp8FlashInferPrefillAttState if prefill else flashinfer.Fp8FlashInferDecodeAttState + att = state_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + wrapper = SimpleNamespace(run=lambda *args, **kwargs: observed.update(kwargs)) + att.prefill_wrapper = att.decode_wrapper = wrapper + + q = torch.zeros((2, selected.config.head_num, selected.config.head_dim), dtype=torch.bfloat16, device="cuda") + + def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.device | int): + return torch.empty(shape, dtype=dtype, device=device) + + if prefill: + att.prefill_att( + q=q, + k=k, + v=v, + att_control=AttControl(logical_layer_id=8), + alloc_func=allocate, + ) + else: + att.decode_att( + q=q, + k=k, + v=v, + att_control=AttControl(logical_layer_id=8), + alloc_func=allocate, + ) + if kind == "static_head": + heads = selected.config.head_num + torch.testing.assert_close(observed["q_scales"], selected.q_scales[1]) + torch.testing.assert_close(observed["k_descale"], selected.scales[1, :heads].expand(2, heads)) + torch.testing.assert_close(observed["v_descale"], selected.scales[1, heads:].expand(2, heads)) + else: + assert observed["k_scale"] == selected.cpu_scales[1, 0] + assert observed["v_scale"] == selected.cpu_scales[1, 1] + assert not hasattr(infer, "full_att_state") + + +@pytest.mark.parametrize("head_num", [1, 3]) +@pytest.mark.parametrize("scale_mode", ["per_tensor", "per_head"]) +def test_fp8_split_storage_matches_calibrated_values_and_combined_kernel(head_num: int, scale_mode: str): + from lightllm.common.attention_state.full_att import StaticFP8AttentionConfig, StaticFP8FullAttentionState + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv_fp8 import destindex_copy_kv_fp8 + + scale_count = 2 if scale_mode == "per_tensor" else 2 * head_num + scales = ( + torch.arange( + start=1, + end=scale_count + 1, + dtype=torch.float32, + device="cuda", + ).unsqueeze(0) + / 4 + ) + state = StaticFP8FullAttentionState( + layout_config=StaticFP8AttentionConfig( + logical_layer_groups=((5,),), + global_head_num=head_num, + head_dim=128, + scale_mode=scale_mode, + ), + page_size=4, + device="cuda", + ) + bind_test_fp8_calibration( + state=state, + scales=scales, + q_scales=torch.ones((1, head_num), dtype=torch.float32, device="cuda") if scale_mode == "per_head" else None, + ) + state.allocate_storage(8) + state.k_storage.zero_() + state.v_storage.zero_() + indices = torch.tensor([1, 6], dtype=torch.int32, device="cuda") + kv = torch.randn((2, 2 * head_num, 128), dtype=torch.bfloat16, device="cuda") + state.write_full_att_kv(logical_layer_id=5, token_indices=indices, kv=kv) + k, v = state.get_layer_view(5) + assert k.untyped_storage().data_ptr() != v.untyped_storage().data_ptr() + actual = torch.cat((k[indices], v[indices]), dim=1) + head_scales = scales[0].repeat_interleave(head_num) if scale_mode == "per_tensor" else scales[0] + expected = (kv.float() / head_scales[None, :, None]).clamp(-448, 448).to(torch.float8_e4m3fn) + assert torch.equal(actual, expected.view(torch.uint8)) + assert not k[[0, 2, 3, 4, 5, 7, 8, 9, 10, 11]].any() + assert not v[[0, 2, 3, 4, 5, 7, 8, 9, 10, 11]].any() + + # 算子的原合并输出入口仍应保持相同的量化结果。 + combined = torch.zeros((12, 2 * head_num, 128), dtype=torch.uint8, device="cuda") + destindex_copy_kv_fp8( + KV=kv, + DestLoc=indices, + scales=scales[0], + Out=combined.view(torch.float8_e4m3fn), + is_per_tensor_quant=scale_mode == "per_tensor", + ) + assert torch.equal(combined[indices], actual) diff --git a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py new file mode 100644 index 0000000000..8e2ee61214 --- /dev/null +++ b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py @@ -0,0 +1,235 @@ +"""独立 FP8/E8M0 MLA：显式选择、固定编码、完整 payload 搬运和 padding 保留。""" + +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.attention_state.full_att import ( + FP8E8M0PackedMLAConfig, + FP8E8M0PackedMLAState, + PackedMLALogicalConfig, + FullAttentionState, + PackedMLAFullAttentionState, + create_packed_mla_state, +) +from lightllm.common.attention_state.types import StateTransferLayout + + +if TYPE_CHECKING: + from easydict import EasyDict + + +def make_state(device: str = "meta", capacity: int = 16, tp_rank: int = 0) -> FP8E8M0PackedMLAState: + state = create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0, 4), (8,)), + latent_dim=512, + format="packed_group64_e8m0", + tp_rank=tp_rank, + tp_world_size=2, + ), + page_size=4, + device=device, + ) + state.allocate_storage(token_capacity=capacity) + return state + + +def test_factory_selects_independent_encoding_and_meta_geometry(pd_args: EasyDict) -> None: + pd_args.cpu_cache_token_page_size = 4 + state = make_state() + assert type(state) is FP8E8M0PackedMLAState + assert FP8E8M0PackedMLAState.__bases__ == (FullAttentionState,) + assert not isinstance(state, PackedMLAFullAttentionState) + assert isinstance(state.config, FP8E8M0PackedMLAConfig) + assert type(state.config).__bases__ == (object,) + assert not hasattr(state.config, "head_dim") + assert state.get_state_type() == "mla_packed_fp8_group64_e8m0" + assert state.storage.shape == (2, 20, 1, 720) and state.storage.is_meta + assert state.config.payload_bytes == 716 and state.get_token_bytes() == 1440 + assert state.get_storage_layer_index(logical_layer_id=0) == state.get_storage_layer_index(logical_layer_id=4) + assert state.get_layer_view(logical_layer_id=8).kv.shape == (20, 1, 584) + assert state.get_layer_view(logical_layer_id=8).indexer_k.shape == (20, 1, 132) + (region,) = state.get_cpu_transfer_regions() + assert region.shape == (2, 4, 1, 716) and region.byte_end == 2 * 4 * 716 + + # 默认仍是原来的 group128/FP32 编码，不能再通过 512 维自动猜测 GLM 格式。 + with pytest.raises(AssertionError, match="latent 576"): + create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0,),), + latent_dim=512, + format="packed_group128_fp8", + ), + device="meta", + ) + with pytest.raises(AssertionError, match="512-dimensional"): + create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0,),), + latent_dim=576, + format="packed_group64_e8m0", + ), + device="meta", + ) + + +@pytest.mark.parametrize("draft", [False, True]) +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) +def test_glm_main_and_draft_explicitly_select_encoding(draft: bool, kv_type: str) -> None: + from lightllm.models.glm5_next.model import Glm5NextTpPartModel + from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + + config = dict( + num_hidden_layers=4, + num_attention_heads=4, + hidden_size=2048, + kv_lora_rank=512, + index_kpool=4, + index_head_dim=128, + layer_types=["linear_attention"] * 3 + ["deepseek_sparse_attention"], + linear_attn_config=dict(num_heads=4, head_dim=128, short_conv_kernel_size=4), + ) + args = SimpleNamespace( + tp=2, dp=1, llm_kv_type=kv_type, page_size=4, data_type="bfloat16", linear_att_ssm_data_type="float32" + ) + model_class = Glm5NextMTPModel if draft else Glm5NextTpPartModel + logical_ids = (14,) if draft else (10, 11, 12, 13) + (state,) = model_class.create_full_att_states( + model_config=config, + args=args, + model_logical_layer_ids=logical_ids, + tp_rank=1, + tp_world_size=2, + ) + assert type(state) is (FP8E8M0PackedMLAState if kv_type == "fp8kv_dsa" else PackedMLAFullAttentionState) + assert state.config.logical_layer_groups == ((logical_ids[-1],),) + assert state.storage is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fixed_quantizer_preserves_indexer_and_other_rows() -> None: + from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv + + state = make_state(device="cuda") + state.storage.fill_(173) + indices = torch.tensor([1, 7, 12], dtype=torch.int32, device="cuda") + kv = torch.randn((3, 1, 512), dtype=torch.bfloat16, device="cuda") + expected = torch.empty((3, 1, 584), dtype=torch.uint8, device="cuda") + write_fp8_e8m0_packed_kv(kv=kv, token_indices=torch.arange(3, dtype=torch.int32, device="cuda"), packed_kv=expected) + state.write_full_att_kv(logical_layer_id=4, token_indices=indices, kv=kv) + view = state.get_layer_view(logical_layer_id=0) + torch.testing.assert_close(view.kv[indices], expected, rtol=0, atol=0) + assert view.indexer_k.eq(173).all() + assert state.storage[..., state.config.payload_bytes :].eq(173).all() + assert state.storage[1].eq(173).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +@pytest.mark.parametrize("tp_rank", [0, 1]) +@pytest.mark.parametrize("reorder_axes", [False, True]) +def test_full_payload_roundtrip_skips_padding( + transfer: str, + tp_rank: int, + reorder_axes: bool, + pd_args: EasyDict, +) -> None: + pd_args.cpu_cache_token_page_size, pd_args.pd_kv_page_size = 4, 16 + source = make_state(device="cuda") + target = make_state(device="cuda", tp_rank=tp_rank) + source.storage.random_(0, 256) + target.storage.fill_(173) + regions = source.get_cpu_transfer_regions() if transfer == "cpu" else source.get_pd_transfer_regions() + if reorder_axes: + regions = tuple( + replace( + region, + axes=("token", "layer", "head", "dim"), + shape=(region.shape[1], region.shape[0], *region.shape[2:]), + ) + for region in regions + ) + page_tokens = 4 if transfer == "cpu" else 16 + layout = StateTransferLayout( + state_name=source.state_name, + tokens_per_page=page_tokens, + byte_start=0, + byte_end=regions[0].byte_end, + regions=regions, + ) + pages = torch.full( + (5, layout.byte_end + 32), + 99, + dtype=torch.uint8, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", + ) + page_view = pages[:, 16:-16] + src_ids = torch.tensor([11, 0, 4, 3, 13, 6, 5, 8, 2], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([15, 2, 3, 7, 12, 9, 1, 8, 0], dtype=torch.int64, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + if transfer == "cpu": + source.write_to_cpu( + token_indices=src_ids, cpu_page_buffers=page_view, cpu_page_indices=[3, 0, 4], layout=layout + ) + target.read_from_cpu( + cpu_page_buffers=page_view, cpu_page_indices=[3, 0, 4], token_indices=dst_ids, layout=layout + ) + else: + source.write_to_pd_page(token_indices=src_ids, pd_page_buffers=page_view, pd_page_index=3, layout=layout) + target.read_from_pd_page(pd_page_buffers=page_view, token_indices=dst_ids, pd_page_index=3, layout=layout) + stream.synchronize() + torch.testing.assert_close(target.storage[:, dst_ids, :, :716], source.storage[:, src_ids, :, :716], rtol=0, atol=0) + assert target.storage[..., 716:].eq(173).all() + assert target.storage[:, 16:].eq(173).all() + assert pages[:, :16].eq(99).all() and pages[:, -16:].eq(99).all() + assert pages[1].eq(99).all() and pages[2].eq(99).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_cpu_ready_pages_remain_unchanged(pd_args: EasyDict) -> None: + pd_args.cpu_cache_token_page_size = 4 + source = make_state(device="cuda") + source.storage.fill_(7) + regions = source.get_cpu_transfer_regions() + layout = StateTransferLayout( + state_name=source.state_name, tokens_per_page=4, byte_start=0, byte_end=regions[0].byte_end, regions=regions + ) + pages = torch.full((3, layout.byte_end), 99, dtype=torch.uint8, pin_memory=True) + source.write_to_cpu( + token_indices=torch.arange(9, dtype=torch.int32, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=[2, 0, 1], + layout=layout, + cpu_page_ready_flags=[False, True, False], + ) + torch.cuda.synchronize() + assert pages[0].eq(99).all() + assert pages[2].eq(7).all() + partial_page = pages[1].view(2, 4, 1, 716) + assert partial_page[:, 0].eq(7).all() + assert partial_page[:, 1:].eq(99).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_local_and_dp_copy_include_entire_physical_row() -> None: + source = make_state(device="cuda") + target = make_state(device="cuda", capacity=32) + source.storage.random_(0, 256) + target.storage.fill_(173) + src_ids = torch.tensor([0, 3, 11], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([5, 20, 31], dtype=torch.int64, device="cuda") + target.copy_tokens_from_dp(src_token_indices=src_ids, dst_token_indices=dst_ids, remote_dp_full_att_state=source) + torch.testing.assert_close(target.storage[:, dst_ids], source.storage[:, src_ids], rtol=0, atol=0) + source.copy_full_att_tokens( + src_indices=src_ids, dst_indices=torch.tensor([4, 6, 8], dtype=torch.int64, device="cuda") + ) + torch.testing.assert_close(source.storage[:, [4, 6, 8]], source.storage[:, src_ids], rtol=0, atol=0) diff --git a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py new file mode 100644 index 0000000000..0197b17daa --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -0,0 +1,412 @@ +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +"""CPU 页直接拷贝：真实索引、跨页布局、原始编码与未参与区域的保持。""" + +from dataclasses import replace + +import pytest +import torch + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers +from unit_tests.common.att_resources.test_state_identity import allocated_state +from unit_tests.common.att_resources.test_quantized_full_state import transfer_layout +from lightllm.common.attention_state.types import AxisMapping +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("page_device", ["cpu", "cuda"]) +def test_narrow_head_fields_merge_before_page_transfer(page_device: str, monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.attention_state.full_att import gqa_transfer + + config = GQAAttentionConfig( + logical_layer_groups=((0,), (1,)), + global_head_num=4, + head_dim=1, + dtype=torch.bfloat16, + ) + source = GQAFullAttentionState(layout_config=config, page_size=4, device="cuda") + target = GQAFullAttentionState(layout_config=config, page_size=4, device="cuda") + source.allocate_storage(token_capacity=8) + target.allocate_storage(token_capacity=12) + src_ids = torch.tensor([5, 1, 7], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([9, 0, 6], dtype=torch.int64, device="cuda") + expected = [] + for src_tensor, dst_tensor in ((source.k_storage, target.k_storage), (source.v_storage, target.v_storage)): + src_tensor.view(torch.uint8).random_(0, 256) + dst_tensor.view(torch.uint8).fill_(173) + snapshot = dst_tensor.clone() + snapshot[:, dst_ids] = src_tensor[:, src_ids] + expected.append(snapshot) + + layout = transfer_layout(source) + pages = torch.full( + size=(2, layout.byte_end + 32), + fill_value=91, + dtype=torch.uint8, + device=page_device, + pin_memory=page_device == "cpu", + ) + page_view = pages[:, 16:-16] + launches = [] + kernel = gqa_transfer._copy_gqa_pages + + class CaptureLaunch: + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs): + launches.append((kwargs["HEADS"], kwargs["WIDTH"], kwargs["ELEMENT_BYTES"])) + return kernel[grid](**kwargs) + + return launch + + monkeypatch.setattr(gqa_transfer, "_copy_gqa_pages", CaptureLaunch()) + if page_device == "cpu": + source.write_to_cpu( + token_indices=src_ids, + cpu_page_buffers=page_view, + cpu_page_indices=[1], + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=page_view, + cpu_page_indices=[1], + token_indices=dst_ids, + layout=layout, + ) + else: + source.write_to_pd_page( + token_indices=src_ids, + pd_page_buffers=page_view, + pd_page_index=1, + layout=layout, + ) + target.read_from_pd_page( + pd_page_buffers=page_view, + token_indices=dst_ids, + pd_page_index=1, + layout=layout, + ) + torch.cuda.synchronize() + # 单个 BF16 head 只有 2 字节；合并 4 个 head 后才能按 INT64 搬运。 + assert launches == [(1, 1, 8)] * 4 + for actual, reference in zip((target.k_storage, target.v_storage), expected): + assert torch.equal(actual.view(torch.uint8), reference.view(torch.uint8)) + assert (pages[0] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, -16:] == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "kind", + [ + "gqa", + "mla", + "int4", + "int8", + "static_head", + "static_tensor", + "packed", + "packed_fp8", + "packed_glm", + ], +) +@pytest.mark.parametrize("token_major", [False, True]) +def test_direct_cpu_pages_preserve_bytes_without_gather(kind: str, token_major: bool, monkeypatch: pytest.MonkeyPatch): + source = allocated_state(kind, device="cuda") + target = allocated_state(kind, device="cuda") + for tensor in full_state_test_buffers(source).values(): + tensor.view(torch.uint8).random_(0, 256) + for tensor in full_state_test_buffers(target).values(): + tensor.view(torch.uint8).fill_(173) + layout = transfer_layout(source) + if token_major: + regions = [] + for region in layout.regions: + axes = ("token", "layer", *region.axes[2:]) + regions.append(replace(region, axes=axes, shape=tuple(region.shape[region.axes.index(a)] for a in axes))) + layout = replace(layout, regions=tuple(regions)) + # 页池本身是更大联合页的视图，检验真实 page stride，而不是按字段长度推算下一页。 + pages = torch.full( + size=(4, layout.byte_end + 64), + fill_value=91, + dtype=torch.uint8, + pin_memory=True, + ) + page_view = pages[:, 32:-32] + src_ids = torch.tensor([6, 6, 1, 1, 7, 7, 0, 0, 3, 3], dtype=torch.int32, device="cuda")[::2] + dst_ids = torch.tensor([2, 2, 6, 6, 4, 4, 1, 1, 7, 7], dtype=torch.int64, device="cuda")[::2] + expected = {name: tensor[:, dst_ids].clone() for name, tensor in full_state_test_buffers(target).items()} + source_fields = full_state_test_fields(source, full_state_test_buffers(source)) + expected_fields = full_state_test_fields(target, expected) + for name, field in source_fields.items(): + expected_fields[name].tensor.copy_(field.tensor[:, src_ids]) + + def reject_gather(*args: object, **kwargs: object) -> None: + pytest.fail("CPU transfer must not allocate gathered KV or scatter it back") + + tensor_to = torch.Tensor.to + tensor_contiguous = torch.Tensor.contiguous + + def reject_index_to(tensor: torch.Tensor, *args: object, **kwargs: object): + if tensor is src_ids or tensor is dst_ids: + pytest.fail("token indices must be passed directly to the kernel") + return tensor_to(tensor, *args, **kwargs) + + def reject_index_contiguous(tensor: torch.Tensor, *args: object, **kwargs: object): + if tensor is src_ids or tensor is dst_ids: + pytest.fail("token index views must not be materialized") + return tensor_contiguous(tensor, *args, **kwargs) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "to", reject_index_to) + patch.setattr(torch.Tensor, "contiguous", reject_index_contiguous) + patch.setattr(torch.Tensor, "index_select", reject_gather) + patch.setattr(torch.Tensor, "index_copy_", reject_gather) + with torch.cuda.stream(stream): + source.write_to_cpu( + token_indices=src_ids, + cpu_page_buffers=page_view, + cpu_page_indices=[2, 0], + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=page_view, + cpu_page_indices=[2, 0], + token_indices=dst_ids, + layout=layout, + ) + stream.synchronize() + for name, tensor in full_state_test_buffers(target).items(): + assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected[name].view(torch.uint8)) + assert (tensor[:, [0, 3, 5, 8, 9, 10, 11]].contiguous().view(torch.uint8) == 173).all() + assert (pages[:, :32] == 91).all() and (pages[:, -32:] == 91).all() + assert (pages[[1, 3]] == 91).all() + for region in layout.regions: + external = page_view[0, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + # 最后一页只有一个有效 token，其余槽位仍保留填充值。 + unused = [slice(None)] * len(region.shape) + unused[region.axes.index("token")] = slice(1, None) + assert (external[tuple(unused)].contiguous().view(torch.uint8) == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_direct_cpu_pages_only_update_selected_head_slice_and_skip_replica_writes(): + state = allocated_state("gqa", device="cuda") + state.k_storage.fill_(3) + state.v_storage.fill_(7) + layout = transfer_layout(state) + regions = [] + for region in layout.regions: + slices = tuple( + AxisMapping( + axis="head", + local_start=1, + external_start=0, + length=1, + ) + if item.axis == "head" + else item + for item in region.axis_mappings + ) + regions.append(replace(region, axis_mappings=slices)) + layout = replace(layout, regions=tuple(regions)) + pages = torch.zeros((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) + ids = torch.tensor([1, 5], dtype=torch.int64, device="cuda") + state.write_to_cpu( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=layout, + ) + torch.cuda.synchronize() + snapshot = pages.clone() + state.k_storage.fill_(19) + state.v_storage.fill_(19) + read_only = replace(layout, regions=tuple(replace(region, write_enabled=False) for region in layout.regions)) + state.write_to_cpu( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=read_only, + ) + state.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[0], + token_indices=ids, + layout=read_only, + ) + torch.cuda.synchronize() + assert torch.equal(pages, snapshot) + assert (state.k_storage[:, ids, 1] == 3).all() and (state.v_storage[:, ids, 1] == 7).all() + assert (state.k_storage[:, :, 0] == 19).all() and (state.v_storage[:, :, 0] == 19).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_direct_cpu_pages_accept_registered_external_memory(): + """共享缓存通过 cudaHostRegister 注册，而非 torch 的 pinned allocator 分配。""" + import ctypes + import mmap + + state = allocated_state("gqa", device="cuda") + state.k_storage.fill_(3) + state.v_storage.fill_(7) + layout = transfer_layout(state) + memory = mmap.mmap(-1, layout.byte_end) + pages = torch.frombuffer(memory, dtype=torch.uint8).view(1, -1) + cuda = ctypes.CDLL("/usr/local/cuda/targets/x86_64-linux/lib/libcudart.so") + cuda.cudaHostRegister.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint] + cuda.cudaHostUnregister.argtypes = [ctypes.c_void_p] + pointer = ctypes.c_void_p(pages.data_ptr()) + assert cuda.cudaHostRegister(pointer, layout.byte_end, 3) == 0 + try: + assert pages.is_pinned() + ids = torch.tensor([1, 5], dtype=torch.int32, device="cuda") + state.write_to_cpu( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=layout, + ) + state.k_storage.zero_() + state.v_storage.zero_() + state.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[0], + token_indices=ids, + layout=layout, + ) + torch.cuda.synchronize() + assert (state.k_storage[:, ids] == 3).all() and (state.v_storage[:, ids] == 7).all() + finally: + torch.cuda.synchronize() + assert cuda.cudaHostUnregister(pointer) == 0 + del pages + memory.close() + + +@pytest.mark.parametrize("kind", ["gqa", "int8", "static_head", "mla", "packed_fp8"]) +@pytest.mark.parametrize("write", [True, False]) +def test_cpu_cache_transfer_rejects_cpu_state(kind: str, write: bool): + state = allocated_state(kind, device="cpu") + layout = transfer_layout(state) + pages = torch.zeros((1, layout.byte_end), dtype=torch.uint8) + ids = torch.tensor([0], dtype=torch.int64) + with pytest.raises(AssertionError, match="CUDA State storage"): + if write: + state.write_to_cpu( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=layout, + ) + else: + state.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=[0], + token_indices=ids, + layout=layout, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("token_count", [0, 1]) +def test_gqa_single_cpu_page_and_empty_transfer(token_count: int) -> None: + source = allocated_state("gqa", device="cuda") + target = allocated_state("gqa", device="cuda") + source.k_storage.fill_(3) + source.v_storage.fill_(7) + target.k_storage.fill_(19) + target.v_storage.fill_(19) + layout = transfer_layout(source) + pages = torch.full( + size=(3, layout.byte_end), + fill_value=91, + dtype=torch.uint8, + pin_memory=True, + ) + src_ids = torch.tensor(data=[5], dtype=torch.int32, device="cuda")[:token_count] + dst_ids = torch.tensor(data=[2], dtype=torch.int64, device="cuda")[:token_count] + page_indices = [1] if token_count else [] + + source.write_to_cpu( + token_indices=src_ids, + cpu_page_buffers=pages, + cpu_page_indices=page_indices, + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=pages, + cpu_page_indices=page_indices, + token_indices=dst_ids, + layout=layout, + ) + torch.cuda.synchronize() + if token_count: + assert (target.k_storage[:, 2] == 3).all() and (target.v_storage[:, 2] == 7).all() + assert (target.k_storage[:, [0, 1, 3, 4, 5, 6, 7]] == 19).all() + assert (target.v_storage[:, [0, 1, 3, 4, 5, 6, 7]] == 19).all() + else: + assert (target.k_storage == 19).all() and (target.v_storage == 19).all() + assert (pages[1] == 91).all() + assert (pages[[0, 2]] == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "kind", + [ + "gqa", + "mla", + "int4", + "int8", + "static_head", + "static_tensor", + "packed", + "packed_fp8", + "packed_glm", + ], +) +def test_full_state_interfaces_reject_cpu_token_indices(kind: str) -> None: + state = allocated_state(kind, device="cuda") + layout = transfer_layout(state) + cpu_ids = torch.tensor([0], dtype=torch.int32) + cuda_ids = torch.tensor([1], dtype=torch.int64, device="cuda") + cpu_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, pin_memory=True) + pd_pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.write_full_att_kv(state.config.logical_layer_groups[0][0], cpu_ids, torch.empty(0, device="cuda")) + for src, dst in ((cpu_ids, cuda_ids), (cuda_ids, cpu_ids)): + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.copy_full_att_tokens(src, dst) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.copy_tokens_from_dp(src, dst, state) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.write_to_cpu( + token_indices=cpu_ids, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + layout=layout, + ) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.read_from_cpu( + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + token_indices=cpu_ids, + layout=layout, + ) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.write_to_pd_page( + token_indices=cpu_ids, + pd_page_buffers=pd_pages, + pd_page_index=0, + layout=layout, + ) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.read_from_pd_page( + pd_page_buffers=pd_pages, + token_indices=cpu_ids, + pd_page_index=0, + layout=layout, + ) diff --git a/unit_tests/common/att_resources/test_full_att_dp_transfer.py b/unit_tests/common/att_resources/test_full_att_dp_transfer.py new file mode 100644 index 0000000000..78bc924405 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_dp_transfer.py @@ -0,0 +1,221 @@ +"""DP 直接索引复制：完整编码、真实 stride、非默认流和跨设备 IPC。""" + +from __future__ import annotations + +from multiprocessing.connection import Connection +from multiprocessing.reduction import ForkingPickler +import traceback + +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention_state.full_att import StaticFP8FullAttentionState +from lightllm.common.attention_state.full_att.dp_transfer import copy_gqa_tokens +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers, rebuild_full_state +from unit_tests.common.att_resources.test_capacity import definition + + +SCHEMES = ("gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_dp_narrow_head_fields_merge_before_copy(monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.common.attention_state.full_att import dp_transfer + + source = torch.empty(size=(2, 17, 4, 1), dtype=torch.bfloat16, device="cuda") + target = torch.empty(size=(2, 23, 4, 1), dtype=torch.bfloat16, device="cuda") + source.view(torch.uint8).random_(0, 256) + target.view(torch.uint8).fill_(173) + src_ids = torch.tensor([12, 0, 12], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([19, 1, 6], dtype=torch.int64, device="cuda") + expected = target.clone() + expected[:, dst_ids] = source[:, src_ids] + kernel = dp_transfer._copy_indexed_token_storages + launches = [] + + class CaptureLaunch: + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs): + launches.append((kwargs["SHAPES"], kwargs["ELEMENT_BYTES"])) + return kernel[grid](**kwargs) + + return launch + + monkeypatch.setattr(dp_transfer, "_copy_indexed_token_storages", CaptureLaunch()) + copy_gqa_tokens( + src_k_storage=source, + src_v_storage=source, + dst_k_storage=target, + dst_v_storage=torch.empty_like(target), + src_token_indices=src_ids, + dst_token_indices=dst_ids, + ) + torch.cuda.synchronize() + assert launches == [(((2, 1, 1), (2, 1, 1)), (8, 8))] + assert torch.equal(target.view(torch.uint8), expected.view(torch.uint8)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", SCHEMES) +@pytest.mark.parametrize("source_device", [0, 1]) +@pytest.mark.parametrize("token_count", [0, 1, 513]) +def test_dp_copy_preserves_complete_rows_without_temporary_payload( + kind: str, + source_device: int, + token_count: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + if source_device >= torch.cuda.device_count(): + pytest.skip("two CUDA devices required") + torch.cuda.set_device(0) + if source_device != 0: + if not torch.cuda.can_device_access_peer(0, source_device): + pytest.skip("CUDA peer access required") + # 仅测试直接跨卡 Tensor 的场景需要预先开 peer；生产通过 IPC 装配，不在复制时建映射。 + from cuda.bindings import runtime + + (error,) = runtime.cudaDeviceEnablePeerAccess(source_device, 0) + assert error in (runtime.cudaError_t.cudaSuccess, runtime.cudaError_t.cudaErrorPeerAccessAlreadyEnabled) + + template = definition(kind) + source = rebuild_full_state(template, token_capacity=1024, device=f"cuda:{source_device}") + target = rebuild_full_state(template, token_capacity=1200, device="cuda:0") + source_buffers = full_state_test_buffers(source) + target_buffers = full_state_test_buffers(target) + expected, original = {}, {} + for name, tensor in source_buffers.items(): + tensor.view(torch.uint8).random_(0, 256) + target_buffers[name].view(torch.uint8).fill_(173) + original[name] = tensor.cpu().view(torch.uint8) + expected[name] = target_buffers[name].cpu().view(torch.uint8) + + # 源索引允许重复，目标唯一；两端都带 storage offset 和非单位 stride。 + src_list = [(i * 29) % 257 for i in range(token_count)] + dst_list = [(i * 7) % 1200 for i in range(token_count)] + src_dtype = torch.int32 if source_device == 0 else torch.int64 + dst_dtype = torch.int64 if source_device == 0 else torch.int32 + src_ids = torch.empty(size=(token_count * 2 + 1,), dtype=src_dtype, device=f"cuda:{source_device}")[1::2] + dst_ids = torch.empty(size=(token_count * 3 + 2,), dtype=dst_dtype, device="cuda:0")[2::3] + src_ids.copy_(torch.tensor(src_list, dtype=src_dtype, device=f"cuda:{source_device}")) + dst_ids.copy_(torch.tensor(dst_list, dtype=dst_dtype, device="cuda:0")) + for name in expected: + expected[name][:, dst_list] = original[name][:, src_list] + calibration = [] + if isinstance(target, StaticFP8FullAttentionState): + for tensor in (target.scales, target.q_scales): + if tensor is not None: + calibration.append((tensor, tensor.clone())) + + # 显式依赖事件，不让测试偶然依赖默认流；跨设备源索引也必须等生产完成。 + producer_done = torch.cuda.Event() + producer_done.record(torch.cuda.current_stream(device=source_device)) + stream = torch.cuda.Stream(device=0) + stream.wait_event(producer_done) + stream.wait_stream(torch.cuda.current_stream(device=0)) + + def reject_intermediate(*args: object, **kwargs: object) -> None: + pytest.fail("DP copies must use supplied CUDA indices and storage directly") + + with monkeypatch.context() as patch: + for method in ("index_select", "index_copy_", "to", "tolist", "cpu", "contiguous"): + patch.setattr(torch.Tensor, method, reject_intermediate) + with torch.cuda.stream(stream): + target.copy_tokens_from_dp( + src_token_indices=src_ids, + dst_token_indices=dst_ids, + remote_dp_full_att_state=source, + ) + stream.synchronize() + for name, tensor in target_buffers.items(): + # 比较全部存储：包含未选槽位、HOLD 页、packed padding 和任意浮点位模式。 + assert torch.equal(tensor.cpu().view(torch.uint8), expected[name]) + assert torch.equal(source_buffers[name].cpu().view(torch.uint8), original[name]) + for tensor, unchanged in calibration: + assert torch.equal(tensor, unchanged) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("column_stride", [1, 2]) +@pytest.mark.parametrize("offset", [0, 1, 2, 4]) +def test_dp_copy_uses_real_storage_strides(column_stride: int, offset: int) -> None: + # 宽类型 view 的 8/4/2 字节对齐及不连续末轴回退；源目标层/token/head stride 不同。 + source_base = torch.empty(size=(2, 19, 3, 32), dtype=torch.bfloat16, device="cuda") + target_base = torch.empty(size=(2, 23, 3, 48), dtype=torch.bfloat16, device="cuda") + source_base.view(torch.uint8).random_(0, 256) + target_base.view(torch.uint8).fill_(173) + source = source_base[:, ::2, :, offset : offset + 8 * column_stride : column_stride] + target = target_base[:, ::3, :, offset : offset + 8 * column_stride : column_stride] + expected = target_base.clone() + expected[:, ::3, :, offset : offset + 8 * column_stride : column_stride][:, [4, 0, 6]] = source[:, [7, 2, 7]] + copy_gqa_tokens( + src_k_storage=source, + src_v_storage=source, + dst_k_storage=target, + dst_v_storage=torch.empty_like(target), + src_token_indices=torch.tensor([7, 2, 7], dtype=torch.int32, device="cuda"), + dst_token_indices=torch.tensor([4, 0, 6], dtype=torch.int32, device="cuda"), + ) + torch.cuda.synchronize() + assert torch.equal(target_base.view(torch.uint8), expected.view(torch.uint8)) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_dp_copy_through_production_peer_ipc(monkeypatch: pytest.MonkeyPatch) -> None: + from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor + + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip("CUDA peer access required") + states = [rebuild_full_state(definition(kind), token_capacity=16, device="cuda:1") for kind in SCHEMES] + for state in states: + for tensor in full_state_test_buffers(state).values(): + tensor.view(torch.uint8).random_(0, 256) + source_ids = torch.tensor([3, 99, 1, 99, 3], dtype=torch.int32, device="cuda:1")[::2] + torch.cuda.synchronize(device=1) + monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", reduce_tensor.__code__) + payload = bytes(ForkingPickler.dumps((states, source_ids))) + context = mp.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_ipc_copy_worker, args=(child, payload)) + process.start() + child.close() + try: + assert parent.poll(60), "DP peer worker did not finish" + error = parent.recv() + assert error is None, error + parent.send(None) + process.join(10) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + + +def _ipc_copy_worker(connection: Connection, payload: bytes) -> None: + try: + torch.cuda.set_device(0) + sources, src_ids = ForkingPickler.loads(payload) + for kind, source in zip(SCHEMES, sources): + target = rebuild_full_state(definition(kind), token_capacity=24, device="cuda:0") + target_buffers = full_state_test_buffers(target) + for tensor in target_buffers.values(): + tensor.view(torch.uint8).fill_(173) + expected = {name: tensor.cpu().view(torch.uint8) for name, tensor in target_buffers.items()} + for name, tensor in full_state_test_buffers(source).items(): + expected[name][:, [8, 2, 5]] = tensor.cpu().view(torch.uint8)[:, [3, 1, 3]] + target.copy_tokens_from_dp( + src_token_indices=src_ids, + dst_token_indices=torch.tensor([8, 2, 5], dtype=torch.int64, device="cuda:0"), + remote_dp_full_att_state=source, + ) + torch.cuda.synchronize() + for name, tensor in target_buffers.items(): + assert torch.equal(tensor.cpu().view(torch.uint8), expected[name]) + connection.send(None) + connection.recv() + except Exception: + connection.send(traceback.format_exc()) + finally: + connection.close() diff --git a/unit_tests/common/att_resources/test_full_att_transfer.py b/unit_tests/common/att_resources/test_full_att_transfer.py new file mode 100644 index 0000000000..8c11087762 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_transfer.py @@ -0,0 +1,338 @@ +"""各 full-attention 编码的统一 CPU/PD kernel：无中间 KV、跳步循环与字段完整性。""" + +from __future__ import annotations + +from dataclasses import replace +from functools import partial +from importlib import import_module +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.attention_state.full_att import StaticFP8FullAttentionState +from lightllm.common.attention_state.types import StateTransferLayout +from unit_tests.common.att_resources.model_helpers import ( + full_state_test_buffers, + full_state_test_fields, + rebuild_full_state, +) +from unit_tests.common.att_resources.test_capacity import definition + +if TYPE_CHECKING: + from easydict import EasyDict + from multiprocessing.connection import Connection + + +SCHEMES = { + "int4": "group_quantized", + "int8": "group_quantized", + "static_head": "static_fp8", + "static_tensor": "static_fp8", + "mla": "packed", + "packed": "packed_mla", + "packed_fp8": "packed_mla", + "packed_glm": "packed_mla", +} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", list(SCHEMES)) +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +@pytest.mark.parametrize("grid_num", [1, 36]) +@pytest.mark.parametrize("token_count", [1, 513]) +def test_full_att_pages_copy_all_fields_with_bounded_grid( + kind: str, + transfer: str, + grid_num: int, + token_count: int, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + page_tokens = 128 if transfer == "cpu" else 1024 + pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = page_tokens + template = definition(kind) + source = rebuild_full_state(template, token_capacity=1024, device="cuda") + target = rebuild_full_state(template, token_capacity=1024, device="cuda") + source_buffers = full_state_test_buffers(source) + target_buffers = full_state_test_buffers(target) + for buffer in source_buffers.values(): + buffer.view(torch.uint8).random_(0, 256) + for buffer in target_buffers.values(): + buffer.view(torch.uint8).fill_(173) + calibration = [] + if isinstance(source, StaticFP8FullAttentionState): + for state in (source, target): + for tensor in (state.scales, state.q_scales): + if tensor is not None: + calibration.append((tensor, tensor.clone())) + + regions = source.get_cpu_transfer_regions() if transfer == "cpu" else source.get_pd_transfer_regions() + if grid_num == 1: + # 调换层/token 轴，验证寻址使用真实 stride；region 顺序也不要求与存储成员相同。 + regions = tuple( + replace( + region, + axes=("token", "layer", *region.axes[2:]), + shape=(region.shape[1], region.shape[0], *region.shape[2:]), + ) + for region in reversed(regions) + ) + layout = StateTransferLayout( + state_name=source.state_name, + tokens_per_page=page_tokens, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + page_count = (token_count + page_tokens - 1) // page_tokens + page_indices = [3, 0, 4, 1, 6][:page_count] + pages = torch.full( + size=(8, layout.byte_end + 32), + fill_value=91, + dtype=torch.uint8, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", + ) + page_view = pages[:, 16:-16] + src_ids = torch.randperm(n=1024, device="cuda")[:token_count].to(dtype=torch.int32) + dst_order = torch.randperm(n=1024, device="cuda") + dst_ids = dst_order[:token_count] + expected = {name: tensor[:, dst_ids].clone() for name, tensor in target_buffers.items()} + source_fields = full_state_test_fields(source, source_buffers) + expected_fields = full_state_test_fields(target, expected) + for name, field in source_fields.items(): + expected_fields[name].tensor.copy_(other=field.tensor[:, src_ids]) + + # 直接覆盖 State 调用的入口参数，读写仍通过生产 State 接口提交。 + scheme = SCHEMES[kind] + state_module = import_module(name=type(source).__module__) + transfer_name = f"transfer_{scheme}_pages" + monkeypatch.setattr(state_module, transfer_name, partial(getattr(state_module, transfer_name), grid_num=grid_num)) + + def reject_gather(*args: object, **kwargs: object) -> None: + pytest.fail("page transfers must not gather/scatter KV or copy token indices to host") + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "index_select", reject_gather) + patch.setattr(torch.Tensor, "index_copy_", reject_gather) + patch.setattr(torch.Tensor, "tolist", reject_gather) + with torch.cuda.stream(stream): + if transfer == "cpu": + source.write_to_cpu( + token_indices=src_ids, + cpu_page_buffers=page_view, + cpu_page_indices=page_indices, + layout=layout, + ) + target.read_from_cpu( + cpu_page_buffers=page_view, + cpu_page_indices=page_indices, + token_indices=dst_ids, + layout=layout, + ) + else: + source.write_to_pd_page( + token_indices=src_ids, + pd_page_buffers=page_view, + pd_page_index=3, + layout=layout, + ) + target.read_from_pd_page( + pd_page_buffers=page_view, + token_indices=dst_ids, + pd_page_index=3, + layout=layout, + ) + stream.synchronize() + + # 比较整行原始字节：packed padding 沿用目标填充值，字段数据来自源存储。 + for name, tensor in target_buffers.items(): + assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected[name].view(torch.uint8)) + assert (tensor[:, dst_order[token_count:]].contiguous().view(torch.uint8) == 173).all() + assert (tensor[:, 1024:].view(torch.uint8) == 173).all() + for tensor, before in calibration: + assert torch.equal(tensor, before) + assert (pages[[index for index in range(8) if index not in page_indices]] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, -16:] == 91).all() + last_count = token_count - (page_count - 1) * page_tokens + for region in regions: + last_page = ( + page_view[page_indices[-1], region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + ) + unused = [slice(None)] * 4 + unused[region.axes.index("token")] = slice(last_count, None) + assert (last_page[tuple(unused)].contiguous().view(torch.uint8) == 91).all() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_all_full_att_schemes_transfer_through_peer_ipc( + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """各编码字段通过生产 IPC peer 映射跨 GPU 搬运，padding 保留目标值。""" + from multiprocessing.reduction import ForkingPickler + import torch.multiprocessing as mp + + from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor + + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip("CUDA peer access required") + pd_args.pd_kv_page_size = 4 + states, layouts, expected_buffers = [], [], [] + for kind in SCHEMES: + state = rebuild_full_state(definition(kind), token_capacity=8, device="cuda:1") + buffers = full_state_test_buffers(state) + for tensor in buffers.values(): + tensor.view(torch.uint8).random_(0, 256) + tensor[:, 5].view(torch.uint8).fill_(173) + expected = {name: tensor[:, 5:6].clone() for name, tensor in buffers.items()} + source_fields = full_state_test_fields(state, buffers) + expected_fields = full_state_test_fields(state, expected) + for name, field in source_fields.items(): + expected_fields[name].tensor.copy_(other=field.tensor[:, 1:2]) + regions = state.get_pd_transfer_regions() + layouts.append( + StateTransferLayout( + state_name=state.state_name, + tokens_per_page=4, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + ) + states.append(state) + expected_buffers.append(expected) + torch.cuda.synchronize(device=1) + monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", reduce_tensor.__code__) + payload = bytes(ForkingPickler.dumps(states)) + context = mp.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_peer_transfer_worker, args=(child, payload, layouts)) + process.start() + child.close() + try: + assert parent.poll(60), "PD peer worker did not finish" + error = parent.recv() + assert error is None, error + for state, expected in zip(states, expected_buffers): + for name, tensor in full_state_test_buffers(state).items(): + assert torch.equal(tensor[:, 5:6].contiguous().view(torch.uint8), expected[name].view(torch.uint8)) + parent.send(None) + process.join(10) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + + +def _peer_transfer_worker( + connection: Connection, + payload: bytes, + layouts: list[StateTransferLayout], +) -> None: + from multiprocessing.reduction import ForkingPickler + import traceback + + try: + torch.cuda.set_device(device=0) + states = ForkingPickler.loads(payload) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + source_ids = torch.tensor(data=[1], dtype=torch.int32, device="cuda:0") + target_ids = torch.tensor(data=[5], dtype=torch.int64, device="cuda:0") + for state, layout in zip(states, layouts): + pages = torch.empty(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda:0") + state.write_to_pd_page( + token_indices=source_ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + state.read_from_pd_page( + pd_page_buffers=pages, + token_indices=target_ids, + pd_page_index=1, + layout=layout, + ) + stream.synchronize() + connection.send(None) + connection.recv() + except Exception: + connection.send(traceback.format_exc()) + finally: + connection.close() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", ["gqa", *SCHEMES]) +@pytest.mark.parametrize( + "readies", [[False, True, False], [True, False, True], [True, True, True], [False, False, False]] +) +@pytest.mark.parametrize("grid_num", [1, 36]) +def test_cpu_ready_pages_preserve_fields_and_logical_token_positions( + kind: str, + readies: list[bool], + grid_num: int, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + pd_args.cpu_cache_token_page_size = 5 + state = rebuild_full_state(definition(kind), token_capacity=32, device="cuda") + # 单 program 验证跳过 ready 块后仍会处理后续未就绪块；默认 grid 保持生产配置。 + state_module = import_module(name=type(state).__module__) + scheme = "gqa" if kind == "gqa" else SCHEMES[kind] + transfer_name = f"transfer_{scheme}_pages" + monkeypatch.setattr(state_module, transfer_name, partial(getattr(state_module, transfer_name), grid_num=grid_num)) + for buffer in full_state_test_buffers(state).values(): + buffer.view(torch.uint8).random_(0, 256) + regions = state.get_cpu_transfer_regions() + layout = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=5, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + # 页编号不连续，token 索引带 stride；最后一页只写三个 token。 + page_indices = [3, 0, 4] + token_indices = torch.arange(26, dtype=torch.int32, device="cuda")[::2] + baseline = torch.full( + size=(6, layout.byte_end + 32), + fill_value=173, + dtype=torch.uint8, + pin_memory=True, + ) + pages = torch.full_like(baseline, 173, pin_memory=True) + state.write_to_cpu( + token_indices=token_indices, + cpu_page_buffers=baseline[:, 16:-16], + cpu_page_indices=page_indices, + layout=layout, + ) + torch.cuda.synchronize() + expected = pages.clone() + for page, ready in zip(page_indices, readies): + if not ready: + expected[page].copy_(baseline[page]) + + def reject_allocation(*args: object, **kwargs: object) -> None: + pytest.fail("all-ready writes must not allocate index or ready tensors") + + with monkeypatch.context() as patch: + if all(readies): + patch.setattr(torch, "tensor", reject_allocation) + state.write_to_cpu( + token_indices=token_indices, + cpu_page_buffers=pages[:, 16:-16], + cpu_page_indices=page_indices, + layout=layout, + cpu_page_ready_flags=readies, + ) + torch.cuda.synchronize() + # 整页比较同时覆盖 KV、量化 scale、packed indexer 字段及前后 padding。 + assert torch.equal(pages, expected) diff --git a/unit_tests/common/att_resources/test_full_state.py b/unit_tests/common/att_resources/test_full_state.py new file mode 100644 index 0000000000..1339ebb4b0 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_state.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + + +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention_state.full_att import ( + FullAttentionState, + GQAFullAttentionState, + MLAFullAttentionState, + GQAAttentionConfig, + MLAConfig, +) +from lightllm.common.attention_state.types import ( + AxisMapping, + StateTransferLayout, + TransferRegion, +) + + +def full_state(cls: type[FullAttentionState] = GQAFullAttentionState, device: str | torch.device | int = "cpu"): + if cls is MLAFullAttentionState: + config = MLAConfig( + logical_layer_groups=((2,), (5,)), + latent_dim=4, + latent_dtype=torch.float32, + ) + else: + config = GQAAttentionConfig( + logical_layer_groups=((2,), (5,)), + global_head_num=1, + head_dim=4, + dtype=torch.float32, + ) + prepared_full_state = cls(config, page_size=4, device=device) + prepared_full_state.allocate_storage(12) + return prepared_full_state + + +def layout_for(state: FullAttentionState): + regions = [] + offset = 16 + for name in full_state_test_fields(state, full_state_test_buffers(state)): + regions.append( + TransferRegion( + field=name, + byte_start=offset, + byte_end=offset + 128, + dtype=torch.float32, + axes=("token", "layer", "head", "dim"), + shape=(4, 2, 1, 4), + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=0, + length=1, + ), + AxisMapping( + axis="dim", + local_start=0, + external_start=0, + length=4, + ), + ), + ) + ) + offset += 128 + return StateTransferLayout( + state_name=state.state_name, + tokens_per_page=4, + byte_start=0, + byte_end=offset + 16, + regions=tuple(regions), + ) + + +@pytest.mark.parametrize("cls", [GQAFullAttentionState, MLAFullAttentionState]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_full_attention_write_copy_and_pd( + cls: type[FullAttentionState], device: str | torch.device | int, monkeypatch: pytest.MonkeyPatch +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + state = full_state(cls, device) + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + ids = torch.tensor([1, 3], dtype=torch.int32, device=device) + data = torch.arange(2 * read_kv_storage(state).shape[2] * 4, dtype=torch.float32, device=device).reshape( + 2, read_kv_storage(state).shape[2], 4 + ) + if device == "cuda": + state.write_full_att_kv(2, ids, data) + state.write_full_att_kv(5, ids, data + 100) + else: + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.write_full_att_kv(2, ids, data) + # CPU 仅构造复制和传输测试的数据，不执行推理 KV 写入。 + for layer, values in ((0, data), (1, data + 100)): + head_start = 0 + for buffer in full_state_test_buffers(state).values(): + head_end = head_start + buffer.shape[2] + buffer[layer, ids.long()] = values[:, head_start:head_end] + head_start = head_end + + def reject_host_list(tensor: torch.Tensor): + raise AssertionError("token copies must not convert indices to Python lists") + + if device == "cuda": + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "tolist", reject_host_list) + state.copy_full_att_tokens(ids, torch.tensor([8, 9], dtype=torch.int32, device=device)) + assert torch.equal(read_kv_storage(state)[:, [1, 3]], read_kv_storage(state)[:, [8, 9]]) + else: + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + state.copy_full_att_tokens(ids, torch.tensor([8, 9], dtype=torch.int32, device=device)) + layout = layout_for(state) + pages = torch.full( + size=(3, layout.byte_end), + fill_value=255, + dtype=torch.uint8, + device=device, + ) + if device == "cpu": + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + state.write_to_pd_page( + token_indices=ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + state.read_from_pd_page( + pd_page_buffers=pages, + token_indices=ids, + pd_page_index=1, + layout=layout, + ) + return + state.write_to_pd_page( + token_indices=ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + target = full_state(cls, device) + for storage_buffer in full_state_test_buffers(target).values(): + storage_buffer.zero_() + target.read_from_pd_page( + pd_page_buffers=pages, + token_indices=torch.tensor([5, 6], dtype=torch.long, device=device), + pd_page_index=1, + layout=layout, + ) + assert torch.equal(read_kv_storage(target)[:, [5, 6]], read_kv_storage(state)[:, [1, 3]]) + assert (pages[0] == 255).all() and (pages[2] == 255).all() + assert not read_kv_storage(target)[:, 0].any() + + +@pytest.mark.parametrize("state_offset, extra_page_bytes", [(0, 0), (32, 64)]) +def test_cpu_partial_last_page(state_offset: int, extra_page_bytes: int): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + state = full_state(device="cuda") + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.copy_(torch.arange(storage_buffer.numel()).reshape_as(storage_buffer)) + layout = layout_for(state) + layout = replace( + layout, + byte_start=layout.byte_start + state_offset, + byte_end=layout.byte_end + state_offset, + regions=tuple( + replace(r, byte_start=r.byte_start + state_offset, byte_end=r.byte_end + state_offset) + for r in layout.regions + ), + ) + buffer = torch.full( + size=(4, layout.byte_end + extra_page_bytes), + fill_value=252, + dtype=torch.uint8, + pin_memory=True, + ) + state.write_to_cpu( + token_indices=torch.tensor([0, 1, 2, 3, 7], dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[2, 0], + layout=layout, + ) + target = full_state(device="cuda") + for storage_buffer in full_state_test_buffers(target).values(): + storage_buffer.zero_() + target.read_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[2, 0], + token_indices=torch.tensor([6, 7, 8, 9, 10], dtype=torch.long, device="cuda"), + layout=layout, + ) + assert torch.equal(read_kv_storage(target)[:, [6, 7, 8, 9, 10]], read_kv_storage(state)[:, [0, 1, 2, 3, 7]]) + assert (buffer[[1, 3]] == 252).all() + assert (buffer[:, : layout.byte_start] == 252).all() and (buffer[:, layout.byte_end :] == 252).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_invalid_late_region_shape_cannot_partially_write_pd_page(): + state = full_state(device="cuda") + layout = layout_for(state) + region = layout.regions[-1] + region = replace(region, byte_end=region.byte_end - region.dtype.itemsize) + bad = replace(layout, regions=(*layout.regions[:-1], region)) + # 直接构造 Tensor 视图时，字节数与 shape 不匹配仍会在提交复制前失败。 + buffer = torch.full( + size=(layout.byte_end + 256,), + fill_value=53, + dtype=torch.uint8, + device="cuda", + ) + with pytest.raises(RuntimeError, match="shape"): + state.write_to_pd_page( + token_indices=torch.tensor([0], dtype=torch.long, device="cuda"), + pd_page_buffers=buffer.unsqueeze(0), + pd_page_index=0, + layout=bad, + ) + assert (buffer == 53).all() + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_dp_copies_all_layers_in_storage_order(device: str | torch.device | int): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + config = GQAAttentionConfig( + logical_layer_groups=((1,), (3,), (5,), (7,), (9,), (11,)), + global_head_num=1, + head_dim=4, + dtype=torch.float32, + ) + source = GQAFullAttentionState(config, page_size=4, device=device) + source.allocate_storage(12) + target = GQAFullAttentionState(config, page_size=4, device=device) + target.allocate_storage(12) + for storage_buffer in full_state_test_buffers(source).values(): + storage_buffer.copy_(torch.arange(storage_buffer.numel(), device=device).reshape_as(storage_buffer)) + for storage_buffer in full_state_test_buffers(target).values(): + storage_buffer.fill_(-1) + if device == "cpu": + with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): + target.copy_tokens_from_dp( + src_token_indices=torch.tensor([2, 0]), + dst_token_indices=torch.tensor([6, 7]), + remote_dp_full_att_state=source, + ) + return + target.copy_tokens_from_dp( + src_token_indices=torch.tensor([2, 0], dtype=torch.long, device=device), + dst_token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), + remote_dp_full_att_state=source, + ) + torch.testing.assert_close(read_kv_storage(target)[:, [6, 7]], read_kv_storage(source)[:, [2, 0]]) + assert (read_kv_storage(target)[:, 8:] == -1).all() + assert (read_kv_storage(target)[:, :6] == -1).all() + + # 拒绝将单层数据广播到多层，失败时目标 token 保持原值。 + single_layer_source = GQAFullAttentionState( + replace(config, logical_layer_groups=config.logical_layer_groups[:1]), page_size=4, device=device + ) + single_layer_source.allocate_storage(12) + for storage_buffer in full_state_test_buffers(single_layer_source).values(): + storage_buffer.fill_(9) + with pytest.raises(AssertionError, match="equivalent TP state configurations"): + target.copy_tokens_from_dp( + src_token_indices=torch.tensor([2, 0], dtype=torch.long, device=device), + dst_token_indices=torch.tensor([8, 9], dtype=torch.long, device=device), + remote_dp_full_att_state=single_layer_source, + ) + assert (read_kv_storage(target)[:, 8:] == -1).all() diff --git a/unit_tests/common/att_resources/test_full_storage_initialization.py b/unit_tests/common/att_resources/test_full_storage_initialization.py new file mode 100644 index 0000000000..26d4425ac7 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -0,0 +1,120 @@ +"""Real model initialization, compute writes, and profiled allocation.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +import json + +from unit_tests.common.att_resources.model_helpers import prepare_model_states + +import pytest +import torch + +from lightllm.common.attention_state.full_att import MLAFullAttentionState, PackedMLAFullAttentionState +from lightllm.utils.envs_utils import get_env_start_args +from unit_tests.common.att_resources.test_capacity import allocated_bytes +from unit_tests.server.router.model_infer.test_cpu_resource_layout import CASES, make_runtime, runtime + +if TYPE_CHECKING: + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from pathlib import Path + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("runtime"), +] + + +@pytest.mark.parametrize("kind,encoding,model_cls", CASES) +def test_model_storage_write_and_copy_preserve_all_fields( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], +): + _, (resources,) = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=model_cls, + tp=1, + ) + (state,) = resources.get_full_attention_states("main") + config = state.config + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + ids = torch.tensor([0, 4], dtype=torch.int32, device="cuda") + is_mla = isinstance(state, (MLAFullAttentionState, PackedMLAFullAttentionState)) + dim = config.latent_dim if is_mla else config.head_dim + kv = torch.randn( + size=(2, 1 if is_mla else 2 * config.head_num, dim), + device="cuda", + dtype=torch.bfloat16, + ) + state.write_full_att_kv(config.logical_layer_groups[-1][0], ids, kv) + if len(config.logical_layer_groups) > 1: + assert not read_kv_storage(state)[0].any() + resources.copy_full_att_tokens(ids, torch.tensor([1, 5], dtype=torch.int32, device="cuda")) + for name, field in full_state_test_fields(state, full_state_test_buffers(state)).items(): + if "token" in field.axes: + assert torch.equal( + field.tensor[:, ids].contiguous().view(torch.uint8), + field.tensor[:, [1, 5]].contiguous().view(torch.uint8), + ), name + assert state.capacity == resources.full_att_allocator.size == 16 + assert read_kv_storage(state).shape[1] == 20 # Includes the whole HOLD page. + + +@pytest.mark.parametrize("mode", ["normal", "prefill"]) +@pytest.mark.parametrize("kind,encoding,model_cls", CASES[:8]) +def test_auto_profile_fits_actual_gpu_storage_and_pd_pages( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + mode: str, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], +): + _, (reference,) = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=model_cls, + tp=1, + ) + (state,) = reference.get_full_attention_states("main") + args = get_env_start_args() + args.run_mode, args.pd_kv_page_size, args.pd_kv_page_num = mode, 8, 3 + transfer_bytes = reference.get_pd_operator().get_page_bytes() * 3 if mode == "prefill" else 0 + storage_bytes = allocated_bytes(state) + budget = storage_bytes + transfer_bytes + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (budget, 1 << 30)) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + model = model_cls.__new__(model_cls) + model.args, model.config = args, json.loads((tmp_path / "config.json").read_text()) + model.tp_rank_, model.tp_world_size_ = 0, 1 + model.max_total_token_num, model.max_req_num, model.mem_fraction = None, 1, 1.0 + prepare_model_states(model) + model.att_state_manager.finish_state_registration() + model.att_state_manager.load_full_att_calibration(args=args) + model.att_state_manager.initialize_storage( + args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction + ) + (state,) = model.att_state_manager.get_full_attention_states("main") + assert not hasattr(model, "full_att_state") + assert state.capacity == 16 + assert allocated_bytes(state) == storage_bytes + if mode == "prefill": + assert ( + model.att_state_manager.get_pd_operator().get_page_bytes() == reference.get_pd_operator().get_page_bytes() + ) diff --git a/unit_tests/common/att_resources/test_gqa_pd_kernel.py b/unit_tests/common/att_resources/test_gqa_pd_kernel.py new file mode 100644 index 0000000000..2870ca126e --- /dev/null +++ b/unit_tests/common/att_resources/test_gqa_pd_kernel.py @@ -0,0 +1,283 @@ +"""GQA PD 直接搬运：原始位模式、物理页选择、TP 切片及有效 token 范围。""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.types import StateTransferLayout + +if TYPE_CHECKING: + from easydict import EasyDict + from multiprocessing.connection import Connection + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("token_major", [False, True]) +@pytest.mark.parametrize("token_count", [1, 3]) +def test_gqa_pd_direct_copy_between_tp_layouts( + dtype: torch.dtype, + token_major: bool, + token_count: int, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + pd_args.pd_kv_page_size = 4 + # P 保存完整 heads，D 只恢复 rank 1 的一半 heads；两侧外部页格式相同。 + source = GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=((2,), (5,)), + global_head_num=4, + head_dim=128, + dtype=dtype, + ), + page_size=4, + device="cuda", + ) + target = GQAFullAttentionState( + layout_config=replace(source.config, tp_world_size=2, tp_rank=1), + page_size=4, + device="cuda", + ) + for state in (source, target): + state.allocate_storage(token_capacity=8) + for tensor in (source.k_storage, source.v_storage): + tensor.view(torch.uint8).random_(0, 256) + for tensor in (target.k_storage, target.v_storage): + tensor.view(torch.uint8).fill_(173) + layouts = [] + for state in (source, target): + regions = [] + for region in state.get_pd_transfer_regions(): + axes = ("token", "layer", "head", "dim") if token_major else region.axes + regions.append( + replace( + region, + byte_start=region.byte_start + 16, + byte_end=region.byte_end + 16, + axes=axes, + shape=tuple(region.shape[region.axes.index(axis)] for axis in axes), + ) + ) + layouts.append( + StateTransferLayout( + state_name=state.state_name, + tokens_per_page=4, + byte_start=16, + byte_end=max(region.byte_end for region in regions), + regions=tuple(regions), + ) + ) + layout, target_layout = layouts + pages = torch.full( + size=(3, layout.byte_end + 64), + fill_value=91, + dtype=torch.uint8, + device="cuda", + ) + page_view = pages[:, 32:-32] # 页池视图的真实跨度大于可见宽度。 + src_ids = torch.tensor(data=[6, 9, 1, 9, 7, 9], dtype=torch.int32, device="cuda")[::2][:token_count] + dst_ids = torch.tensor(data=[2, 6, 4], dtype=torch.int64, device="cuda")[:token_count] + expected = [ + tensor[:, src_ids, 2:].contiguous().view(torch.uint8) for tensor in (source.k_storage, source.v_storage) + ] + + def reject_gather(*args: object, **kwargs: object) -> None: + pytest.fail("PD transfer must not gather or scatter intermediate KV tensors") + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "index_select", reject_gather) + patch.setattr(torch.Tensor, "index_copy_", reject_gather) + patch.setattr(torch.Tensor, "tolist", reject_gather) + with torch.cuda.stream(stream): + source.write_to_pd_page( + token_indices=src_ids, + pd_page_buffers=page_view, + pd_page_index=1, + layout=layout, + ) + target.read_from_pd_page( + pd_page_buffers=page_view, + token_indices=dst_ids, + pd_page_index=1, + layout=target_layout, + ) + stream.synchronize() + for tensor, expected_bytes in zip((target.k_storage, target.v_storage), expected): + assert torch.equal(tensor[:, dst_ids].contiguous().view(torch.uint8), expected_bytes) + assert ( + tensor[:, [index for index in range(12) if index not in [2, 6, 4][:token_count]]] + .contiguous() + .view(torch.uint8) + == 173 + ).all() + assert (pages[[0, 2]] == 91).all() + assert (pages[:, :48] == 91).all() and (pages[:, -32:] == 91).all() + for region in layout.regions: + external = page_view[1, region.byte_start : region.byte_end].view(dtype).view(region.shape) + unused = [slice(None)] * 4 + unused[region.axes.index("token")] = slice(token_count, None) + assert (external[tuple(unused)].contiguous().view(torch.uint8) == 91).all() + + # 反向验证分片写出：rank 1 只修改全局 heads [2, 4)，rank 0 区域保持原始字节。 + pages.fill_(91) + target.write_to_pd_page( + token_indices=dst_ids, + pd_page_buffers=page_view, + pd_page_index=2, + layout=target_layout, + ) + for region, expected_bytes in zip(target_layout.regions, expected): + external = page_view[2, region.byte_start : region.byte_end].view(dtype).view(region.shape) + external = external.permute(tuple(region.axes.index(axis) for axis in ("layer", "token", "head", "dim"))) + assert torch.equal(external[:, :token_count, 2:].contiguous().view(torch.uint8), expected_bytes) + assert (external[:, :, :2].contiguous().view(torch.uint8) == 91).all() + assert (external[:, token_count:].contiguous().view(torch.uint8) == 91).all() + assert (pages[:2] == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_gqa_pd_head_slice_and_replica_reads(pd_args: EasyDict) -> None: + pd_args.pd_kv_page_size = 4 + # MQA 的 KV head 在 TP ranks 间复制，非首 rank 禁止写出但仍能读回。 + state = GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=1, + head_dim=128, + dtype=torch.bfloat16, + tp_world_size=2, + tp_rank=1, + ), + page_size=4, + device="cuda", + ) + state.allocate_storage(token_capacity=8) + state.k_storage.fill_(19) + state.v_storage.fill_(19) + regions = state.get_pd_transfer_regions() + layout = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=4, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + pages = torch.zeros(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda") + for region, value in zip(regions, (3, 7)): + pages[1, region.byte_start : region.byte_end].view(region.dtype).fill_(value) + snapshot = pages.clone() + ids = torch.tensor(data=[5, 1], dtype=torch.int32, device="cuda") + state.write_to_pd_page( + token_indices=ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + state.read_from_pd_page( + pd_page_buffers=pages, + token_indices=ids, + pd_page_index=1, + layout=layout, + ) + assert torch.equal(pages, snapshot) + assert (state.k_storage[:, ids] == 3).all() and (state.v_storage[:, ids] == 7).all() + assert (state.k_storage[:, [0, 2, 3, 4, 6, 7]] == 19).all() + assert (state.v_storage[:, [0, 2, 3, 4, 6, 7]] == 19).all() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +def test_gqa_pd_direct_copy_through_peer_ipc(pd_args: EasyDict, monkeypatch: pytest.MonkeyPatch) -> None: + """使用生产的 IPC peer 映射，从传输 GPU 直接读写另一块 GPU 上的 K/V。""" + from multiprocessing.reduction import ForkingPickler + import torch.multiprocessing as mp + + from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor + + if not torch.cuda.can_device_access_peer(0, 1): + pytest.skip("CUDA peer access required") + pd_args.pd_kv_page_size = 4 + state = GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=2, + head_dim=128, + dtype=torch.bfloat16, + ), + page_size=4, + device="cuda:1", + ) + state.allocate_storage(token_capacity=8) + state.k_storage.zero_() + state.v_storage.zero_() + state.k_storage[:, 1].fill_(3) + state.v_storage[:, 1].fill_(7) + torch.cuda.synchronize(device=1) + regions = state.get_pd_transfer_regions() + layout = StateTransferLayout( + state_name=state.state_name, + tokens_per_page=4, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", reduce_tensor.__code__) + payload = bytes(ForkingPickler.dumps(state)) + context = mp.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_peer_pd_worker, args=(child, payload, layout)) + process.start() + child.close() + try: + assert parent.poll(45), "PD peer worker did not finish" + error = parent.recv() + assert error is None, error + assert (state.k_storage[:, 5] == 3).all() and (state.v_storage[:, 5] == 7).all() + parent.send(None) + process.join(10) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + + +def _peer_pd_worker(connection: Connection, payload: bytes, layout: StateTransferLayout) -> None: + from multiprocessing.reduction import ForkingPickler + import traceback + + try: + torch.cuda.set_device(device=0) + state = ForkingPickler.loads(payload) + pages = torch.zeros(size=(2, layout.byte_end), dtype=torch.uint8, device="cuda:0") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + state.write_to_pd_page( + token_indices=torch.tensor(data=[1], dtype=torch.int32, device="cuda:0"), + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + state.read_from_pd_page( + pd_page_buffers=pages, + token_indices=torch.tensor(data=[5], dtype=torch.int64, device="cuda:0"), + pd_page_index=1, + layout=layout, + ) + stream.synchronize() + assert not pages[0].any() + connection.send(None) + connection.recv() # 保持导入的存储存活，直到父进程检查原 GPU 上的数据。 + except Exception: + connection.send(traceback.format_exc()) + finally: + connection.close() diff --git a/unit_tests/common/att_resources/test_gqa_transfer.py b/unit_tests/common/att_resources/test_gqa_transfer.py new file mode 100644 index 0000000000..ac19dc39ab --- /dev/null +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -0,0 +1,192 @@ +"""GQA 共用搬运 kernel 在受限 grid 下完整覆盖多轮、跨页和尾块数据。""" + +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.full_att.gqa.gqa_transfer import transfer_gqa_pages +from lightllm.common.attention_state.types import StateTransferLayout + +if TYPE_CHECKING: + from easydict import EasyDict + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("grid_num", [1, 7, 36]) +@pytest.mark.parametrize("page_device", ["cpu", "cuda"]) +def test_gqa_grid_stride_covers_all_tokens(grid_num: int, page_device: str, pd_args: "EasyDict") -> None: + pd_args.pd_kv_page_size = 32 + config = GQAAttentionConfig( + logical_layer_groups=((0,), (2,), (4,)), + global_head_num=3, + head_dim=128, + dtype=torch.bfloat16, + ) + source = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") + target = GQAFullAttentionState(layout_config=config, page_size=32, device="cuda") + source.allocate_storage(token_capacity=128) + target.allocate_storage(token_capacity=128) + for tensor in (source.k_storage, source.v_storage): + tensor.view(torch.uint8).random_(0, 256) + for tensor in (target.k_storage, target.v_storage): + tensor.view(torch.uint8).fill_(173) + regions = source.get_pd_transfer_regions() + layout = StateTransferLayout( + state_name=source.state_name, + tokens_per_page=32, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + # 3 layers × 97 tokens × 3 heads × 128 dims，超过 36×256，且最后一块只有 128 个元素。 + token_count = 97 + src_ids = torch.randperm(n=128, device="cuda")[:token_count].to(dtype=torch.int32) + dst_order = torch.randperm(n=128, device="cuda") + dst_ids = dst_order[:token_count] + pages = torch.full( + size=(6, layout.byte_end + 32), + fill_value=91, + dtype=torch.uint8, + device=page_device, + pin_memory=page_device == "cpu", + ) + page_view = pages[:, 16:-16] + page_indices = [3, 0, 4, 1] + transfer_gqa_pages( + k_storage=source.k_storage, + v_storage=source.v_storage, + token_indices=src_ids, + page_buffers=page_view, + page_indices=page_indices, + layout=layout, + write=True, + grid_num=grid_num, + ) + transfer_gqa_pages( + k_storage=target.k_storage, + v_storage=target.v_storage, + token_indices=dst_ids, + page_buffers=page_view, + page_indices=page_indices, + layout=layout, + write=False, + grid_num=grid_num, + ) + torch.cuda.synchronize() + for src, dst in ((source.k_storage, target.k_storage), (source.v_storage, target.v_storage)): + assert torch.equal( + dst[:, dst_ids].contiguous().view(torch.uint8), src[:, src_ids].contiguous().view(torch.uint8) + ) + assert (dst[:, dst_order[token_count:]].contiguous().view(torch.uint8) == 173).all() + assert (dst[:, 128:].view(torch.uint8) == 173).all() + assert (pages[[2, 5]] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, -16:] == 91).all() + for region in regions: + last_page = page_view[1, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + assert (last_page[:, 1:].contiguous().view(torch.uint8) == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("page_device", ["cpu", "cuda"]) +@pytest.mark.parametrize("num_stages", [1, 3]) +@pytest.mark.parametrize( + "byte_start, extra_bytes, dim_stride, element_bytes", + [ + (8, 32, 1, 8), + (2, 32, 1, 2), + (4, 32, 1, 4), + (8, 34, 1, 2), + (8, 36, 1, 4), + (8, 32, 2, 2), + ], +) +def test_wide_transfer_preserves_unaligned_pages_and_strided_fields( + page_device: str, + num_stages: int, + byte_start: int, + extra_bytes: int, + dim_stride: int, + element_bytes: int, + pd_args: "EasyDict", + monkeypatch: pytest.MonkeyPatch, +) -> None: + from lightllm.common.attention_state.full_att import gqa_transfer + + pd_args.pd_kv_page_size = 2 + config = GQAAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=2, + head_dim=128, + dtype=torch.bfloat16, + ) + source = GQAFullAttentionState(layout_config=config, page_size=2, device="cuda") + target = GQAFullAttentionState(layout_config=config, page_size=2, device="cuda") + source.allocate_storage(token_capacity=8) + target.allocate_storage(token_capacity=8) + for name in ("k_storage", "v_storage"): + storage = getattr(source, name) + if dim_stride != 1: + storage = torch.empty(size=(*storage.shape[:-1], 256), dtype=storage.dtype, device="cuda") + storage.view(torch.uint8).random_(0, 256) + setattr(source, name, storage[..., ::dim_stride]) + getattr(target, name).view(torch.uint8).fill_(173) + regions = source.get_pd_transfer_regions() + layout = StateTransferLayout( + state_name=source.state_name, + tokens_per_page=2, + byte_start=0, + byte_end=max(region.byte_end for region in regions), + regions=regions, + ) + # 页起点、页跨度和最内维 stride 分别破坏宽类型对齐，回退必须保持完整字段及保护区。 + pages = torch.full( + size=(3, layout.byte_end + extra_bytes), + fill_value=91, + dtype=torch.uint8, + device=page_device, + pin_memory=page_device == "cpu", + ) + page_view = pages[:, byte_start : byte_start + layout.byte_end] + src_ids = torch.tensor(data=[3, 3, 1, 1, 7, 7], dtype=torch.int32, device="cuda")[::2] + dst_ids = torch.tensor(data=[6, 2, 4], dtype=torch.int64, device="cuda") + expected = {} + for name in ("k_storage", "v_storage"): + expected[name] = getattr(target, name).clone() + expected[name][:, dst_ids] = getattr(source, name)[:, src_ids] + + actual_kernel = gqa_transfer._copy_gqa_pages + write_units = [] + + class CaptureLaunch: + def __getitem__(self, grid: tuple[int]): + def launch(**kwargs): + if kwargs["WRITE"]: + write_units.append(kwargs["ELEMENT_BYTES"]) + return actual_kernel[grid](**kwargs) + + return launch + + monkeypatch.setattr(gqa_transfer, "_copy_gqa_pages", CaptureLaunch()) + for state, ids, write in ((source, src_ids, True), (target, dst_ids, False)): + transfer_gqa_pages( + k_storage=state.k_storage, + v_storage=state.v_storage, + token_indices=ids, + page_buffers=page_view, + page_indices=[2, 0], + layout=layout, + write=write, + num_stages=num_stages, + ) + torch.cuda.synchronize() + assert write_units == [element_bytes, element_bytes] + for name in ("k_storage", "v_storage"): + assert torch.equal(getattr(target, name).view(torch.uint8), expected[name].view(torch.uint8)) + assert (pages[1] == 91).all() + assert (pages[:, :byte_start] == 91).all() + assert (pages[:, byte_start + layout.byte_end :] == 91).all() + for region in regions: + tail_page = page_view[0, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + assert (tail_page[:, 1:].contiguous().view(torch.uint8) == 91).all() diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py new file mode 100644 index 0000000000..15d29bd85d --- /dev/null +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -0,0 +1,117 @@ +import pytest + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.common.attention_state.bounded_att import BoundedAttentionState +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("draft_layers", [0, 1, 3]) +@pytest.mark.parametrize( + "layer_types,expected_indices,linear_layers", + [ + ( + ("linear_attention", "linear_attention", "linear_attention", "full_attention") * 2, + (0, 1, 2, 0, 3, 4, 5, 1), + 6, + ), + ( + ( + "deepseek_sparse_attention", + "linear_attention", + "linear_attention", + "deepseek_sparse_attention", + "linear_attention", + ), + (0, 0, 1, 1, 2), + 3, + ), + ( + ("linear_attention", "full_attention", "linear_attention") * 2 + ("full_attention",), + (0, 0, 1, 2, 1, 3, 2), + 4, + ), + ], +) +def test_hybrid_layer_mapping_and_mtp_states( + layer_types: tuple[str, ...], + expected_indices: tuple[int, ...], + linear_layers: int, + draft_layers: int, +) -> None: + model_config = dict( + num_hidden_layers=len(layer_types), + hidden_size=128, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + linear_num_key_heads=1, + linear_num_value_heads=1, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + full_attention_interval=4, + layer_types=layer_types, + ) + args = StartArgs(data_type="bfloat16", linear_att_ssm_data_type="float32") + config = Qwen3NextResourceConfig.from_model_config(model_config, args) + logical_layer_ids = tuple(range(len(layer_types))) + (full_state,) = Qwen3NextTpPartModel.create_full_att_states( + model_config=model_config, + args=args, + model_logical_layer_ids=logical_layer_ids, + ) + (bounded_state,) = Qwen3NextTpPartModel.create_bounded_att_states( + model_config=model_config, + args=args, + model_logical_layer_ids=logical_layer_ids, + device="meta", + ) + manager = AttentionStateManager(request_capacity=1) + manager.register_state(model_role="main", state=full_state) + manager.register_state(model_role="main", state=bounded_state) + # draft 各自创建独立 State，不再把其层数和物理行追加到 main 的配置中。 + for draft_index in range(draft_layers): + (state,) = Qwen3_5MTPModel.create_full_att_states( + model_config=model_config, + args=args, + model_logical_layer_ids=(len(layer_types) + draft_index,), + ) + manager.register_state(model_role=f"draft{draft_index}", state=state) + manager.finish_state_registration() + + assert config.all_layer_num == len(layer_types) + full_layers = len(layer_types) - linear_layers + assert config.get_full_att_layer_indices() == tuple( + layer for layer, kind in enumerate(layer_types) if kind != "linear_attention" + ) + assert len(full_state.config.logical_layer_groups) == full_layers + assert bounded_state.conv_state.shape[0] == bounded_state.ssm_state.shape[0] == linear_layers + assert sum(len(state.config.logical_layer_groups) for state in manager.get_full_attention_states()) == ( + full_layers + draft_layers + ) + for layer, (kind, expected) in enumerate(zip(layer_types, expected_indices)): + is_linear = kind == "linear_attention" + state = manager.get_state( + model_role="main", + state_type_filter="", + logical_layer_id=layer, + state_class=BoundedAttentionState if is_linear else FullAttentionState, + ) + if is_linear: + conv, ssm = state.get_layer_view(logical_layer_id=layer) + assert conv.storage_offset() == expected * state.conv_state.stride(0) + assert ssm.storage_offset() == expected * state.ssm_state.stride(0) + else: + assert state.get_storage_layer_index(logical_layer_id=layer) == expected + for draft in range(draft_layers): + state = manager.get_state( + model_role=f"draft{draft}", + state_type_filter="gqa", + logical_layer_id=len(layer_types) + draft, + ) + assert state is not full_state + assert state.get_storage_layer_index(logical_layer_id=len(layer_types) + draft) == 0 diff --git a/unit_tests/common/att_resources/test_layer_groups.py b/unit_tests/common/att_resources/test_layer_groups.py new file mode 100644 index 0000000000..0f55fe5986 --- /dev/null +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -0,0 +1,164 @@ +"""逻辑层共享物理 KV：容量、视图、标定索引和外部页传输保持一致。""" + +from dataclasses import is_dataclass, replace +from typing import TYPE_CHECKING + +import pytest +import torch +from torch.utils._pytree import tree_leaves + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.models.gemma4.model import Gemma4TpPartModel +from unit_tests.common.att_resources.model_helpers import ( + full_state_test_buffers, + full_state_test_fields, + rebuild_full_state, + replace_logical_layer_groups, +) +from unit_tests.common.att_resources.test_capacity import definition, allocated_bytes + +if TYPE_CHECKING: + from easydict import EasyDict + + +SCHEMES = ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] + + +@pytest.mark.parametrize("kind", SCHEMES) +def test_shared_layers_use_two_physical_stores(kind: str, pd_args: "EasyDict") -> None: + template = definition(kind) + config = replace_logical_layer_groups(template.config, ((0, 4), (8, 11))) + state = rebuild_full_state( + template=template, + config=config, + token_capacity=8, + device="cpu", + ) + assert all(state.contains_logical_layer(logical_layer_id) for logical_layer_id in (0, 4, 8, 11)) + assert not state.contains_logical_layer(1) + assert all(buffer.shape[0] == 2 for buffer in full_state_test_buffers(state).values()) + assert state.get_token_bytes() == template.get_token_bytes() + assert allocated_bytes(state, token_only=True) == state.get_token_bytes() * (8 + state.page_size) + + for getter in (state.get_cpu_transfer_regions, state.get_pd_transfer_regions): + assert all(region.shape[region.axes.index("layer")] == 2 for region in getter()) + + # 同组返回相同地址，不同组返回独立数据；包括量化 scale 和 packed indexer 的计算视图。 + views = [] + for layer_id in (0, 4, 8, 11): + view = state.get_layer_view(logical_layer_id=layer_id) + if is_dataclass(view): + view = tuple(vars(view).values()) + views.append([(tensor.data_ptr(), tensor.shape, tensor.stride()) for tensor in tree_leaves(view)]) + assert views[0] == views[1] + assert views[2] == views[3] + assert views[0] != views[2] + + assert [state.get_storage_layer_index(logical_layer_id=layer) for layer in (0, 4, 8, 11)] == [0, 0, 1, 1] + + +@pytest.mark.parametrize("kind", SCHEMES) +@pytest.mark.parametrize("groups", [((),), ((0,), (0,)), ((0, 0),), (0, 4)]) +def test_invalid_layer_groups_are_rejected(kind: str, groups: tuple) -> None: + with pytest.raises(AssertionError): + replace_logical_layer_groups(definition(kind).config, groups) + + +def test_model_local_shared_groups_preserve_global_logical_layer_ids(pd_args: "EasyDict") -> None: + pd_args.data_type = "bfloat16" + config = dict( + num_hidden_layers=6, + num_attention_heads=4, + num_key_value_heads=2, + hidden_size=512, + head_dim=128, + global_head_dim=128, + num_kv_shared_layers=2, + layer_types=["sliding_attention", "full_attention"] * 3, + ) + (state,) = Gemma4TpPartModel.create_full_att_states( + model_config=config, + args=pd_args, + model_logical_layer_ids=(10, 11, 12, 13, 14, 15), + ) + # 共享消费者 14/15 分别使用最后一个同类型生产层 12/13，保持四份物理层存储。 + assert state.config.logical_layer_groups == ((10,), (11,), (12, 14), (13, 15)) + latent_dtype = state.config.latent_dtype if hasattr(state.config, "latent_dtype") else state.config.dtype + assert state.get_token_bytes() == 4 * 2 * 2 * 128 * latent_dtype.itemsize + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", SCHEMES) +@pytest.mark.parametrize("transfer", ["cpu", "pd"]) +def test_shared_layers_transfer_each_physical_layer_once( + kind: str, + transfer: str, + pd_args: "EasyDict", +) -> None: + pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 + template = definition(kind) + source = rebuild_full_state( + template=template, + config=replace_logical_layer_groups(template.config, ((0, 4), (8,))), + token_capacity=8, + device="cuda", + ) + target = rebuild_full_state( + template=template, + config=replace_logical_layer_groups(template.config, ((20, 24), (28,))), + token_capacity=8, + device="cuda", + ) + source_buffers = full_state_test_buffers(source) + target_buffers = full_state_test_buffers(target) + for buffer in source_buffers.values(): + buffer.view(torch.uint8).random_(0, 256) + for buffer in target_buffers.values(): + buffer.zero_() + managers = [] + for state in (source, target): + manager = AttentionStateManager(request_capacity=1) + manager.register_state(model_role="main", state=state) + manager.finish_state_registration() + managers.append(manager) + src_ids = torch.tensor([0, 3, 6], dtype=torch.int32, device="cuda") + dst_ids = torch.tensor([1, 4, 7], dtype=torch.int32, device="cuda") + if transfer == "cpu": + pages = torch.zeros( + (2, managers[0].get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) + managers[0].get_cpu_cache_operator().write_req( + token_indices=src_ids, + cpu_page_buffers=pages, + cpu_page_indices=[1], + layouts=managers[0].get_cpu_cache_operator().get_page_layouts(), + ) + managers[1].get_cpu_cache_operator().read_req( + cpu_page_buffers=pages, + cpu_page_indices=[1], + token_indices=dst_ids, + layouts=managers[1].get_cpu_cache_operator().get_page_layouts(), + ) + else: + pages = torch.zeros((2, managers[0].get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + managers[0].get_pd_operator().write_page( + pd_page_buffers=pages, + pd_page_index=1, + kind="kv", + token_indices=src_ids, + layouts=managers[0].get_pd_operator().get_pd_full_att_page_layouts(), + ) + managers[1].get_pd_operator().read_page( + pd_page_buffers=pages, + pd_page_index=1, + kind="kv", + token_indices=dst_ids, + layouts=managers[1].get_pd_operator().get_pd_full_att_page_layouts(), + ) + torch.cuda.synchronize() + source_fields = full_state_test_fields(source, source_buffers) + target_fields = full_state_test_fields(target, target_buffers) + for name, field in source_fields.items(): + expected = field.tensor[:, src_ids].contiguous().view(torch.uint8) + actual = target_fields[name].tensor[:, dst_ids].contiguous().view(torch.uint8) + assert torch.equal(actual, expected) diff --git a/unit_tests/common/att_resources/test_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py new file mode 100644 index 0000000000..c28186809b --- /dev/null +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -0,0 +1,216 @@ +"""按角色、层归属及方案查询资源，不在推理层保存绑定状态。""" + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att import GQAAttentionConfig, FullAttentionState, GQAFullAttentionState +from lightllm.common.attention_state.bounded_att import KDAState, GDNState +from unit_tests.common.att_resources.test_bounded_state import make_state +from unit_tests.common.att_resources.test_window_compressor_state import make_window +from unit_tests.common.att_resources.test_capacity import definition +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.common.basemodel.layer_infer.transformer_layer_infer import TransformerLayerInfer + + +def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: + model = LlamaTpPartModel.__new__(LlamaTpPartModel) + model.args, model.config = SimpleNamespace(), {} + model.tp_rank_, model.tp_world_size_ = 0, 1 + model.layers_infer = [SimpleNamespace(layer_num_=0, indexer=SimpleNamespace())] + model.max_req_num = 0 + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + state = GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=128, + dtype=torch.bfloat16, + ), + device="meta", + ) + monkeypatch.setattr( + LlamaTpPartModel, + "create_full_att_states", + classmethod(lambda cls, **kwargs: [state]), + ) + + model.prepare_attention_states() + + assert model.att_state_manager.main_gqa_bf16_state is state + assert not hasattr(model.layers_infer[0], "full_att_state_name") + assert not hasattr(model.layers_infer[0].indexer, "full_att_state_name") + + +@pytest.mark.parametrize("role", ["main", "draft0"]) +def test_query_selects_role_and_scheme_and_preserves_shared_layers(role: str, monkeypatch: pytest.MonkeyPatch) -> None: + manager = AttentionStateManager(request_capacity=1) + for model_role in ("main", "draft0"): + for dtype, groups in ((torch.bfloat16, ((0, 2),)), (torch.float16, ((3,),))): + manager.register_state( + model_role, + GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=groups, + global_head_num=1, + head_dim=128, + dtype=dtype, + ), + device="meta", + ), + ) + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.is_mtp_draft_model = role != "main" + model.mtp_previous_draft_models = [] + model.max_req_num = 0 + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + model.main_model = SimpleNamespace(att_state_manager=manager) + monkeypatch.setattr("lightllm.common.basemodel.basemodel.AttentionStateManager", lambda request_capacity: manager) + model.config, model.args = {}, SimpleNamespace(mtp_mode=None) + model.tp_rank_, model.tp_world_size_ = 0, 1 + model.layers_infer = [] + for logical_layer_id in range(4): + layer = TransformerLayerInfer.__new__(TransformerLayerInfer) + layer.layer_num_ = logical_layer_id + model.layers_infer.append(layer) + + model.prepare_attention_states() + + assert all(not hasattr(layer, "full_att_state_name") for layer in model.layers_infer) + manager.finish_state_registration() + state = manager.get_state(role, "gqa", 0) + assert state.state_name == f"{role}_gqa_bf16_state" + assert manager.get_state(role, "gqa", 2) is state + assert manager.get_state(role, "gqa", 3).state_name == f"{role}_gqa_fp16_state" + + +@pytest.mark.parametrize("finalized", [False, True]) +def test_prefix_and_logical_layer_select_quantized_state(finalized: bool) -> None: + from dataclasses import replace + + manager = AttentionStateManager(request_capacity=1) + plain = definition("gqa") + quantized = definition("int8") + quantized = type(quantized)( + layout_config=replace_logical_layer_groups(quantized.config, ((10, 12),)), + device="meta", + ) + manager.register_state("main", plain) + manager.register_state("main", quantized) + manager.register_state("draft0", definition("gqa")) + logical_layer_id = plain.config.logical_layer_groups[0][0] + if not finalized: + with pytest.raises(AssertionError, match="State registration must finish"): + manager.get_state("main", "gqa", logical_layer_id) + manager.finish_state_registration() + + assert manager.get_state("main", "gqa", logical_layer_id) is plain + assert manager.get_state("main", "gqa", 10) is quantized + assert manager.get_state("main", "int8", 10) is quantized + assert manager.get_state("main", "gqa", 12) is quantized + assert manager.get_state("draft0", "gqa", logical_layer_id) is not plain + with pytest.raises(AssertionError, match="expected exactly one"): + manager.get_state("main", "gqa", 99) + + +def test_selection_requires_finished_creation_and_reuses_cache() -> None: + manager = AttentionStateManager(request_capacity=1) + plain = definition("gqa") + manager.register_state("main", plain) + logical_layer_id = plain.config.logical_layer_groups[0][0] + with pytest.raises(AssertionError, match="State registration must finish"): + manager.get_state("main", "gqa", logical_layer_id) + manager.finish_state_registration() + assert manager.get_state("main", "gqa", logical_layer_id) is plain + hits = manager.get_state.cache_info().hits + assert manager.get_state("main", "gqa", logical_layer_id) is plain + assert manager.get_state.cache_info().hits == hits + 1 + + +def test_selection_rejects_ambiguous_finished_members() -> None: + manager = AttentionStateManager(request_capacity=1) + plain = definition("gqa") + manager.register_state("main", plain) + manager.register_state("main", definition("int8")) + manager.finish_state_registration() + with pytest.raises(AssertionError, match="expected exactly one"): + manager.get_state("main", "gqa", plain.config.logical_layer_groups[0][0]) + + +@pytest.mark.parametrize("state_class", [KDAState, GDNState]) +def test_bounded_state_layer_ownership(state_class: type[KDAState] | type[GDNState]) -> None: + state = make_state(cls=state_class) + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", state) + manager.finish_state_registration() + assert state.contains_logical_layer(1) and state.contains_logical_layer(3) + assert not state.contains_logical_layer(0) + assert manager.get_state("main", state.get_state_type(), 3) is state + + +def test_mla_selection_excludes_same_layer_tail() -> None: + from dataclasses import replace + + manager = AttentionStateManager(request_capacity=3) + tail = make_window() + template = definition("mla") + state = type(template)( + layout_config=replace_logical_layer_groups(template.config, ((3, 7),)), + device="meta", + ) + manager.register_state("main", state) + manager.register_state("main", tail) + manager.finish_state_registration() + assert tail.contains_logical_layer(3) and tail.contains_logical_layer(7) + assert not tail.contains_logical_layer(0) + assert manager.get_state("main", "mla", 7) is state + assert manager.get_state("main", "compressor", 7) is tail + assert ( + manager.get_state( + model_role="main", + state_type_filter="", + logical_layer_id=7, + state_class=FullAttentionState, + ) + is state + ) + with pytest.raises(AssertionError, match="expected exactly one"): + manager.get_state("main", "", 7) + + +@pytest.mark.parametrize("role", ["draft1", "draft10"]) +def test_query_matches_complete_model_role(role: str) -> None: + manager = AttentionStateManager(request_capacity=1) + for model_name in ("draft1", "draft10"): + manager.register_state(model_name, definition("gqa")) + manager.finish_state_registration() + assert manager.get_state(role, "gqa", 0).state_name == f"{role}_gqa_bf16_state" + + +@pytest.mark.parametrize("algorithm", ["kda", "gdn", "compressor"]) +def test_precision_variants_have_distinct_names_and_can_coexist(algorithm: str) -> None: + from dataclasses import replace + + manager = AttentionStateManager(request_capacity=3) + template = ( + make_window() if algorithm == "compressor" else make_state(cls=KDAState if algorithm == "kda" else GDNState) + ) + states = [] + for dtype in (torch.bfloat16, torch.float32): + config = replace(template.config, **({"dtype": dtype} if algorithm == "compressor" else {"conv_dtype": dtype})) + state = type(template)(config=config, request_capacity=3, device="cpu") + manager.register_state("main", state) + states.append(state) + manager.finish_state_registration() + logical_layer_id = template.config.logical_layer_ids[0] + for state in states: + assert manager.get_state("main", state.get_state_type(), logical_layer_id) is state + assert getattr(manager, state.state_name) is state + assert states[0].state_name != states[1].state_name + with pytest.raises(AssertionError, match="expected exactly one"): + manager.get_state("main", algorithm, logical_layer_id) diff --git a/unit_tests/common/att_resources/test_manager_helpers.py b/unit_tests/common/att_resources/test_manager_helpers.py new file mode 100644 index 0000000000..990559a072 --- /dev/null +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention_state.allocators import FullAttKVAllocator, StatePageAllocator +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.types import PageKind, StateTransferLayout +from unit_tests.common.att_resources.manager_fixtures import MainDraftTestResourceManager +from unit_tests.common.att_resources.test_full_state import full_state, layout_for +from unit_tests.common.att_resources.test_bounded_state import page_layout + +if TYPE_CHECKING: + from easydict import EasyDict + + +@pytest.fixture(autouse=True) +def manager_cpu_page_args(allocator_environment: None, monkeypatch: pytest.MonkeyPatch) -> None: + """本模块手工构造的 CPU 联合页使用 4-token 规格，启动参数须与布局一致。""" + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setattr(get_env_start_args(), "cpu_cache_token_page_size", 4) + + +def composition(device: str | torch.device | int = "cpu"): + return MainDraftTestResourceManager(device) + + +def layouts_for(manager: AttentionStateManager, kind: str = "cpu"): + layouts = [] + if kind != "att_state": + layouts += [layout_for(state) for state in (manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)] + if kind != "kv": + layouts += [ + page_layout(state) + for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state) + ] + + offset = 0 + result = {} + for layout in layouts: + state_end = offset + layout.byte_end + result[layout.state_name] = replace( + layout, + byte_start=offset + layout.byte_start, + byte_end=state_end, + tokens_per_page=None if kind == "att_state" else 4, + regions=tuple( + replace(r, byte_start=r.byte_start + offset, byte_end=r.byte_end + offset) for r in layout.regions + ), + ) + offset = state_end + return result + + +def init_source(manager: AttentionStateManager): + for i, state in enumerate((manager.main_gqa_fp32_state, manager.draft0_mla_fp32_state)): + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.copy_( + torch.arange(storage_buffer.numel(), device=storage_buffer.device).reshape_as(storage_buffer) + i + ) + for i, state in enumerate((manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)): + state.conv_state[:, 0].fill_(11 + i) + state.conv_state[:, 1].fill_(23 + i) + state.ssm_state[:, :3].fill_(31 + i) + state.ssm_state[:, 3:6].fill_(47 + i) + slots = manager.bounded_att_big_page_allocator.alloc(2) + manager.save_bounded_att_runtime_to_pages( + req_indices=[0, 1], + state_token_lens=[4, 8], + page_kind="big", + page_slot_ids=slots, + ) + if manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.is_cuda: + torch.cuda.synchronize() + return slots + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") +def test_shared_slots_save_restore_and_propagate_failure(monkeypatch: pytest.MonkeyPatch): + manager = composition(device="cuda") + slots = init_source(manager) + manager.restore_bounded_att_runtime_from_pages( + req_indices=[2], + state_token_lens=[4], + page_kind="big", + page_slot_ids=[slots[0]], + ) + for i, state in enumerate((manager.main_kda_conv_fp32_ssm_fp32_state, manager.draft0_gdn_conv_fp32_ssm_fp32_state)): + assert (state.conv_state[:, 2, :, :2] == 11 + i).all() + assert (state.ssm_state[:, 6] == 31 + i).all() + + def fail(*args: object, **kwargs: object): + raise RuntimeError("simulated copy failure") + + monkeypatch.setattr(manager.draft0_gdn_conv_fp32_ssm_fp32_state, "save_bounded_att_runtime_to_pages", fail) + with pytest.raises(RuntimeError): + manager.save_bounded_att_runtime_to_pages( + req_indices=[0, 1], + state_token_lens=[4, 8], + page_kind="big", + page_slot_ids=slots, + ) + assert manager.bounded_att_big_page_allocator.free_count == 1 + manager.bounded_att_big_page_allocator.free(slots) + assert manager.bounded_att_big_page_allocator.free_count == 3 + + +@pytest.mark.parametrize("operation", ["save_bounded_att_runtime_to_pages", "restore_bounded_att_runtime_from_pages"]) +def test_page_operations_forward_each_request_length(operation: str, monkeypatch: pytest.MonkeyPatch): + manager = composition() + calls = [] + + def record_lengths( + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + calls.append((req_indices, state_token_lens, page_kind, page_slot_ids)) + + for state in manager.get_bounded_attention_states(): + monkeypatch.setattr(state, operation, record_lengths) + getattr(manager, operation)( + req_indices=[0, 2], + state_token_lens=[17, 29], + page_kind="big", + page_slot_ids=[0, 1], + ) + assert calls == [([0, 2], [17, 29], "big", [0, 1])] * len(manager.get_bounded_attention_states()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for token indices") +def test_cpu_restore_forwards_absolute_length_for_partial_prefix(monkeypatch: pytest.MonkeyPatch): + manager = composition() + layouts = layouts_for(manager) + manager.bounded_att_cpu_load_big_page_slot = 2 + for state in manager.get_full_attention_states(): + monkeypatch.setattr(state, "read_from_cpu", lambda **kwargs: None) + for state in manager.get_bounded_attention_states(): + monkeypatch.setattr(state, "read_state_pages_from_cpu", lambda **kwargs: None) + restored = [] + + def record_restore( + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + restored.append((req_indices, state_token_lens, page_kind, page_slot_ids)) + + monkeypatch.setattr(manager, "restore_bounded_att_runtime_from_pages", record_restore) + # 只加载绝对位置 [64, 70)，本批长度 6；恢复位置必须是 70，且末尾碎页使用 LOAD。 + manager.get_cpu_cache_operator().read_req( + cpu_page_buffers=torch.empty((2, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8), + cpu_page_indices=[0, 1], + token_indices=torch.arange(6, device="cuda"), + layouts=layouts, + big_page_slot_ids=[0], + req_idx=1, + state_token_len=70, + ) + assert restored == [([1], [70], "big", [2])] + + +@pytest.mark.parametrize("operation", ["write_to_pd_page", "read_from_pd_page"]) +def test_pd_bounded_only_forwards_prefix_length(operation: str, monkeypatch: pytest.MonkeyPatch): + manager = composition() + layouts = layouts_for(manager, kind="att_state") + lengths = [] + + def record_length( + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + lengths.append((req_idx, state_token_len)) + + for state in manager.get_bounded_attention_states(): + monkeypatch.setattr(state, operation, record_length) + getattr(manager.get_pd_operator(), "write_page" if operation == "write_to_pd_page" else "read_page")( + pd_page_buffers=torch.empty((1, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8), + pd_page_index=0, + kind="att_state", + layouts=layouts, + req_idx=2, + state_token_len=29, + ) + assert lengths == [(2, 29)] * len(manager.get_bounded_attention_states()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("token_count", [2, 4, 6, 8]) +def test_cpu_restore_owns_partial_slot_and_restores_final_runtime(token_count: int): + source, target = composition("cuda"), composition("cuda") + source_slots = init_source(source) + layouts = layouts_for(source) + page_count = (token_count + 3) // 4 + page_ids = [2, 0][:page_count] + pages = torch.zeros((3, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) + small_slots = [] + if token_count % 4: + small_slots = source.bounded_att_small_page_allocator.alloc(1) + source.save_bounded_att_runtime_to_pages( + req_indices=[page_count - 1], + state_token_lens=[token_count], + page_kind="small", + page_slot_ids=small_slots, + ) + torch.cuda.synchronize() + source.get_cpu_cache_operator().write_req( + token_indices=torch.arange(token_count, device="cuda"), + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=layouts, + big_page_slot_ids=source_slots[: token_count // 4], + small_page_slot_ids=small_slots, + ) + torch.cuda.synchronize() + + # 保留末槽给 LOAD；调用方只分配完整页，碎页不占普通页额度。 + target.bounded_att_big_page_allocator = StatePageAllocator(3, reserved_slots=1) + target.bounded_att_cpu_load_big_page_slot = 2 + allocator = target.bounded_att_big_page_allocator + slots = allocator.alloc(token_count // 4) + original_slots = slots.copy() + free_count = allocator.free_count + target.get_cpu_cache_operator().read_req( + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + token_indices=torch.arange(4, 4 + token_count, device="cuda"), + state_token_len=token_count, + layouts=layouts, + big_page_slot_ids=slots, + req_idx=2, + ) + torch.cuda.synchronize() + assert slots == original_slots + assert allocator.free_count == free_count + for src, dst in zip(source.get_bounded_attention_states(), target.get_bounded_attention_states()): + pool = dst.big_pages_pool + for src_slot, dst_slot in zip(source_slots, slots): + torch.testing.assert_close(pool.conv_state[dst_slot], src.big_pages_pool.conv_state[src_slot]) + torch.testing.assert_close(pool.ssm_state[dst_slot], src.big_pages_pool.ssm_state[src_slot]) + final_slot = source_slots[page_count - 1] + torch.testing.assert_close(dst.conv_state[:, 2, :, :2].cpu(), src.big_pages_pool.conv_state[final_slot]) + torch.testing.assert_close(dst.ssm_state[:, 6].cpu(), src.big_pages_pool.ssm_state[final_slot]) + assert not dst.conv_state[:, :2].any() + if token_count % 4: + torch.testing.assert_close(dst.big_pages_pool.conv_state[2], src.big_pages_pool.conv_state[final_slot]) + torch.testing.assert_close(dst.big_pages_pool.ssm_state[2], src.big_pages_pool.ssm_state[final_slot]) + for src, dst in zip(source.get_full_attention_states(), target.get_full_attention_states()): + torch.testing.assert_close(read_kv_storage(dst)[:, 4 : 4 + token_count], read_kv_storage(src)[:, :token_count]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "readies", [[True, False, False], [False, True, False], [False, False, True], [True, True, True]] +) +def test_cpu_ready_joint_pages_preserve_full_and_bounded_fields(readies: list[bool]): + source = composition("cuda") + big_slots = init_source(source) + small_slots = source.bounded_att_small_page_allocator.alloc(1) + source.save_bounded_att_runtime_to_pages( + req_indices=[1], + state_token_lens=[6], + page_kind="small", + page_slot_ids=small_slots, + ) + torch.cuda.synchronize() + layouts = layouts_for(source) + page_ids = [2, 0, 3] + token_ids = torch.arange(10, device="cuda") + baseline = torch.full( + size=(4, max(layout.byte_end for layout in layouts.values())), + fill_value=139, + dtype=torch.uint8, + pin_memory=True, + ) + pages = torch.full_like(baseline, 139, pin_memory=True) + source.get_cpu_cache_operator().write_req( + token_indices=token_ids, + cpu_page_buffers=baseline, + cpu_page_indices=page_ids, + layouts=layouts, + big_page_slot_ids=big_slots, + small_page_slot_ids=small_slots, + ) + torch.cuda.synchronize() + expected = pages.clone() + for page, ready in zip(page_ids, readies): + if not ready: + expected[page].copy_(baseline[page]) + # ready 页不能访问其页池快照；负的越界槽位可检测是否在构造视图前跳过。 + masked_big_slots = [-999 if readies[i] else slot for i, slot in enumerate(big_slots)] + masked_small_slots = [-999] if readies[2] else small_slots + source.get_cpu_cache_operator().write_req( + token_indices=token_ids, + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=layouts, + cpu_page_ready_flags=readies, + big_page_slot_ids=masked_big_slots, + small_page_slot_ids=masked_small_slots, + ) + torch.cuda.synchronize() + assert torch.equal(pages, expected) diff --git a/unit_tests/common/att_resources/test_model_calibration.py b/unit_tests/common/att_resources/test_model_calibration.py new file mode 100644 index 0000000000..9544084bb6 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -0,0 +1,261 @@ +"""统一标定阶段不依赖 State 注册顺序，不在布局定义阶段加载参数。""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from dataclasses import replace +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.server.core.objs.start_args_type import StartArgs + +if TYPE_CHECKING: + from pathlib import Path + from easydict import EasyDict + + +@pytest.mark.parametrize("encoding", ["fp8kv_sph", "fp8kv_spt"]) +def test_calibration_binds_all_model_states_after_creation(tmp_path: Path, encoding: str): + config = dict( + architectures=["Qwen3NextForCausalLM"], + num_hidden_layers=8, + num_attention_heads=2, + num_key_value_heads=1, + hidden_size=32, + head_dim=16, + full_attention_interval=4, + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=8, + linear_value_head_dim=8, + linear_conv_kernel_dim=3, + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + scales = torch.arange(1, 9, dtype=torch.float32) * 0.125 + calibration = dict( + architectures="Qwen3NextForCausalLM", + num_head=1, + num_layers=8, + scales_shape=[8, 2], + scales=scales[:, None].expand(-1, 2).tolist(), + q_calibration=dict(num_head=1, scales_shape=[8, 1], scales=(scales[:, None] * 2).tolist()), + quant_type="per_head" if encoding == "fp8kv_sph" else "per_tensor", + ) + calibration_path = tmp_path / "calibration.json" + calibration_path.write_text(json.dumps(calibration)) + args = StartArgs( + model_dir=str(tmp_path), + data_type="bfloat16", + llm_kv_type=encoding, + kv_quant_calibration_config_path=str(calibration_path), + page_size=4, + mtp_mode="vanilla_with_att", + ) + definitions = ( + (Qwen3NextTpPartModel, config, 8, args, (3, 7), (0, 2)), + ( + Qwen3NextTpPartModel, + dict(config, num_hidden_layers=6, full_attention_interval=2), + 6, + args, + (9, 11, 13), + (2, 5), + ), + (Qwen3NextTpPartModel, dict(config, num_hidden_layers=3), 3, args, (), (5, 5)), + ( + LlamaTpPartModel, + dict(config, num_hidden_layers=3), + 3, + replace(args, mtp_mode="vanilla_no_att"), + (), + (5, 5), + ), + # Native MTP keeps the main config but its resource hook describes one full-attention layer. + (Qwen3_5MTPModel, config, 1, args, (20,), (5, 6)), + (LlamaTpPartModel, dict(config, num_hidden_layers=2), 2, args, (21, 22), (6, 8)), + ) + models, layer_start = [], 0 + for cls, model_config, layer_count, model_args, _, _ in definitions: + model = cls.__new__(cls) + model.is_mtp_draft_model = bool(models) + model.config, model.args = model_config.copy(), model_args + model.tp_rank_, model.tp_world_size_ = 0, 1 + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(layer_start, layer_start + layer_count)] + if models: + model.main_model, model.mtp_previous_draft_models = models[0], models[1:].copy() + models.append(model) + layer_start += layer_count + + # 反向创建、注册，标定仍按 main 和 draft 的数值顺序绑定。 + manager = AttentionStateManager(request_capacity=1) + for index in reversed(range(len(models))): + model = models[index] + model.att_state_manager = manager + states = [] + if not (model.is_mtp_draft_model and model.args.mtp_mode in ("vanilla_no_att", "eagle_no_att")): + states = model.create_full_att_states( + model_config=model.config, + args=model.args, + model_logical_layer_ids=tuple(layer.layer_num_ for layer in model.layers_infer), + tp_rank=model.tp_rank_, + tp_world_size=model.tp_world_size_, + ) + for state in states: + assert state.state_name is None + assert state.scales is None and state.q_scales is None + manager.register_state(model.att_state_role, state) + manager.finish_state_registration() + manager.load_full_att_calibration(args=args) + for index, model in enumerate(models): + states = manager.get_full_attention_states(model.att_state_role) + state = states[0] if states else None + expected_layers, (start, end) = definitions[index][-2:] + if not expected_layers: + assert state is None + continue + assert state.config.logical_layer_groups == tuple((layer,) for layer in expected_layers) + assert read_kv_storage(state) is None + assert torch.equal(state.scales[:, 0], scales[start:end]) + if encoding == "fp8kv_sph": + assert torch.equal(state.q_scales[:, 0], scales[start:end] * 2) + + +@pytest.mark.parametrize("tp_world_size", [1, 4]) +def test_calibration_reads_once_and_handles_shared_layers_numeric_roles_and_tp( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + tp_world_size: int, +): + from builtins import open as original_open + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + + # 第 0 行属于普通精度的共享层；其余行分别属于 main、draft2、draft10 的 FP8 State。 + scales = torch.arange(1, 17, dtype=torch.float32).reshape(4, 4) + q_scales = scales[:, :2] + 20 + calibration_path = tmp_path / "calibration.json" + calibration_path.write_text( + json.dumps( + dict( + num_layers=4, + num_head=2, + quant_type="per_head", + scales_shape=[4, 4], + scales=scales.tolist(), + q_calibration=dict(num_head=2, scales_shape=[4, 2], scales=q_scales.tolist()), + ) + ) + ) + args = StartArgs(kv_quant_calibration_config_path=str(calibration_path)) + reads = [] + + def count_open(file: str, *args: object, **kwargs: object): + if file == str(calibration_path): + reads.append(file) + return original_open(file, *args, **kwargs) + + monkeypatch.setattr("builtins.open", count_open) + for tp_rank in range(tp_world_size): + manager = AttentionStateManager(request_capacity=1) + expected = [] + # 特意打乱角色及同角色 State 的注册顺序。 + for role, groups, row, encoding in ( + ("draft10", ((5,),), 3, "fp8kv_sph"), + ("main", ((2, 3),), 1, "fp8kv_sph"), + ("draft2", ((4,),), 2, "fp8kv_sph"), + ("main", ((0, 1),), 0, "None"), + ): + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=groups, + global_head_num=4, + head_dim=16, + dtype=torch.bfloat16, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + kv_type=encoding, + device="meta", + ) + manager.register_state(role, state) + if encoding != "None": + expected.append((state, row)) + with pytest.raises(AssertionError, match="State registration"): + manager.load_full_att_calibration(args=args) + manager.finish_state_registration() + manager.load_full_att_calibration(args=args) + heads = 4 // tp_world_size + local_heads = slice(tp_rank * heads, (tp_rank + 1) * heads) + for state, row in expected: + expected_k = scales[row : row + 1, :2].repeat_interleave(2, dim=-1)[:, local_heads] + expected_v = scales[row : row + 1, 2:].repeat_interleave(2, dim=-1)[:, local_heads] + expected_q = q_scales[row : row + 1].repeat_interleave(2, dim=-1)[:, local_heads] + assert torch.equal(state.scales, torch.cat((expected_k, expected_v), dim=-1)) + assert torch.equal(state.q_scales, expected_q) + assert state.k_storage is None + assert len(reads) == tp_world_size # 每个管理器只读取一次，不随 State 数量增加。 + + +@pytest.mark.parametrize("encoding", ["fp8kv_spt", "fp8kv_sph"]) +def test_cpu_layout_does_not_require_calibration_file( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + cpu_args: EasyDict, + encoding: str, +): + from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts + + config = dict( + model_type="llama", + num_attention_heads=2, + hidden_size=32, + num_hidden_layers=2, + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + cpu_args.model_dir = str(tmp_path) + cpu_args.llm_kv_type = encoding + cpu_args.data_type = "bfloat16" + cpu_args.kv_quant_calibration_config_path = str(tmp_path / "does-not-exist.json") + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: LlamaTpPartModel) + layouts = build_cpu_cache_layouts() + assert len(layouts) == 1 + (layout,) = layouts.values() + assert layout.byte_end > layout.byte_start + assert {region.field for region in layout.regions} == {"k", "v"} + + +@pytest.mark.parametrize("encoding", ["fp8kv_spt", "fp8kv_sph"]) +def test_deferred_calibration_keeps_unit_scale_policy(encoding: str): + from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state + + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=((0, 1), (2,)), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + ), + kv_type=encoding, + device="cpu", + ) + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + with pytest.raises(AssertionError, match="calibration must be bound"): + state.allocate_storage(4) + if encoding == "fp8kv_sph": + with pytest.raises(AssertionError, match="requires a KV calibration file"): + manager.load_full_att_calibration(args=StartArgs()) + else: + manager.load_full_att_calibration(args=StartArgs()) + assert torch.equal(state.scales, torch.ones((2, 2))) + state.allocate_storage(4) + assert state.k_storage.shape[0] == 2 diff --git a/unit_tests/common/att_resources/test_model_config_normalization.py b/unit_tests/common/att_resources/test_model_config_normalization.py new file mode 100644 index 0000000000..e3658f3cbe --- /dev/null +++ b/unit_tests/common/att_resources/test_model_config_normalization.py @@ -0,0 +1,166 @@ +"""模型启动和 CPU cache 必须使用同一份有效配置。""" + +from copy import deepcopy +import importlib +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel import TpPartBaseModel +from lightllm.models.glm4_moe_lite_mtp.model import Glm4MoeLiteMTPModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.qwen3_5_dflash.model import Qwen3_5DFlashModel +from lightllm.models.qwen3_5_dspark.model import Qwen3_5DSparkModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.models.qwen3_dflash2.model import Qwen3DFlash2Model +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("model_class", [Qwen3DFlash2Model, Qwen3_5DFlashModel, Qwen3_5DSparkModel]) +@pytest.mark.parametrize("nested", [False, True]) +def test_dflash_runtime_and_cpu_state_use_overridden_config( + model_class: type[TpPartBaseModel], + tmp_path: Path, + nested: bool, +) -> None: + text_config = { + "num_hidden_layers": 32, + "n_layer": 32, + "num_attention_heads": 8, + "n_head": 8, + "hidden_size": 1024, + "num_key_value_heads": 4, + "rope_parameters": {"rope_theta": 10000, "partial_rotary_factor": 0.5}, + "dflash_config": {"num_hidden_layers": 4, "num_attention_heads": 4, "hidden_size": 512}, + } + raw_config = {"text_config": text_config} if nested else text_config + original = deepcopy(raw_config) + (tmp_path / "config.json").write_text(json.dumps(raw_config)) + + # CPU 服务只调用类接口；模型启动从同一文件走真实 _init_config 调用链。 + cpu_config = model_class.normalize_model_config(raw_config) + model = model_class.__new__(model_class) + model.weight_dir_ = str(tmp_path) + model.finetune_config = None + model._init_config() + assert model.config == cpu_config + assert cpu_config["num_hidden_layers"] == cpu_config["n_layer"] == 4 + assert cpu_config["num_attention_heads"] == cpu_config["n_head"] == 4 + assert cpu_config["hidden_size"] == cpu_config["n_embd"] == cpu_config["n_embed"] == 512 + assert model_class.normalize_model_config(cpu_config) == cpu_config + assert raw_config == original + + args = StartArgs(data_type="bfloat16", llm_kv_type="None", page_size=4) + (state,) = model_class.create_full_att_states( + model_config=cpu_config, + args=args, + model_logical_layer_ids=tuple(range(cpu_config["num_hidden_layers"])), + ) + assert state.config.logical_layer_groups == ((0,), (1,), (2,), (3,)) + assert state.config.head_dim == 128 + + +@pytest.mark.parametrize("model_class", [Glm4MoeLiteMTPModel, Glm5NextMTPModel, Qwen3_5MTPModel]) +def test_native_mtp_runtime_matches_normalized_config( + model_class: type[TpPartBaseModel], + tmp_path: Path, +) -> None: + text_config = { + "num_hidden_layers": 32, + "num_nextn_predict_layers": 1, + "hidden_size": 1024, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "full_attention_interval": 4, + "mhc": True, + "scoring_func": "sigmoid", + } + original = deepcopy(text_config) + raw_config = {"text_config": text_config, "vision_config": {}} + (tmp_path / "config.json").write_text(json.dumps(raw_config)) + config = model_class.normalize_model_config(raw_config) + model = model_class.__new__(model_class) + model.weight_dir_ = str(tmp_path) + model.finetune_config = None + model.tp_world_size_ = 1 + model.main_model = SimpleNamespace(config=text_config) + model._init_config() + + assert model.config == config + assert config["num_hidden_layers"] == config["n_layer"] == 1 + assert model_class.normalize_model_config(config) == config + assert text_config == original + if model_class is Qwen3_5MTPModel: + assert config["full_attention_interval"] == 1 + else: + assert config["mtp_checkpoint_layer_start"] == 32 + + +@pytest.mark.parametrize("model_class", [Glm4MoeLiteMTPModel, Glm5NextMTPModel]) +def test_native_mtp_weights_keep_checkpoint_layer_start( + model_class: type[TpPartBaseModel], + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = model_class.normalize_model_config( + { + "num_hidden_layers": 32, + "num_nextn_predict_layers": 1, + "hidden_size": 1024, + } + ) + module = importlib.import_module(model_class.pre_and_post_weight_class.__module__) + names = [] + + def init_base(self: object, data_type: torch.dtype, network_config: dict[str, object]) -> None: + self.data_type_ = data_type + + def record_weight(**kwargs: object) -> SimpleNamespace: + names.append(kwargs.get("weight_names", kwargs.get("weight_name"))) + return SimpleNamespace() + + monkeypatch.setattr(module.PreAndPostLayerWeight, "__init__", init_base) + monkeypatch.setattr(module, "ROWMMWeight", record_weight) + monkeypatch.setattr(module, "RMSNormWeight", record_weight) + quant_cfg = SimpleNamespace(get_quant_method=lambda layer, name: None) + pre_post_weight = model_class.pre_and_post_weight_class( + data_type=torch.bfloat16, + network_config=config, + quant_cfg=quant_cfg, + ) + assert len(names) == 4 + assert all(name.startswith("model.layers.32.") for name in names) + + model = model_class.__new__(model_class) + model.config = config + model.data_type = torch.bfloat16 + model.quant_cfg = quant_cfg + model.main_model = SimpleNamespace( + pre_post_weight=SimpleNamespace( + wte_weight_=object(), + lm_head_weight_=object(), + final_norm_weight_=object(), + ) + ) + model.pre_and_post_weight_class = lambda *args, **kwargs: pre_post_weight + layers = [] + + def record_layer( + layer_num: int, + data_type: torch.dtype, + network_config: dict[str, object], + quant_cfg: object, + ) -> object: + layers.append(layer_num) + return SimpleNamespace(verify_load=lambda: None) + + model.transformer_weight_class = record_layer + if model_class is Glm4MoeLiteMTPModel: + model.weight_dir_ = "unused" + model.weight_dict = None + pre_post_weight.verify_load = lambda: None + monkeypatch.setattr(importlib.import_module(model_class.__module__), "load_hf_weights", lambda *a, **kw: None) + model._init_weights() + assert layers == [32] diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py new file mode 100644 index 0000000000..af8dbfcde0 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states +from lightllm.common.attention_state.full_att import FullAttentionState +from unit_tests.common.att_resources.model_helpers import owned_states +import dataclasses + +import pytest +import torch + +from lightllm.common.req_manager import ReqManager +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import ( + get_added_mtp_kv_layer_num, + get_env_start_args, + get_unique_server_name, + set_env_start_args, + set_unique_server_name, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def clear_mtp_capacity_cache(): + get_added_mtp_kv_layer_num.cache_clear() + yield + get_added_mtp_kv_layer_num.cache_clear() + + +@pytest.mark.parametrize("glm", [False, True]) +@pytest.mark.parametrize("mtp", [0, 2]) +@pytest.mark.parametrize("mode", ["normal", "prefill"]) +def test_real_model_resource_hooks_and_shared_draft( + monkeypatch: pytest.MonkeyPatch, + glm: bool, + mtp: int, + mode: str, +): + args = StartArgs( + run_mode=mode, + data_type="bfloat16", + page_size=4, + pd_kv_page_size=128, + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + linear_att_cache_size=2, + cpu_cache_token_page_size=8, + mtp_step=mtp, + mtp_mode="eagle_with_att" if mtp else None, + ) + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) + set_env_start_args(dataclasses.asdict(args)) + get_env_start_args.cache_clear() + set_unique_server_name(args) + get_unique_server_name.cache_clear() + model_type = Glm5NextTpPartModel if glm else Qwen3NextTpPartModel + model = model_type.__new__(model_type) + model.args, model.data_type = args, torch.bfloat16 + model.tp_world_size_, model.tp_rank_ = 1, 0 + model.max_req_num, model.max_total_token_num = 3, None + model.max_seq_length, model.batch_max_tokens = 32, 16 + model.mem_fraction, model.num_kv_heads = 0.9, 2 + model.config = { + "num_attention_heads": 4, + "hidden_size": 64, + "num_key_value_heads": 2, + "head_dim": 16, + "n_layer": 4, + "num_hidden_layers": 4, + "full_attention_interval": 2, + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + "linear_conv_kernel_dim": 3, + } + if glm: + model.config.update( + num_key_value_heads=1, + kv_lora_rank=512, + index_kpool=4, + index_head_dim=128, + layer_types=["linear_attention"] * 3 + ["deepseek_sparse_attention"], + linear_attn_config={"num_heads": 4, "head_dim": 8, "short_conv_kernel_size": 3, "kda_layers": [0, 1, 2]}, + ) + observed = [] + + def check_runtime_before_profile(device: str | torch.device | int): + resources = model.att_state_manager + states = resources.get_full_attention_states() + assert main_linear_state(resources).conv_state.is_cuda + assert main_linear_state(resources).big_pages_pool is None + assert resources.full_att_allocator is None + assert not hasattr(resources, "pd_page_bytes") + assert not hasattr(resources, "pd_transfer_layouts") + assert read_kv_storage(states[0]) is None + observed.append(main_linear_state(resources).conv_state.data_ptr()) + reserve = 0 + if mode == "prefill": + reserve = resources.get_pd_operator().get_page_bytes() * args.pd_kv_page_num + return resources.get_full_att_token_bytes() * (32 + args.page_size) + reserve, 0 + + monkeypatch.setattr(torch.cuda, "mem_get_info", check_runtime_before_profile) + prepare_model_states(model) + drafts = [] + if mtp: + draft_type = Glm5NextMTPModel if glm else Qwen3_5MTPModel + draft = draft_type.__new__(draft_type) + draft.main_model, draft.mtp_previous_draft_models = model, [] + draft.args, draft.data_type = args, torch.bfloat16 + draft.tp_rank_, draft.tp_world_size_ = 0, 1 + draft.max_req_num, draft.max_total_token_num = model.max_req_num, None + draft.config = dict( + model.config, + num_hidden_layers=1, + n_layer=1, + full_attention_interval=1, + ) + draft.max_seq_length, draft.batch_max_tokens = model.max_seq_length, model.batch_max_tokens + prepare_model_states(draft, logical_layer_ids=(4,)) + assert draft.att_state_manager is model.att_state_manager + assert read_kv_storage(owned_states(draft.att_state_manager.__dict__, FullAttentionState)[0]) is None + drafts.append(draft) + models = (model, *drafts) + req_manager = ReqManager( + max_request_num=model.max_req_num, + max_sequence_length=max(item.get_req_manager_max_seq_len() for item in models), + ) + for item in models: + item.bind_req_manager(req_manager=req_manager) + assert all(item.req_manager is req_manager for item in models) + model.att_state_manager.finish_state_registration() + model.att_state_manager.load_full_att_calibration(args=args) + model.att_state_manager.initialize_storage( + args, token_capacity=model.max_total_token_num, mem_fraction=model.mem_fraction + ) + model.req_manager.bind_hold_page(model.att_state_manager.full_att_allocator.hold_token_indices) + resources = model.att_state_manager + (full_state,) = resources.get_full_attention_states("main") + assert type(model.req_manager) is ReqManager + assert not hasattr(model.req_manager, "req_to_ssm_state") + assert resources.full_att_allocator.size == full_state.capacity == 32 + assert not hasattr(model, "mem_manager") + assert not hasattr(model, "full_att_state") + full_interval = 4 if glm else model.config["full_attention_interval"] + assert full_state.config.logical_layer_groups == tuple( + (layer,) for layer in range(full_interval - 1, model.config["num_hidden_layers"], full_interval) + ) + assert observed == [main_linear_state(resources).conv_state.data_ptr()] + assert resources.bounded_att_big_page_allocator.size == 6 + assert resources.bounded_att_small_page_allocator.size == 2 + assert not hasattr(model, "mtp_state_cursors") + assert (main_linear_state(resources).mtp_state_index is not None) == bool(mtp) + if mode == "prefill": + assert resources.get_pd_operator().get_pd_full_att_page_layouts() + assert resources.get_pd_operator().get_pd_bounded_att_page_layouts() + if mtp: + draft = drafts[0] + assert draft.req_manager is model.req_manager + assert draft.att_state_manager is resources + assert draft.att_state_manager.full_att_allocator is resources.full_att_allocator + assert not hasattr(draft, "full_att_state") + (draft_state,) = resources.get_full_attention_states("draft0") + assert draft_state is not full_state + assert draft_state.capacity == full_state.capacity + assert {layer for group in draft_state.config.logical_layer_groups for layer in group}.isdisjoint( + layer for group in full_state.config.logical_layer_groups for layer in group + ) + assert not hasattr(draft, "mem_manager") + assert not hasattr(draft, "mtp_state_cursors") + assert observed == [main_linear_state(resources).conv_state.data_ptr()] diff --git a/unit_tests/common/att_resources/test_model_resources.py b/unit_tests/common/att_resources/test_model_resources.py new file mode 100644 index 0000000000..a571817fae --- /dev/null +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from typing import TYPE_CHECKING +from unit_tests.common.att_resources.model_helpers import main_linear_state +import dataclasses +import json + +import pytest +import torch + +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.model_helpers import create_full_attention_roles +from lightllm.common.attention_state.full_att import ( + GQAFullAttentionState, + GQAAttentionConfig, +) +from lightllm.common.attention_state.full_att import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, +) +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from lightllm.models.glm5_next.config import Glm5NextResourceConfig +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args + +if TYPE_CHECKING: + from easydict import EasyDict + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def make_resources( + glm: bool, + tp: int, + rank: int, + mtp: int, + draft_layers: int = 0, +): + gk, gv = (4, 4) if glm else (2, 4) + cls = Glm5NextResourceConfig if glm else Qwen3NextResourceConfig + config = cls( + tp_world_size=tp, + full_att_all_num_kv_heads=1 if glm else 2, + full_att_head_dim=512 if glm else 16, + global_linear_k_heads=gk, + global_linear_v_heads=gv, + head_linear_k_dim=8, + head_linear_v_dim=8, + conv_kernel_size=3, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4 if glm else 2, + all_layer_num=4, + ) + if glm: + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + logical_config=PackedMLALogicalConfig( + logical_layer_groups=tuple((layer,) for layer in (3, *range(4, 4 + draft_layers))), + format="packed_bf16", + tp_rank=rank, + tp_world_size=tp, + latent_dim=512, + ) + ), + page_size=4, + ) + full.allocate_storage(16) + else: + full = GQAFullAttentionState( + GQAAttentionConfig( + logical_layer_groups=tuple((layer,) for layer in (1, 3, *range(4, 4 + draft_layers))), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + tp_rank=rank, + tp_world_size=tp, + ), + page_size=4, + ) + full.allocate_storage(16) + main, draft = create_full_attention_roles(full, 4) + resources = build_hybrid_resources( + config=config, + request_capacity=3, + mtp_step=mtp, + tp_rank=rank, + full_state=main, + full_allocator=FullAttKVAllocator(16, 4), + draft_full_state=draft, + big_page_capacity=4, + small_page_capacity=2, + ) + return resources, resources.get_full_attention_states("main")[0] + + +@pytest.mark.parametrize("glm", [False, True]) +@pytest.mark.parametrize("tp", [1, 2, 4]) +@pytest.mark.parametrize("mtp", [0, 2]) +def test_model_cpu_layout_roundtrip_tp_shards_replicas_and_tail( + glm: bool, + tp: int, + mtp: int, + cpu_args: EasyDict, +): + sources = [ + make_resources( + glm=glm, + tp=tp, + rank=rank, + mtp=mtp, + ) + for rank in range(tp) + ] + targets = [ + make_resources( + glm=glm, + tp=tp, + rank=rank, + mtp=mtp, + ) + for rank in range(tp) + ] + expected = [] + for rank, (resources, full) in enumerate(sources): + cfg = main_linear_state(resources).config + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.fill_(1 if glm else rank // max(1, tp // 2) + 1) + state = main_linear_state(resources) + local_start = 0 + for local_heads, global_heads in ( + (cfg.num_k_heads, cfg.global_k_heads), + (cfg.num_k_heads, cfg.global_k_heads), + (cfg.num_v_heads, cfg.global_v_heads), + ): + start = rank * local_heads // max(1, tp * local_heads // global_heads) + for h in range(local_heads): + state.conv_state[:, 0, local_start + h * 8 : local_start + (h + 1) * 8].fill_(start + h + 3) + local_start += local_heads * 8 + for h in range(cfg.num_v_heads): + state.ssm_state[:, : mtp + 1, h].fill_(rank * cfg.num_v_heads + h + 7) + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + resources.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[slot], + ) + torch.cuda.synchronize() + expected.append((state.big_pages_pool.conv_state[slot].clone(), state.big_pages_pool.ssm_state[slot].clone())) + size = sources[0][0].get_cpu_cache_operator().get_page_bytes() + buffer = torch.zeros((2, size), dtype=torch.uint8, pin_memory=True) + for resources, full in sources: + resources.get_cpu_cache_operator().write_req( + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[1], + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + big_page_slot_ids=[0], + ) + torch.cuda.synchronize() + for rank, (resources, full) in enumerate(targets): + if glm: + resources.main_compressor_indexer_w4_s4_bf16_state.buffer.fill_(13) + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + resources.get_cpu_cache_operator().read_req( + cpu_page_buffers=buffer, + cpu_page_indices=[1], + token_indices=torch.arange( + start=8, + end=16, + dtype=torch.long, + device="cuda", + ), + state_token_len=8, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + big_page_slot_ids=[slot], + req_idx=1, + ) + torch.testing.assert_close(main_linear_state(resources).conv_state[:, 1, :, :2].cpu(), expected[rank][0]) + torch.testing.assert_close(main_linear_state(resources).ssm_state[:, mtp + 1].cpu(), expected[rank][1]) + src_fields = full_state_test_fields(sources[rank][1], full_state_test_buffers(sources[rank][1])) + dst_fields = full_state_test_fields(full, full_state_test_buffers(full)) + for name in src_fields: + assert torch.equal( + src_fields[name].tensor[:, :8].contiguous().view(torch.uint8), + dst_fields[name].tensor[:, 8:16].contiguous().view(torch.uint8), + ) + if glm: + # 无重叠窗口在组边界不依赖旧历史，CPU 恢复无需修改其数据。 + assert (resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 1] == 13).all() + assert (resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, [0, 2]] == 13).all() + assert not buffer[0].any() + + +def test_cpu_layout_query_does_not_require_cpu_cache_enabled(monkeypatch: pytest.MonkeyPatch): + args = StartArgs(enable_cpu_cache=False, cpu_cache_token_page_size=8) # 未启用缓存时也可查询全局页配置。 + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + get_env_start_args.cache_clear() + resources, _ = make_resources( + glm=True, + tp=1, + rank=0, + mtp=0, + ) + assert all(layout.tokens_per_page == 8 for layout in resources.get_cpu_cache_operator().get_page_layouts().values()) + get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_model_storage_definition.py b/unit_tests/common/att_resources/test_model_storage_definition.py new file mode 100644 index 0000000000..ed6b891603 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -0,0 +1,276 @@ +"""Model-specific storage geometry must also describe external cache pages.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state + +from typing import TYPE_CHECKING +import dataclasses +import json +from functools import partial + +from unit_tests.common.att_resources.model_helpers import prepare_model_states + +import pytest +import torch + +from lightllm.models.bloom.model import BloomTpPartModel +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.qwen2.model import Qwen2TpPartModel +from lightllm.models.starcoder.model import StarcoderTpPartModel +from lightllm.models.starcoder2.model import Starcoder2TpPartModel +from lightllm.models.mistral.model import MistralTpPartModel +from lightllm.models.mixtral.model import MixtralTpPartModel +from lightllm.models.gemma_2b.model import Gemma_2bTpPartModel +from lightllm.models.gemma3.model import Gemma3TpPartModel +from lightllm.models.gemma4.model import Gemma4TpPartModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.test_quantized_full_state import KINDS, make_quant_state + +if TYPE_CHECKING: + from easydict import EasyDict + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from pathlib import Path + + +@pytest.fixture +def args(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + value = StartArgs( + model_dir=str(tmp_path), + data_type="float16", + tp=4, + dp=1, + page_size=4, + enable_cpu_cache=True, + cpu_cache_token_page_size=8, + run_mode="prefill", + pd_kv_page_size=8, + ) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(value))) + get_env_start_args.cache_clear() + yield value + get_env_start_args.cache_clear() + + +CASES = [ + (BloomTpPartModel, {}, 8, 16, 2), + (LlamaTpPartModel, {"num_key_value_heads": 4, "head_dim": 32}, 4, 32, 1), + (Qwen2TpPartModel, {"num_key_value_heads": 2}, 2, 16, 1), + (StarcoderTpPartModel, {}, 1, 16, 1), + (Starcoder2TpPartModel, {"num_key_value_heads": 4}, 4, 16, 1), + (MistralTpPartModel, {"num_key_value_heads": 4, "head_dim": 32}, 4, 32, 1), + (MixtralTpPartModel, {"num_key_value_heads": 4}, 4, 16, 1), + (Gemma_2bTpPartModel, {"num_key_value_heads": 2}, 2, 16, 2), + (Gemma3TpPartModel, {"num_key_value_heads": 4}, 4, 256, 1), + ( + Gemma4TpPartModel, + {"num_key_value_heads": 4, "head_dim": 256, "num_global_key_value_heads": None, "global_head_dim": 512}, + 8, + 256, + 2, + ), + ( + Gemma4TpPartModel, + {"num_key_value_heads": 16, "head_dim": 256, "num_global_key_value_heads": 4, "global_head_dim": 512}, + 16, + 256, + 4, + ), + ( + Gemma4TpPartModel, + { + "num_key_value_heads": 4, + "head_dim": 256, + "global_head_dim": 256, + "num_kv_shared_layers": 1, + "layer_types": ["full_attention", "full_attention"], + }, + 4, + 256, + 1, + ), +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("model_cls,overrides,global_heads,dim,local_heads", CASES) +def test_model_and_cpu_service_use_identical_storage_geometry( + args: StartArgs | EasyDict, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + model_cls: type[TpPartBaseModel], + overrides: dict[str, object], + global_heads: int, + dim: int, + local_heads: int, +): + # Exercise raw HF aliases as well as a wrapped multimodal config. + config = dict( + model_type="llama", + n_head=8, + n_embd=128, + n_layer=2, + **overrides, + ) + wrapped = {"model_type": "gemma4", "text_config": config} if model_cls is Gemma4TpPartModel else config + (tmp_path / "config.json").write_text(json.dumps(wrapped)) + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: model_cls) + service_layout = next(iter(build_cpu_cache_layouts().values())) + for rank in range(args.tp): + model = model_cls.__new__(model_cls) + model.config, model.args = wrapped, args + model.tp_rank_, model.tp_world_size_ = rank, args.tp + model.max_total_token_num, model.max_req_num = 19, 1 + prepare_model_states(model) + resources = model.att_state_manager + resources.finish_state_registration() + resources.initialize_storage(args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction) + (state,) = resources.get_full_attention_states("main") + assert state.config.head_num == local_heads and state.config.head_dim == dim + storage_layers = 2 - (config.get("num_kv_shared_layers") or 0) + assert state.k_storage.shape == state.v_storage.shape == (storage_layers, 20, local_heads, dim) + if storage_layers == 1: + assert state.config.logical_layer_groups == ((0, 1),) + assert resources.get_state("main", "gqa", 0) is resources.get_state("main", "gqa", 1) + assert ( + state.get_layer_view(logical_layer_id=0)[0].data_ptr() + == state.get_layer_view(logical_layer_id=1)[0].data_ptr() + ) + assert read_kv_storage(state).dtype == (torch.bfloat16 if model_cls is Gemma3TpPartModel else torch.float16) + assert resources.full_att_allocator.size == 16 + layout = resources.get_cpu_cache_operator().get_page_layouts()[state.state_name] + assert (layout.byte_start, layout.byte_end) == (service_layout.byte_start, service_layout.byte_end) + # 按物理层数计算传输页；共享消费者不重复增加 KV 或传输容量。 + expected_pd_bytes = ( + storage_layers * args.pd_kv_page_size * 2 * global_heads * dim * read_kv_storage(state).element_size() + ) + assert resources.get_pd_operator().get_page_bytes() == expected_pd_bytes + fields = layout.regions + assert all(region.shape[2] == global_heads for region in fields) + assert not hasattr(model, "mem_manager") + + +@pytest.mark.parametrize("kind", KINDS) +def test_gpu_profile_counts_token_fields_hold_and_transfer_reserve( + monkeypatch: pytest.MonkeyPatch, + kind: str, + pd_args: EasyDict, +): + state = make_quant_state(kind, "cpu") + reference = AttentionStateManager(request_capacity=1) + reference.register_state("main", state) + reference.full_att_allocator = FullAttKVAllocator(state.capacity, state.page_size, pin_memory=False) + reference.finish_state_registration() + args = StartArgs( + page_size=4, + run_mode="prefill", + pd_kv_page_size=8, + pd_kv_page_num=3, + ) + reserve = reference.get_pd_operator().get_page_bytes() * args.pd_kv_page_num + required = state.get_token_bytes() * (32 + args.page_size) + # One byte below the full 32-token budget must drop a complete page. + total = 1 << 30 + free = required - 1 + reserve + total // 4 + observed_devices = [] + + def memory_info(device: str | torch.device | int): + observed_devices.append(device) + return free, total + + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(torch.cuda, "mem_get_info", memory_info) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + monkeypatch.setattr("lightllm.common.attention_state.manager.get_current_device_id", lambda: 3) + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + template = rebuild_full_state( + template=state, + config=state.config, + token_capacity=0, + device="meta", + ) + resources = AttentionStateManager(request_capacity=1) + resources.register_state("main", rebuild_full_state(template, config=template.config)) + resources.finish_state_registration() + resources.initialize_storage( + args=args, + token_capacity=None, + mem_fraction=0.75, + device="cpu", + ) + assert resources.full_att_allocator.size == 28 + assert resources.get_pd_operator().get_page_bytes() == reference.get_pd_operator().get_page_bytes() + assert observed_devices == [torch.device("cuda", 3)] + + +def test_cpu_layout_uses_all_model_factory_states( + args: StartArgs | EasyDict, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + """同一模型返回多种方案和共享层时，CPU 页采用实际 State 列表。""" + from lightllm.common.attention_state.full_att import FullAttentionState, GQAFullAttentionState + + class MixedStorageModel(LlamaTpPartModel): + @classmethod + def create_full_att_states( + cls, + model_config: dict[str, object], + args: StartArgs | EasyDict, + model_logical_layer_ids: tuple[int, ...], + tp_rank: int = 0, + tp_world_size: int = 1, + ) -> list[FullAttentionState]: + (template,) = super().create_full_att_states( + model_config=model_config, + args=args, + model_logical_layer_ids=model_logical_layer_ids, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ) + return [ + GQAFullAttentionState( + layout_config=dataclasses.replace( + template.config, + logical_layer_groups=(model_logical_layer_ids[:2],), + dtype=torch.bfloat16, + ), + page_size=args.page_size, + device="meta", + ), + GQAFullAttentionState( + layout_config=dataclasses.replace( + template.config, + logical_layer_groups=(model_logical_layer_ids[2:],), + dtype=torch.float16, + ), + page_size=args.page_size, + device="meta", + ), + ] + + config = dict( + model_type="llama", + num_attention_heads=8, + hidden_size=128, + num_hidden_layers=3, + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: MixedStorageModel) + cpu_layouts = build_cpu_cache_layouts() + model = MixedStorageModel.__new__(MixedStorageModel) + model.config, model.args = config, args + model.tp_rank_, model.tp_world_size_, model.max_req_num = 0, 1, 1 + prepare_model_states(model) + model.att_state_manager.finish_state_registration() + runtime_layouts = model.att_state_manager.get_cpu_cache_operator().get_page_layouts() + assert set(cpu_layouts) == set(runtime_layouts) == {"main_gqa_bf16_state", "main_gqa_fp16_state"} + assert cpu_layouts == runtime_layouts diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py new file mode 100644 index 0000000000..a50ff4fc71 --- /dev/null +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -0,0 +1,109 @@ +"""Verification commits State-owned recurrent data without model cursor aliases.""" + +from __future__ import annotations + + +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import GDNState +from lightllm.common.attention_state.allocators import StatePageAllocator +from lightllm.server.router.model_infer.mtp_speculative.utils import verify_mtp_tokens +from unit_tests.common.att_resources.test_bounded_state import make_state + + +def make_manager(device: str | torch.device | int): + manager = AttentionStateManager(request_capacity=3) + manager.register_state("main", make_state(device=device)) + manager.register_state("main", make_state(GDNState, device=device)) + manager.register_state("draft0", make_state(GDNState, device=device)) + manager.bounded_att_big_page_allocator = StatePageAllocator(3) + manager.bounded_att_small_page_allocator = StatePageAllocator(2) + manager.finish_state_registration() + return manager + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_commit_selects_state_for_one_role(): + manager = make_manager("cuda") + for state in manager.get_bounded_attention_states(): + for slot in range(3): + state.ssm_state[:, slot].fill_(slot + 10) + state.mtp_state_index.fill_(1) + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="main", + ) + for state in (manager.main_kda_conv_fp32_ssm_fp32_state, manager.main_gdn_conv_fp32_ssm_fp32_state): + assert state.mtp_state_index.tolist() == [2, 1, 1] + selected = state.mtp_state_index[0].item() + assert (state.ssm_state[:, selected] == 12).all() + assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [1, 1, 1] + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([1], dtype=torch.int32, device="cuda"), + role="draft0", + ) + assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 1] + assert manager.main_kda_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [2, 1, 1] + assert manager.main_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [2, 1, 1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_commit_without_recurrent_state_and_invalid_role(): + manager = AttentionStateManager(request_capacity=3) + manager.finish_state_registration() + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="main", + ) + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="draft0", + ) + with pytest.raises(AssertionError, match="unknown model role"): + manager.commit_bounded_att_state( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="target", + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("widths", [(3, 3, 3), (3, 1, 2)]) +def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[int, ...]): + manager = make_manager("cuda") + backend = SimpleNamespace( + model=SimpleNamespace( + att_state_manager=manager, + req_manager=SimpleNamespace( + req_sampling_params_manager=SimpleNamespace( + req_to_next_token_ids=torch.tensor([[-1, 10, 11]] * 3, device="cuda") + ) + ), + ) + ) + reqs, starts, tokens, expected_mask = [], [], [], [] + for req, width, accepted in zip((2, 0, 1), widths, (3, 1, 2)): + starts.append(len(reqs)) + reqs.extend([req] * width) + tokens.extend(99 if offset == accepted - 1 else 10 + offset for offset in range(width)) + expected_mask.extend(int(offset < accepted) for offset in range(width)) + accept_lengths, accepted_index = verify_mtp_tokens( + backend=backend, + next_token_ids=torch.tensor(tokens, device="cuda"), + b_req_idx=torch.tensor(reqs, dtype=torch.int32, device="cuda"), + b_req_mtp_start_loc=torch.tensor(starts, dtype=torch.int32, device="cuda"), + b_mtp_index=torch.tensor([step for width in widths for step in range(width)], dtype=torch.int32, device="cuda"), + ) + assert accept_lengths.tolist() == [3, 1, 2] + assert accepted_index.tolist() == expected_mask + assert manager.main_kda_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] + assert manager.main_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] + assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 0, 0] diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py new file mode 100644 index 0000000000..f66767b9f4 --- /dev/null +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -0,0 +1,320 @@ +"""Exercise production PD transporters between independent CUDA processes.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from unit_tests.common.att_resources.model_helpers import main_linear_state + +import importlib.util +import os +import pickle +import time +import traceback + +import numpy as np +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.pd_io_struct import PDAgentMetadata, PDChunckedTransTask +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page +from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args +from unit_tests.common.att_resources.test_model_resources import make_resources + +if TYPE_CHECKING: + from collections.abc import Sequence + from lightllm.common.attention_state.manager import AttentionStateManager + from lightllm.common.attention_state.full_att.packed_mla.packed_mla_state import PackedMLAFullAttentionState + from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclKVTransporter + from lightllm.server.router.model_infer.mode_backend.pd.nixl_kv_transporter import NixlKVTransporter + from multiprocessing.connection import Connection + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +@pytest.mark.parametrize("backend", ["nccl", "nixl"]) +def test_pd_network_restores_packed_kv_accepted_linear_state_and_full_tail(backend: str): + if backend == "nixl" and importlib.util.find_spec("nixl") is None: + pytest.skip("NIXL unavailable") + context = mp.get_context("spawn") + pipes = [context.Pipe() for _ in range(2)] + processes = [ + context.Process(target=_worker, args=(child, backend, sender)) + for (_, child), sender in zip(pipes, (True, False)) + ] + try: + for process in processes: + process.start() + for _, child in pipes: + child.close() + metadata = [_receive(parent, "metadata") for parent, _ in pipes] + for (parent, _), peer in zip(pipes, reversed(metadata)): + parent.send(peer) + source, target = [_receive(parent, "result") for parent, _ in pipes] + assert source.keys() == target.keys() + for field in source: + np.testing.assert_array_equal(source[field], target[field], err_msg=field) + for parent, _ in pipes: + parent.send(None) + for process in processes: + process.join(timeout=30) + assert process.exitcode == 0 + finally: + for process in processes: + if process.is_alive(): + process.terminate() + if process.pid is not None: + process.join(timeout=10) + if process.is_alive(): + process.kill() + process.join(timeout=10) + for parent, child in pipes: + parent.close() + child.close() + + +def _notification( + transporter: NcclKVTransporter | NixlKVTransporter, + stage: str, + pending: list[PDChunckedTransTask], + key: str | None = None, +): + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + for messages in transporter.get_new_notifs().values(): + pending.extend(pickle.loads(message) for message in messages) + # Notifications are grouped by peer, so a later page's request can be + # returned before the current page's done. Match tasks as production does. + for index, task in enumerate(pending): + if task.error_info is not None: + raise AssertionError(task.error_info) + if task.write_stage == stage and (key is None or task.get_key() == key): + return pending.pop(index) + time.sleep(0.002) + raise TimeoutError(f"PD {stage} notification did not arrive") + + +def _task(peer: PDAgentMetadata, kind: str): + return PDChunckedTransTask( + request_id=123, + start_kv_index=0 if kind == "kv" else 5, + end_kv_index=5, + time_out_secs=60, + pd_master_node_id=0, + prefill_dp_index=0, + decode_dp_index=0, + src_device_id=0, + dst_device_id=1, + mem_indexes=[3, 0, 7, 1, 9] if kind == "kv" else [], + prefill_agent_name=None, + prefill_agent_metadata=None, + prefill_num_pages=None, + prefill_page_reg_desc=None, + decode_agent_name=peer.agent_name, + decode_agent_metadata=peer.agent_metadata, + decode_num_pages=peer.num_pages, + decode_page_reg_desc=peer.page_reg_desc, + first_gen_token_id=42, + first_gen_token_logprob=-0.5, + src_page_index=1 if kind == "kv" else 2, + page_kind=kind, + req_idx=1, + ) + + +def _snapshot( + resources: AttentionStateManager, + full: PackedMLAFullAttentionState, + indices: Sequence[int] | torch.Tensor, + req_idx: int, + cursor: int, +): + result = {} + for name, field in full_state_test_fields(full, full_state_test_buffers(full)).items(): + if "token" in field.axes: + ids = torch.tensor(indices, device=field.tensor.device) + result[name] = field.tensor.index_select(field.axes.index("token"), ids).float().cpu().numpy() + state = main_linear_state(resources) + result["conv"] = state.conv_state[:, req_idx, ..., cursor : cursor + state.config.conv_width].float().cpu().numpy() + result["ssm"] = state.ssm_state[:, req_idx * 3 + cursor].cpu().numpy() + tail = resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, req_idx] + result["tail"] = ( + torch.stack([tail[:, position % tail.shape[1]] for position in range(1, 5)], dim=1).float().cpu().numpy() + ) + return result + + +def _worker(connection: Connection, backend: str, sender: bool): + transporter = None + pending = [] + try: + device = 0 if sender else 1 + torch.cuda.set_device(device) + torch.set_num_threads(2) + set_env_start_args(StartArgs(run_mode="prefill" if sender else "decode", pd_kv_page_size=32)) + get_env_start_args.cache_clear() + resources, full = make_resources( + glm=True, + tp=1, + rank=0, + mtp=2, + ) + buffer = allocate_transfer_buffer(resources, 3, device=f"cuda:{device}") + buffer.fill_(93 if sender else 11) + copy_stream = torch.cuda.Stream() + if backend == "nccl": + from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclKVTransporter + + transporter = NcclKVTransporter( + node_id=os.getpid(), + tp_idx=device, + kv_move_buffer=buffer, + host_ip="127.0.0.1", + ) + else: + from lightllm.server.router.model_infer.mode_backend.pd.nixl_kv_transporter import NixlKVTransporter + + transporter = NixlKVTransporter(os.getpid(), device, buffer) + connection.send( + ( + "metadata", + PDAgentMetadata( + agent_name=transporter.agent_name, + agent_metadata=transporter.agent_metadata, + num_pages=transporter.num_pages, + page_reg_desc=transporter.local_page_mem_desc, + ), + ) + ) + assert connection.poll(60), "peer metadata did not arrive" + peer = connection.recv() + if sender: + for i, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): + if "token" in field.axes: + values = torch.arange(field.tensor.numel(), device=device).reshape(field.tensor.shape) % 17 + i + 1 + field.tensor.copy_(values) + state = main_linear_state(resources) + for i, tensor in enumerate( + (state.conv_state, state.ssm_state, resources.main_compressor_indexer_w4_s4_bf16_state.buffer) + ): + tensor.copy_(torch.arange(tensor.numel(), device=device).reshape(tensor.shape) % (23 + i * 7)) + expected = _snapshot( + resources=resources, + full=full, + indices=[3, 0, 7, 1, 9], + req_idx=1, + cursor=0, + ) + for kind in ("kv", "att_state"): + task = _task(peer, kind) + with torch.cuda.stream(copy_stream): + copy_page( + node_att_managers=[resources], + pd_page_buffers=buffer, + pd_page_index=task.src_page_index, + dp_rank_in_node=0, + tp_world_size=1, + token_indices=torch.tensor(task.mem_indexes, dtype=torch.long, device=device), + task_kind=kind, + req_idx=1, + state_token_len=task.end_kv_index, + write_to_pd_page=True, + ) + # The production sender publishes a page only after its copy event completes. + copy_stream.synchronize() + transporter.send_write_request_task_to_decode_node(task) + ready = _notification( + transporter=transporter, + stage="ready", + pending=pending, + key=task.get_key(), + ) + assert ready.get_key() == task.get_key() + task.dst_page_index = ready.dst_page_index + task.xfer_handle = transporter.write_blocks_paged(task) + deadline = time.monotonic() + 60 + while transporter.check_task_status(task) == "PROC": + assert time.monotonic() < deadline, "PD transfer timed out" + time.sleep(0.002) + assert transporter.check_task_status(task) == "DONE" + transporter.send_write_done_task_to_decode_node(task) + transporter.release_xfer_handle(task.xfer_handle) + else: + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.zero_() + pool = main_linear_state(resources).big_pages_pool + pool.conv_state.fill_(41) + pool.ssm_state.fill_(43) + for kind in ("kv", "att_state"): + request = _notification(transporter, "request", pending) + assert request.page_kind == kind + request.dst_page_index = 2 if kind == "kv" else 1 + transporter.send_write_ready_task_to_prefill_node(request) + done = _notification( + transporter=transporter, + stage="done", + pending=pending, + key=request.get_key(), + ) + assert done.get_key() == request.get_key() + if backend == "nccl": + deadline = time.monotonic() + 60 + while transporter.check_task_status(request) == "PROC": + assert time.monotonic() < deadline, "PD receive timed out" + time.sleep(0.002) + assert transporter.check_task_status(request) == "DONE" + # Match production: unpack on a separate stream after the done notification. + with torch.cuda.stream(copy_stream): + copy_page( + node_att_managers=[resources], + pd_page_buffers=buffer, + pd_page_index=request.dst_page_index, + dp_rank_in_node=0, + tp_world_size=1, + token_indices=torch.tensor([8, 12, 10, 15, 11], dtype=torch.long, device=device) + if kind == "kv" + else torch.tensor([], dtype=torch.long, device=device), + task_kind=kind, + req_idx=2, + state_token_len=request.end_kv_index, + write_to_pd_page=False, + ) + copy_stream.synchronize() + expected = _snapshot( + resources=resources, + full=full, + indices=[8, 12, 10, 15, 11], + req_idx=2, + cursor=0, + ) + assert main_linear_state(resources).mtp_state_index[2] == 0 + assert not main_linear_state(resources).conv_state[:, :2].any() + assert not main_linear_state(resources).ssm_state[:, :6].any() + assert not resources.main_compressor_indexer_w4_s4_bf16_state.buffer[:, :2].any() + assert not read_kv_storage(full)[:, [0, 1, 2, 3]].any() + assert (pool.conv_state == 41).all() + assert (pool.ssm_state == 43).all() + assert (buffer[0] == (93 if sender else 11)).all() + connection.send(("result", expected)) + assert connection.poll(60), "parent did not acknowledge completion" + connection.recv() + except BaseException: + connection.send(("error", traceback.format_exc())) + raise + finally: + if transporter is not None: + transporter.shutdown() + connection.close() + + +def _receive(connection: Connection, stage: str): + assert connection.poll(120), f"PD worker did not report {stage}" + kind, result = connection.recv() + assert kind == stage, result + return result diff --git a/unit_tests/common/att_resources/test_no_attention_draft.py b/unit_tests/common/att_resources/test_no_attention_draft.py new file mode 100644 index 0000000000..2b27dbebab --- /dev/null +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.test_capacity import definition +from unit_tests.common.basemodel.test_overlap_utils import _make_decode_input, _make_prefill_input + + +def draft_model(mode: str): + resources = AttentionStateManager(request_capacity=3) + resources.register_state("main", definition("gqa")) + draft = LlamaTpPartModel.__new__(LlamaTpPartModel) + draft.is_mtp_draft_model = True + draft.args = StartArgs(mtp_mode=mode, data_type="bfloat16") + draft.max_req_num = 2 + draft.layers_infer = [SimpleNamespace(layer_num_=2)] + draft.config = dict( + num_attention_heads=2, + num_key_value_heads=1, + hidden_size=8, + num_hidden_layers=1, + ) + draft.tp_rank_, draft.tp_world_size_ = 0, 1 + main = LlamaTpPartModel.__new__(LlamaTpPartModel) + main.args, main.config = draft.args, dict(draft.config, num_hidden_layers=2) + main.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(2)] + main.att_state_manager = resources + draft.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + draft.main_model, draft.mtp_previous_draft_models = main, [] + return draft, resources + + +@pytest.mark.parametrize("mode", ["vanilla_no_att", "eagle_no_att"]) +@pytest.mark.parametrize("prefill", [False, True]) +def test_ffn_only_draft_has_no_state_or_attention_backend(monkeypatch: pytest.MonkeyPatch, mode: str, prefill: bool): + draft, resources = draft_model(mode) + draft.prepare_attention_states() + resources.finish_state_registration() + draft._init_att_backend() + draft._init_att_backend1() + assert draft.att_state_manager is resources + assert not hasattr(draft, "full_att_state") + assert not resources.get_full_attention_states("draft0") + assert not any(name.startswith("draft0_") and name.endswith("_state") for name in vars(resources)) + + draft.infer_state_class = InferStateInfo + draft.return_all_prompt_logics = False + draft.hidden_collector_prototype = SimpleNamespace(new_instance=lambda: None) + draft.req_manager = object() + draft._select_mem_indexes = lambda value: torch.zeros(value.batch_size, dtype=torch.int32) + monkeypatch.setattr("lightllm.common.basemodel.basemodel.dist_group_manager.get_group", lambda index: None) + model_input = _make_prefill_input([10, 11], req_idx=0, is_decode_req=False) if prefill else _make_decode_input() + state = draft._create_inferstate(model_input) + assert not hasattr(state, "full_att_state") + assert state.att_state_manager is resources + assert state.is_draft_model + assert state.prefill_att_state is None and state.decode_att_state is None + assert state.req_manager is draft.req_manager + state.init_att_state() + if not prefill: + state.copy_for_cuda_graph(draft._create_inferstate(model_input)) + + +@pytest.mark.parametrize("mode", ["vanilla_with_att", "eagle_with_att"]) +def test_attention_draft_prepares_its_own_state_before_shared_allocation(mode: str): + draft, resources = draft_model(mode) + draft.prepare_attention_states() + assert draft.att_state_manager is resources + resources.finish_state_registration() + (state,) = resources.get_full_attention_states("draft0") + assert not hasattr(draft, "full_att_state") + assert state.state_name == "draft0_gqa_bf16_state" + assert state.config.logical_layer_groups == ((2,),) + assert read_kv_storage(state) is None diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py new file mode 100644 index 0000000000..032f7674eb --- /dev/null +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import gc +from typing import TYPE_CHECKING + +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNPagesPool, KDAPagesPool +from lightllm.common.attention_state.bounded_att import utils + +if TYPE_CHECKING: + from easydict import EasyDict + from multiprocessing.queues import Queue + + +@pytest.mark.parametrize("pool_cls", [KDAPagesPool, GDNPagesPool]) +@pytest.mark.parametrize("dp_fetch", [False, True]) +@pytest.mark.parametrize("pin_memory", [False, True]) +def test_page_pool_allocation_strategy( + cpu_args: EasyDict, + pool_cls: type[KDAPagesPool] | type[GDNPagesPool], + dp_fetch: bool, + pin_memory: bool, +) -> None: + if pin_memory and not torch.cuda.is_available(): + pytest.skip("CUDA required for pinned memory") + cpu_args.enable_dp_prompt_cache_fetch = dp_fetch + config = ConvSSMStateConfig( + logical_layer_ids=(0, 2), + global_k_heads=1, + global_v_heads=1, + key_dim=4, + value_dim=4, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + pool = pool_cls(size=2, config=config, pin_memory=pin_memory) + assert pool.conv_state.shape == (2, 2, 12, 3) + assert pool.ssm_state.shape == (2, 2, 1, 4, 4) + assert pool.conv_state.dtype == torch.bfloat16 + assert pool.ssm_state.dtype == torch.float32 + for tensor in (pool.conv_state, pool.ssm_state): + assert tensor.is_shared() == dp_fetch + assert tensor.is_pinned() == pin_memory + tensor.fill_(7) + assert torch.all(tensor == 7) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_shared_registration_lives_until_last_view( + cpu_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + unregistered = [] + unregister = utils._unregister_page_pool_storage + + def record_unregister(ptr: int, device_id: int) -> None: + unregistered.append(ptr) + unregister(ptr=ptr, device_id=device_id) + + monkeypatch.setattr(utils, "_unregister_page_pool_storage", record_unregister) + tensor = utils.create_page_pool_tensor(size=(4096,), dtype=torch.float32) + ptr = tensor.data_ptr() + view = tensor[1:] + del tensor + gc.collect() + assert not unregistered + assert view.is_pinned() + del view + gc.collect() + assert unregistered == [ptr] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_gpu_page_pool_keeps_device_allocation(cpu_args: EasyDict) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + tensor = utils.create_page_pool_tensor(size=(8,), dtype=torch.float32, device="cuda") + assert tensor.device.type == "cuda" + assert not tensor.is_pinned() + + +def test_empty_shared_pool_skips_cuda_registration( + cpu_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + + def unexpected_cuda_init() -> None: + pytest.fail("empty shared storage must not initialize CUDA") + + monkeypatch.setattr(torch.cuda, "current_device", unexpected_cuda_init) + tensor = utils.create_page_pool_tensor(size=(0, 4), dtype=torch.float32) + assert tensor.shape == (0, 4) + assert tensor.is_shared() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +@pytest.mark.parametrize("pool_cls", [KDAPagesPool, GDNPagesPool]) +def test_pinned_page_pool_survives_queue_on_another_gpu( + cpu_args: EasyDict, + pool_cls: type[KDAPagesPool] | type[GDNPagesPool], +) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + torch.cuda.set_device(0) + config = ConvSSMStateConfig( + logical_layer_ids=(0, 2), + global_k_heads=1, + global_v_heads=1, + key_dim=4, + value_dim=4, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + pool = pool_cls(size=2, config=config) + tensors = (pool.conv_state, pool.ssm_state) + pointers = [tensor.data_ptr() for tensor in tensors] + for tensor in tensors: + values = torch.arange(tensor.numel(), device="cuda", dtype=tensor.dtype).reshape(tensor.shape) + tensor.copy_(values, non_blocking=True) + torch.cuda.synchronize() + # 创建端也允许重复调用，不得重复注册或复制存储。 + pool.register_pin_memory() + ctx = mp.get_context("spawn") + queue = ctx.Queue() + result = ctx.Queue() + worker = ctx.Process(target=_read_page_pool_on_other_gpu, args=(queue, result)) + worker.start() + try: + queue.put(pool) + assert result.get(timeout=45) == [(True, False, True), (True, False, True)] + worker.join(timeout=15) + assert worker.exitcode == 0 + for tensor, ptr in zip(tensors, pointers): + assert tensor.data_ptr() == ptr + assert tensor.is_shared() and tensor.is_pinned() + assert tensor.flatten()[0].item() == 1234 + finally: + if worker.is_alive(): + worker.terminate() + worker.join(timeout=5) + queue.close() + result.close() + + +def _read_page_pool_on_other_gpu(queue: Queue, result: Queue) -> None: + torch.cuda.set_device(1) + pool = queue.get(timeout=30) + tensors = (pool.conv_state, pool.ssm_state) + was_pinned = [tensor.is_pinned() for tensor in tensors] + pointers = [tensor.data_ptr() for tensor in tensors] + pool.register_pin_memory() + pool.register_pin_memory() + states = [] + for tensor, ptr, before in zip(tensors, pointers, was_pinned): + assert tensor.data_ptr() == ptr + values = tensor.to(device="cuda", non_blocking=True) + torch.cuda.synchronize() + expected = torch.arange(tensor.numel(), device="cuda", dtype=tensor.dtype).reshape(tensor.shape) + assert torch.equal(values, expected) + tensor.flatten()[0] = 1234 + states.append((tensor.is_shared(), before, tensor.is_pinned())) + result.put(states) diff --git a/unit_tests/common/att_resources/test_paged_storage.py b/unit_tests/common/att_resources/test_paged_storage.py new file mode 100644 index 0000000000..e146261f11 --- /dev/null +++ b/unit_tests/common/att_resources/test_paged_storage.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +import pytest + +from lightllm.common.attention_state.allocators import FullAttKVAllocator + + +def test_kv_cache_allocation_requires_complete_pages(): + allocator = FullAttKVAllocator(16, page_size=4, pin_memory=False) + with pytest.raises(AssertionError, match="complete KV pages"): + allocator.alloc(5) + allocated = allocator.alloc(8) + assert len(allocated) == 8 + assert not set(allocator.hold_token_indices) & set(allocated.tolist()) + allocator.free(allocated) + assert allocator.can_use_mem_size == 16 + + +@pytest.mark.parametrize("size,page_size", [(5, 4), (-4, 4), (4, 0)]) +def test_kv_allocator_capacity_requires_complete_pages(size: int, page_size: int): + with pytest.raises(AssertionError, match="complete pages"): + FullAttKVAllocator(size, page_size=page_size, pin_memory=False) diff --git a/unit_tests/common/att_resources/test_quantized_full_state.py b/unit_tests/common/att_resources/test_quantized_full_state.py new file mode 100644 index 0000000000..de13eb28fc --- /dev/null +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -0,0 +1,388 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + + +from typing import TYPE_CHECKING +from math import prod + +import pytest +import torch + +from lightllm.common.attention_state.full_att import ( + GroupQuantizedAttentionConfig, + GroupQuantizedFullAttentionState, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, + PackedMLAFullAttentionState, +) +from lightllm.common.attention_state.types import ( + AxisMapping, + StateTransferLayout, + TransferRegion, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + from lightllm.common.attention_state.full_att.base_state import FullAttentionState + + +KINDS = ("int4", "int8", "fp8_tensor", "fp8_head", "packed", "packed_fp8") + + +def make_quant_state(kind: str, device: str | torch.device | int): + if kind in ("int4", "int8"): + config = GroupQuantizedAttentionConfig( + logical_layer_groups=((1,), (3,)), + global_head_num=2, + head_dim=16, + scale_dtype=torch.float16, + bits=int(kind[-1]), + ) + prepared_full_state = GroupQuantizedFullAttentionState(config, page_size=4, device=device) + prepared_full_state.allocate_storage(8) + return prepared_full_state + if kind.startswith("fp8_"): + per_head = kind == "fp8_head" + config = StaticFP8AttentionConfig( + logical_layer_groups=((1,), (3,)), + global_head_num=2, + head_dim=16, + scale_mode="per_head" if per_head else "per_tensor", + ) + scales = ( + torch.arange( + start=1, + end=9 if per_head else 5, + dtype=torch.float32, + device=device, + ).reshape(2, -1) + / 16 + ) + prepared_full_state = StaticFP8FullAttentionState(config, page_size=4, device=device) + bind_test_fp8_calibration( + state=prepared_full_state, + scales=scales, + q_scales=torch.ones((2, config.head_num), dtype=torch.float32) if per_head else None, + ) + prepared_full_state.allocate_storage(8) + return prepared_full_state + fp8 = kind == "packed_fp8" + config = PackedMLAAttentionConfig( + logical_config=PackedMLALogicalConfig( + latent_dim=576, + logical_layer_groups=((1,), (3,)), + format="packed_group128_fp8" if fp8 else "packed_bf16", + ) + ) + prepared_full_state = PackedMLAFullAttentionState(config, page_size=4, device=device) + prepared_full_state.allocate_storage(8) + return prepared_full_state + + +def transfer_layout(state: FullAttentionState): + offset, regions = 0, [] + for name, field in full_state_test_fields(state, full_state_test_buffers(state)).items(): + shape = tuple(4 if axis == "token" else dim for axis, dim in zip(field.axes, field.tensor.shape)) + itemsize = field.tensor.element_size() + offset = (offset + itemsize - 1) // itemsize * itemsize + regions.append( + TransferRegion( + field=name, + byte_start=offset, + byte_end=offset + prod(shape) * itemsize, + dtype=field.tensor.dtype, + axes=field.axes, + shape=shape, + axis_mappings=tuple( + AxisMapping( + axis=axis, + local_start=0, + external_start=0, + length=dim, + ) + for axis, dim in zip(field.axes, shape) + if axis not in ("layer", "token") + ), + ) + ) + offset += prod(shape) * itemsize + return StateTransferLayout( + state_name=state.state_name, + tokens_per_page=4, + byte_start=0, + byte_end=(offset + 15) // 16 * 16, + regions=tuple(regions), + ) + + +def assert_token_fields_equal( + source: FullAttentionState, + source_ids: Sequence[int] | torch.Tensor, + target: FullAttentionState, + target_ids: Sequence[int] | torch.Tensor, +): + src = full_state_test_fields(source, full_state_test_buffers(source)) + dst = full_state_test_fields(target, full_state_test_buffers(target)) + for name in src: + assert torch.equal( + src[name].tensor[:, source_ids].contiguous().view(torch.uint8), + dst[name].tensor[:, target_ids].contiguous().view(torch.uint8).to(src[name].tensor.device), + ), name + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_quantized_cpu_pd_dp_and_token_copy_include_every_field( + kind: str, + device: str | torch.device | int, + monkeypatch: pytest.MonkeyPatch, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + source, target = make_quant_state(kind, device), make_quant_state(kind, device) + for i, field in enumerate(full_state_test_fields(source, full_state_test_buffers(source)).values()): + field.tensor.fill_(i + 1) + field.tensor[:, 0].fill_(i + 11) + if device == "cuda": + source.copy_full_att_tokens( + torch.tensor([0, 2], dtype=torch.long, device=device), torch.tensor([4, 5], dtype=torch.long, device=device) + ) + assert_token_fields_equal( + source=source, + source_ids=[0, 2], + target=source, + target_ids=[4, 5], + ) + layout = transfer_layout(source) + cpu = torch.zeros((2, layout.byte_end), dtype=torch.uint8, pin_memory=device == "cuda") + if device == "cpu": + with pytest.raises(AssertionError, match="CUDA State storage"): + source.write_to_cpu( + token_indices=torch.tensor([0, 1, 2, 3, 4], device=device), + cpu_page_buffers=cpu, + cpu_page_indices=[1, 0], + layout=layout, + ) + else: + source.write_to_cpu( + token_indices=torch.tensor([0, 1, 2, 3, 4], dtype=torch.long, device=device), + cpu_page_buffers=cpu, + cpu_page_indices=[1, 0], + layout=layout, + ) + if device == "cuda": + torch.cuda.synchronize() + target.read_from_cpu( + cpu_page_buffers=cpu, + cpu_page_indices=[1, 0], + token_indices=torch.tensor([1, 2, 3, 4, 5], dtype=torch.long, device=device), + layout=layout, + ) + assert_token_fields_equal( + source=source, + source_ids=[0, 1, 2, 3, 4], + target=target, + target_ids=[1, 2, 3, 4, 5], + ) + page = torch.empty(layout.byte_end, dtype=torch.uint8, device=device) + if device == "cpu": + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + source.write_to_pd_page( + token_indices=torch.tensor([0, 3], dtype=torch.long, device=device), + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + layout=layout, + ) + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + target.read_from_pd_page( + pd_page_buffers=page.unsqueeze(0), + token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), + pd_page_index=0, + layout=layout, + ) + else: + source.write_to_pd_page( + token_indices=torch.tensor([0, 3], dtype=torch.long, device=device), + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + layout=layout, + ) + target.read_from_pd_page( + pd_page_buffers=page.unsqueeze(0), + token_indices=torch.tensor([6, 7], dtype=torch.long, device=device), + pd_page_index=0, + layout=layout, + ) + assert_token_fields_equal( + source=source, + source_ids=[0, 3], + target=target, + target_ids=[6, 7], + ) + src_indices = torch.tensor([0, 1], dtype=torch.int32, device=device) + dst_indices = torch.tensor([0, 2], dtype=torch.int64, device=device) + + def reject_host_list(tensor: torch.Tensor) -> list[int]: + pytest.fail("DP token indices must not be converted to Python lists") + + tensor_to = torch.Tensor.to + + def reject_index_conversion(tensor: torch.Tensor, *args, **kwargs): + if tensor is src_indices or tensor is dst_indices: + pytest.fail("DP State must use the supplied CUDA indices directly") + return tensor_to(tensor, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "tolist", reject_host_list) + patch.setattr(torch.Tensor, "to", reject_index_conversion) + if device == "cpu": + with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): + target.copy_tokens_from_dp( + src_token_indices=src_indices, + dst_token_indices=dst_indices, + remote_dp_full_att_state=source, + ) + return + target.copy_tokens_from_dp( + src_token_indices=src_indices, + dst_token_indices=dst_indices, + remote_dp_full_att_state=source, + ) + assert_token_fields_equal( + source=source, + source_ids=[0, 1], + target=target, + target_ids=[0, 2], + ) + + +@pytest.mark.parametrize("kind", KINDS) +def test_cuda_compute_write_preserves_unrelated_fields(kind: str): + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + state = make_quant_state(kind, "cuda") + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + if isinstance(state, PackedMLAFullAttentionState): + state.get_layer_view(1).indexer_k.fill_(29) + dims = state.config.latent_dim if isinstance(state, PackedMLAFullAttentionState) else state.config.head_dim + heads = 1 if isinstance(state, PackedMLAFullAttentionState) else 2 * state.config.head_num + kv = torch.randn((2, heads, dims), dtype=torch.bfloat16, device="cuda") + ids = torch.tensor([1, 6], dtype=torch.int32, device="cuda") + state.write_full_att_kv(1, ids, kv) + if isinstance(state, GroupQuantizedFullAttentionState): + quant = read_kv_storage(state)[0, ids].reshape( + shape=(2, heads, -1, state.config.group_size * state.config.bits // 8), + ) + if state.config.bits == 4: + raw = quant.to(torch.uint8) + quant = torch.stack(((raw & 15).to(torch.int16) - 7, (raw >> 4).to(torch.int16) - 7), -1).flatten(-2) + dequant = (quant.float() * torch.cat((state.k_scales, state.v_scales), dim=2)[0, ids, ..., None]).reshape_as(kv) + assert torch.all( + (dequant - kv.float()).abs() <= torch.cat((state.k_scales, state.v_scales), dim=2)[0, ids].amax() * 2 + ) + elif isinstance(state, StaticFP8FullAttentionState): + scales = state.scales[0] + if state.config.scale_mode == "per_tensor": + scales = scales.repeat_interleave(state.config.head_num) + dequant = read_kv_storage(state)[0, ids].view(torch.float8_e4m3fn).float() * scales[None, :, None] + torch.testing.assert_close( + actual=dequant, + expected=kv.float(), + rtol=0.13, + atol=0.03, + ) + else: + view = state.get_layer_view(1) + assert (view.indexer_k == 29).all() + if state.config.dtype != torch.float8_e4m3fn: + assert torch.equal(view.kv[ids], kv) + else: + fields = full_state_test_fields(state) + assert torch.isfinite(fields["kv_scale"].tensor[0, ids]).all() + assert torch.equal(fields["kv_rope"].tensor[0, ids], kv[..., 512:]) + assert not read_kv_storage(state)[1].any() + + +@pytest.mark.parametrize("kind", ["fp8_tensor", "fp8_head"]) +def test_deferred_storage_preserves_real_fp8_calibration(kind: str): + source = make_quant_state(kind, "cpu") + state = StaticFP8FullAttentionState( + source.config, + page_size=source.page_size, + device="meta", + ) + assert state.scales is None and state.q_scales is None + bind_test_fp8_calibration(state=state, scales=source.scales, q_scales=source.q_scales) + assert read_kv_storage(state) is None and not state.scales.is_meta + assert torch.equal(state.scales, source.scales) + state.device = torch.device("cpu") + state.allocate_storage(source.capacity) + assert torch.equal(state.scales, source.scales) + assert torch.equal(torch.from_numpy(state.cpu_scales), source.scales) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_int4_prefill_with_separate_storage(device, monkeypatch): + import weakref + from types import SimpleNamespace + from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonPrefillAttState + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.common.basemodel.triton_kernel.kv_copy import ppl_int4kv_copy_kv + from lightllm.common.basemodel.triton_kernel.att.prefill_att import context_flashattention_nopad + + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + state = GroupQuantizedFullAttentionState( + GroupQuantizedAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=64, + scale_dtype=torch.float32, + bits=4, + ), + device=device, + ) + state.allocate_storage(8) + (k, k_scale), (v, v_scale) = state.get_layer_view(0) + k.view(torch.uint8).fill_(0x88) + v.view(torch.uint8).fill_(0x99) + k_scale.fill_(0.25) + v_scale.fill_(0.5) + assert k.untyped_storage().data_ptr() != v.untyped_storage().data_ptr() + assert k_scale.untyped_storage().data_ptr() != v_scale.untyped_storage().data_ptr() + owner = InferStateInfo() + owner.total_token_num = owner.max_kv_seq_len = owner.max_q_seq_len = 2 + owner.b_seq_len = torch.tensor([2], dtype=torch.int32, device=device) + owner.b_req_idx = torch.tensor([0], dtype=torch.int32, device=device) + owner.b_ready_cache_len = owner.b_q_start_loc = torch.zeros(1, dtype=torch.int32, device=device) + owner.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, device=device).view(1, -1)) + prefill = Int4kvTritonPrefillAttState( + backend=SimpleNamespace(quant_group_size=8), _infer_state_weak_ref=weakref.ref(owner) + ) + prefill.init_state() + if device == "cpu": + + def dequantize(**kwargs): + assert kwargs["k"] is k and kwargs["v"] is v + assert kwargs["k_scale"] is k_scale and kwargs["v_scale"] is v_scale + kwargs["k_out"].fill_(0.25) + kwargs["v_out"].fill_(1) + + def attention(**kwargs): + kwargs["o"].copy_(kwargs["v"].expand(-1, 2, -1)) + + monkeypatch.setattr(ppl_int4kv_copy_kv, "dequantize_int4kv", dequantize) + monkeypatch.setattr(context_flashattention_nopad, "context_attention_fwd_contiguous_kv", attention) + q = torch.randn(2, 2, 64, device=device) + output = prefill.prefill_att(q, (k, k_scale), (v, v_scale)) + torch.testing.assert_close(output, torch.ones_like(q)) diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py new file mode 100644 index 0000000000..ca97cc7b70 --- /dev/null +++ b/unit_tests/common/att_resources/test_quantized_paged_storage.py @@ -0,0 +1,37 @@ +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +import torch + +from lightllm.common.attention_state.full_att import ( + GroupQuantizedAttentionConfig, + GroupQuantizedFullAttentionState, +) + + +def test_quantized_kv_buffers_include_the_full_hold_page(): + allocations = [] + for bits in (8, 4): + state = GroupQuantizedFullAttentionState( + GroupQuantizedAttentionConfig( + logical_layer_groups=((0,), (1,), (2,)), + global_head_num=2, + head_dim=16, + scale_dtype=torch.float16, + bits=bits, + ), + page_size=4, + device="cpu", + ) + state.allocate_storage(32) + allocations.extend(tuple(buffer.shape) for buffer in full_state_test_buffers(state).values()) + + assert allocations == [ + (3, 36, 2, 16), + (3, 36, 2, 16), + (3, 36, 2, 2), + (3, 36, 2, 2), + (3, 36, 2, 8), + (3, 36, 2, 8), + (3, 36, 2, 2), + (3, 36, 2, 2), + ] diff --git a/unit_tests/common/att_resources/test_role_transfers.py b/unit_tests/common/att_resources/test_role_transfers.py new file mode 100644 index 0000000000..88f02eeaf1 --- /dev/null +++ b/unit_tests/common/att_resources/test_role_transfers.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from typing import TYPE_CHECKING +import pytest +import torch + +from unit_tests.common.att_resources.model_helpers import main_linear_state +from unit_tests.common.att_resources.test_model_resources import make_resources + +if TYPE_CHECKING: + from easydict import EasyDict + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("glm", [False, True]) +def test_main_and_draft_transfers_copy_pd_and_dp_windows( + glm: bool, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +): + pd_args.pd_kv_page_size = 128 + source, _ = make_resources( + glm=glm, + tp=1, + rank=0, + mtp=2, + draft_layers=2, + ) + target, _ = make_resources( + glm=glm, + tp=1, + rank=0, + mtp=2, + draft_layers=2, + ) + for index, role in enumerate(("main", "draft0")): + state = getattr(source, f"{role}_mla_packed_bf16_state" if glm else f"{role}_gqa_bf16_state") + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + if "token" in field.axes: + field.tensor.fill_(11 + index) + if glm: + getattr(source, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer[:, 0].fill_(7 + index) + getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer[:, 2].fill_(19) + main_linear_state(source).conv_state.fill_(3) + main_linear_state(source).ssm_state.fill_(5) + page = torch.empty(source.get_pd_operator().get_page_bytes(), dtype=torch.uint8, device="cuda") + for kind, getter in ( + ("kv", "get_pd_full_att_page_layouts"), + ("att_state", "get_pd_bounded_att_page_layouts"), + ): + source.get_pd_operator().write_page( + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + kind=kind, + layouts=getattr(source.get_pd_operator(), getter)(), + **( + {"token_indices": torch.tensor([0, 1], dtype=torch.long, device="cuda")} + if kind == "kv" + else {"req_idx": 0, "state_token_len": 8} + ), + ) + target.get_pd_operator().read_page( + pd_page_buffers=page.unsqueeze(0), + pd_page_index=0, + kind=kind, + layouts=getattr(target.get_pd_operator(), getter)(), + **( + {"token_indices": torch.tensor([8, 9], dtype=torch.long, device="cuda")} + if kind == "kv" + else {"req_idx": 1, "state_token_len": 8} + ), + ) + src_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + dst_indices = torch.tensor([10, 11], dtype=torch.int64, device="cuda") + + def reject_host_list(tensor: torch.Tensor) -> list[int]: + pytest.fail("DP manager and States must preserve Tensor indices") + + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "tolist", reject_host_list) + target.get_dp_shared_operator().copy_from_dp( + remote_att_manager=source, + src_token_indices=src_indices, + dst_token_indices=dst_indices, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=0, + dst_req_idx=2, + ) + for index, role in enumerate(("main", "draft0")): + state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_gqa_bf16_state") + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + if "token" in field.axes: + assert (field.tensor[:, 8:12] == 11 + index).all() + if glm: + tail = getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer + expected_tail = torch.zeros_like(tail[:, 1]) + for position in range(4, 8): + expected_tail[:, position % tail.shape[2]].fill_(7 + index) + assert torch.equal(tail[:, 1], expected_tail) # PD restores only the prefill logical window. + assert (tail[:, 2] == 7 + index).all() # DP 直接复制相同 MTP 容量的 ring。 + + slot = source.bounded_att_big_page_allocator.alloc(1)[0] + source.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[slot], + ) + torch.cuda.synchronize() + cpu = torch.empty((1, source.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + source.get_cpu_cache_operator().write_req( + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + cpu_page_buffers=cpu, + cpu_page_indices=[0], + layouts=source.get_cpu_cache_operator().get_page_layouts(), + big_page_slot_ids=[slot], + ) + torch.cuda.synchronize() + dst_slot = target.bounded_att_big_page_allocator.alloc(1)[0] + target.get_cpu_cache_operator().read_req( + cpu_page_buffers=cpu, + cpu_page_indices=[0], + token_indices=torch.arange( + start=8, + end=16, + dtype=torch.long, + device="cuda", + ), + state_token_len=8, + layouts=target.get_cpu_cache_operator().get_page_layouts(), + big_page_slot_ids=[dst_slot], + req_idx=1, + ) + for index, role in enumerate(("main", "draft0")): + state = getattr(target, f"{role}_mla_packed_bf16_state" if glm else f"{role}_gqa_bf16_state") + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + if "token" in field.axes: + assert (field.tensor[:, 8:16] == 11 + index).all() + if glm: + # CPU 组边界恢复不修改无重叠窗口，保留先前 PD/DP 的数据。 + tail = getattr(target, f"{role}_compressor_indexer_w4_s4_bf16_state").buffer + for position in range(4, 8): + assert (tail[:, 1, position % tail.shape[2]] == 7 + index).all() + assert (tail[:, 2] == 7 + index).all() diff --git a/unit_tests/common/att_resources/test_runtime_transport.py b/unit_tests/common/att_resources/test_runtime_transport.py new file mode 100644 index 0000000000..1371cfe47e --- /dev/null +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -0,0 +1,571 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import window_states, linear_states + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from typing import TYPE_CHECKING +from unit_tests.common.att_resources.model_helpers import main_linear_state +from types import SimpleNamespace +from contextlib import nullcontext +import pickle +import traceback +import uuid + +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.server.router.model_infer import shared_att_resources_utils +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page +from unit_tests.common.att_resources.test_model_resources import make_resources +from unit_tests.common.att_resources.test_quantized_full_state import KINDS, make_quant_state, assert_token_fields_equal + +if TYPE_CHECKING: + from easydict import EasyDict + from multiprocessing.connection import Connection + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_cpu_pd_and_dp_model_layouts_cover_quantized_fields( + kind: str, + device: str, + pd_args: EasyDict, + monkeypatch: pytest.MonkeyPatch, +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + pd_args.cpu_cache_token_page_size = pd_args.pd_kv_page_size = 4 + src, dst = make_quant_state(kind, device), make_quant_state(kind, device) + managers = [] + for state in (src, dst): + manager = AttentionStateManager(request_capacity=2) + manager.register_state("main", state) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + manager.finish_state_registration() + managers.append(manager) + src, dst = [manager.get_full_attention_states("main")[0] for manager in managers] + for i, field in enumerate(full_state_test_fields(src, full_state_test_buffers(src)).values()): + if "token" in field.axes: + field.tensor.fill_(i + 3) + source_indices = torch.tensor([1, 3], dtype=torch.int32, device=device) + pd_indices = torch.tensor([4, 6], dtype=torch.int64, device=device) + cpu_indices = torch.tensor([5, 7], dtype=torch.int64, device=device) + dp_indices = torch.tensor([0, 2], dtype=torch.int64, device=device) + + def reject_host_list(tensor: torch.Tensor) -> list[int]: + pytest.fail("transfer paths must not convert Tensor indices to Python lists") + + monkeypatch.setattr(torch.Tensor, "tolist", reject_host_list) + page = allocate_transfer_buffer(managers[0], 1, device=device) + if device == "cpu": + with pytest.raises(AssertionError, match="PD pages must reside on CUDA"): + copy_page( + node_att_managers=[managers[0]], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=1, + token_indices=source_indices, + write_to_pd_page=True, + ) + else: + copy_page( + node_att_managers=[managers[0]], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=1, + token_indices=source_indices, + write_to_pd_page=True, + ) + copy_page( + node_att_managers=[managers[1]], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=1, + token_indices=pd_indices, + write_to_pd_page=False, + ) + assert_token_fields_equal( + source=src, + source_ids=[1, 3], + target=dst, + target_ids=[4, 6], + ) + + cpu_pages = torch.empty( + (1, managers[0].get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=device == "cuda" + ) + if device == "cpu": + with pytest.raises(AssertionError, match="CUDA State storage"): + managers[0].get_cpu_cache_operator().write_req( + token_indices=source_indices, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + layouts=managers[0].get_cpu_cache_operator().get_page_layouts(), + ) + else: + managers[0].get_cpu_cache_operator().write_req( + token_indices=source_indices, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + layouts=managers[0].get_cpu_cache_operator().get_page_layouts(), + ) + if device == "cuda": + torch.cuda.synchronize() + managers[1].get_cpu_cache_operator().read_req( + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + token_indices=cpu_indices, + layouts=managers[1].get_cpu_cache_operator().get_page_layouts(), + ) + assert_token_fields_equal( + source=src, + source_ids=[1, 3], + target=dst, + target_ids=[5, 7], + ) + + if device == "cpu": + with pytest.raises(AssertionError, match="DP token indices must reside on CUDA"): + managers[1].get_dp_shared_operator().copy_from_dp( + remote_att_manager=managers[0], + src_token_indices=source_indices, + dst_token_indices=dp_indices, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=0, + dst_req_idx=0, + ) + return + managers[1].get_dp_shared_operator().copy_from_dp( + remote_att_manager=managers[0], + src_token_indices=source_indices, + dst_token_indices=dp_indices, + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=0, + dst_req_idx=0, + ) + assert_token_fields_equal( + source=src, + source_ids=[1, 3], + target=dst, + target_ids=[0, 2], + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("glm", [False, True]) +@pytest.mark.parametrize("source_device", [0, 1]) +@pytest.mark.parametrize("copy_big_pages", [False, True]) +def test_dp_fetch_restores_baseline_runtime_and_full_token_fields( + monkeypatch: pytest.MonkeyPatch, + glm: bool, + source_device: int, + copy_big_pages: bool, +): + if source_device >= torch.cuda.device_count(): + pytest.skip("two CUDA devices required") + from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans as module + + with torch.cuda.device(source_device): + source, source_full = make_resources( + glm=glm, + tp=1, + rank=0, + mtp=2, + ) + with torch.cuda.device(0): + target, target_full = make_resources( + glm=glm, + tp=1, + rank=0, + mtp=2, + ) + if source_device != 0: + # 此测试直接传入另一张卡的 Tensor，未走 IPC 导入；装配阶段先建立 peer 映射。 + # 生产拷贝入口只使用已映射的资源，不负责打开设备间访问。 + from cuda.bindings import runtime + + (error,) = runtime.cudaDeviceEnablePeerAccess(source_device, 0) + assert error in (runtime.cudaError_t.cudaSuccess, runtime.cudaError_t.cudaErrorPeerAccessAlreadyEnabled) + for i, field in enumerate(full_state_test_fields(source_full, full_state_test_buffers(source_full)).values()): + if "token" in field.axes: + field.tensor.fill_(i + 5) + main_linear_state(source).conv_state.fill_(-1) + main_linear_state(source).conv_state[:, 1, ..., :2].fill_(13) + main_linear_state(source).ssm_state[:, 3].fill_(17) + if glm: + source.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 1].fill_(23) + target.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 2].fill_(29) + torch.cuda.synchronize(source_device) + # 同时验证运行态与指定大页复制，保留其他页及小页快照。 + for pool in (main_linear_state(target).big_pages_pool, main_linear_state(target).small_pages_pool): + pool.conv_state.fill_(41) + pool.ssm_state.fill_(43) + src_big_page_ids, dst_big_page_ids = [], [] + if copy_big_pages: + src_big_page_ids = source.bounded_att_big_page_allocator.alloc(1) + dst_big_page_ids = target.bounded_att_big_page_allocator.alloc(1) + source_pool = main_linear_state(source).big_pages_pool + source_pool.conv_state[src_big_page_ids[0]].fill_(51) + source_pool.ssm_state[src_big_page_ids[0]].fill_(53) + req = SimpleNamespace( + req_idx=2, + cur_kv_len=4, + hold_kv_len=8, + hybrid_len_to_big_page_id={}, + ) + task = module.TransTask( + dst_req=req, + src_dp_rank=1, + src_req_idx=1, + src_token_indices=torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=f"cuda:{source_device}"), + dst_token_indices=torch.tensor([8, 9, 10, 11]), + src_big_page_ids=src_big_page_ids, + dst_big_page_ids=dst_big_page_ids, + ) + syncs = [] + + def barrier(group: str): + torch.cuda.synchronize(0) + torch.cuda.synchronize(source_device) + assert req.cur_kv_len == 4 + assert not req.hybrid_len_to_big_page_id + syncs.append(group) + + monkeypatch.setattr(module.dist, "barrier", barrier) + worker = module.DPKVSharedMoudle.__new__(module.DPKVSharedMoudle) + worker.backend = SimpleNamespace( + model=SimpleNamespace(att_state_manager=target), + dp_att_state_managers={1: source}, + is_hybrid_att_model=True, + is_master_in_dp=False, + node_nccl_group="node", + ) + worker.big_page_tokens = 4 + worker.kv_trans([task]) + assert syncs == ["node"] and req.cur_kv_len == 8 + assert_token_fields_equal( + source=source_full, + source_ids=[0, 1, 2, 3], + target=target_full, + target_ids=[8, 9, 10, 11], + ) + assert (main_linear_state(target).conv_state[:, 2, ..., :2] == 13).all() + assert (main_linear_state(target).ssm_state[:, 6] == 17).all() + assert main_linear_state(target).mtp_state_index[2] == 0 + assert not main_linear_state(target).conv_state[:, :2].any() + big_pool = main_linear_state(target).big_pages_pool + if copy_big_pages: + page_id = dst_big_page_ids[0] + assert (big_pool.conv_state[page_id] == 51).all() + assert (big_pool.ssm_state[page_id] == 53).all() + assert req.hybrid_len_to_big_page_id == {8: page_id} + else: + assert not req.hybrid_len_to_big_page_id + untouched = [page for page in range(big_pool.size) if page not in dst_big_page_ids] + assert (big_pool.conv_state[untouched] == 41).all() + assert (big_pool.ssm_state[untouched] == 43).all() + assert (main_linear_state(target).small_pages_pool.conv_state == 41).all() + assert (main_linear_state(target).small_pages_pool.ssm_state == 43).all() + if glm: + assert torch.equal( + target.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 2].cpu(), + source.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 1].cpu(), + ) + assert (source.main_compressor_indexer_w4_s4_bf16_state.buffer[:, 1] == 23).all() + assert not target.main_compressor_indexer_w4_s4_bf16_state.buffer[:, :2].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("p2p", [False, True]) +@pytest.mark.parametrize( + "share_pools,ipc_kinds", + [ + (False, ("pd",)), + (True, ("pd", "dp")), + (True, ("dp", "pd")), + ], +) +def test_actual_rank_storage_ipc_separates_pd_and_dp_exports( + monkeypatch: pytest.MonkeyPatch, + p2p: bool, + share_pools: bool, + ipc_kinds: tuple[str, ...], + cpu_args: EasyDict, +): + cpu_args.enable_dp_prompt_cache_fetch = share_pools + from lightllm.utils.envs_utils import get_unique_server_name + from lightllm.utils.shm_utils import create_or_link_shm + + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "att-ipc-" + uuid.uuid4().hex) + get_unique_server_name.cache_clear() + monkeypatch.setattr(shared_att_resources_utils, "get_current_rank_in_node", lambda: 0) + if p2p and torch.cuda.device_count() < 2: + pytest.skip("two CUDA devices required for IPC peer mapping") + monkeypatch.setattr(shared_att_resources_utils, "kv_trans_use_p2p", lambda: p2p) + monkeypatch.setattr(mp.reductions.reduce_tensor, "__code__", mp.reductions.reduce_tensor.__code__) + source_device = int(p2p) + with torch.cuda.device(source_device): + resources, full = make_resources( + glm=True, + tp=1, + rank=0, + mtp=2, + ) + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.fill_(9) + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer[:, 4].zero_() + main_linear_state(resources).conv_state.fill_(11) + resources.main_compressor_indexer_w4_s4_bf16_state.buffer.fill_(7) + for state in linear_states(resources): + for pool in (state.big_pages_pool, state.small_pages_pool): + if pool is not None: + pool.conv_state.zero_() + pool.ssm_state.zero_() + slot = resources.bounded_att_big_page_allocator.alloc(1) + with torch.cuda.device(source_device): + resources.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=slot, + ) + mapping = torch.zeros((3, 16), dtype=torch.int32, device=f"cuda:{source_device}") + mapping[1, 0] = 7 + torch.cuda.synchronize(source_device) + for kind in ipc_kinds: + shared_att_resources_utils.publish_att_resources( + resources=resources, + request_mapping=mapping, + kind=kind, + reader_count=1, + ) + ctx = mp.get_context("spawn") + parent, child = ctx.Pipe() + process = ctx.Process(target=_ipc_worker, args=(child, ipc_kinds)) + process.start() + child.close() + try: + assert parent.poll(45), "CUDA IPC child did not finish" + error = parent.recv() + assert error is None, error + assert (read_kv_storage(full)[:, 4] == 9).all() + window = resources.main_compressor_indexer_w4_s4_bf16_state + assert not window.buffer[:, 1, :, : window.config.candidate_dim].any() + assert torch.isneginf(window.buffer[:, 1, :, window.config.candidate_dim :]).all() + assert (main_linear_state(resources).conv_state == 11).all() + assert (main_linear_state(resources).big_pages_pool.conv_state[slot[0]] == 11).all() + assert not main_linear_state(resources).big_pages_pool.ssm_state[slot[0]].any() + if share_pools: + assert main_linear_state(resources).small_pages_pool.conv_state.flatten()[0].item() == 37 + assert resources.bounded_att_big_page_allocator.free_count == 1 + parent.send(None) + process.join(10) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + for kind in ipc_kinds: + shm = create_or_link_shm(f"rank_storage_{kind}_0", expected_size=-1, force_mode="link") + shm.close() + shm.unlink() + get_unique_server_name.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("source_tp,target_tp", [(1, 4), (4, 1), (2, 8), (8, 2)]) +def test_qwen_pd_reshards_full_and_linear_replicated_heads(source_tp: int, target_tp: int, pd_args: EasyDict): + pd_args.pd_kv_page_size = 128 + source = [ + make_resources( + glm=False, + tp=source_tp, + rank=rank, + mtp=2, + ) + for rank in range(source_tp) + ] + target = [ + make_resources( + glm=False, + tp=target_tp, + rank=rank, + mtp=2, + ) + for rank in range(target_tp) + ] + for rank, (manager, full) in enumerate(source): + for field_index, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): + for local_head in range(full.config.head_num): + head = rank * full.config.head_num // max(1, source_tp // 2) + local_head + field.tensor[:, :, local_head].fill_(100 * field_index + head) + cfg = main_linear_state(manager).config + offset = 0 + for group, (local_heads, global_heads) in enumerate( + ((cfg.num_k_heads, 2), (cfg.num_k_heads, 2), (cfg.num_v_heads, 4)) + ): + start = rank * local_heads // max(1, source_tp // global_heads) + for head in range(local_heads): + main_linear_state(manager).conv_state[:, 1, offset + head * 8 : offset + (head + 1) * 8, :2].fill_( + 10 * group + start + head + ) + offset += local_heads * 8 + start = rank * cfg.num_v_heads // max(1, source_tp // 4) + for head in range(cfg.num_v_heads): + main_linear_state(manager).ssm_state[:, 3, head].fill_(30 + start + head) + page = allocate_transfer_buffer(source[0][0], 1, device="cuda") + for kind in ("kv", "att_state"): + copy_page( + node_att_managers=[m for m, _ in source], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=source_tp, + token_indices=torch.tensor([0, 3], dtype=torch.long, device="cuda"), + task_kind=kind, + req_idx=1, + state_token_len=8, + write_to_pd_page=True, + ) + copy_page( + node_att_managers=[m for m, _ in target], + pd_page_buffers=page, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=target_tp, + token_indices=torch.tensor([4, 7], dtype=torch.long, device="cuda"), + task_kind=kind, + req_idx=2, + state_token_len=8, + write_to_pd_page=False, + ) + for rank, (manager, full) in enumerate(target): + for field_index, field in enumerate(full_state_test_fields(full, full_state_test_buffers(full)).values()): + for local_head in range(full.config.head_num): + head = rank * full.config.head_num // max(1, target_tp // 2) + local_head + assert (field.tensor[:, [4, 7], local_head] == 100 * field_index + head).all() + cfg = main_linear_state(manager).config + offset = 0 + for group, (local_heads, global_heads) in enumerate( + ((cfg.num_k_heads, 2), (cfg.num_k_heads, 2), (cfg.num_v_heads, 4)) + ): + start = rank * local_heads // max(1, target_tp // global_heads) + for head in range(local_heads): + actual = main_linear_state(manager).conv_state[:, 2, offset + head * 8 : offset + (head + 1) * 8, :2] + assert (actual == 10 * group + start + head).all() + offset += local_heads * 8 + start = rank * cfg.num_v_heads // max(1, target_tp // 4) + for head in range(cfg.num_v_heads): + assert (main_linear_state(manager).ssm_state[:, 6, head] == 30 + start + head).all() + assert main_linear_state(manager).mtp_state_index[2] == 0 + + +@pytest.mark.parametrize("kind", ["pd", "dp"]) +def test_ipc_readers_accept_different_pickle_lengths(kind: str, monkeypatch: pytest.MonkeyPatch) -> None: + """各读者独立导出；模拟 IPC 计数变化引起的长度差异，并验证补齐后仍可逐份读取。""" + resources = SimpleNamespace( + get_pd_operator=lambda: SimpleNamespace(create_ipc_copy=lambda: {"kind": "pd"}), + get_dp_shared_operator=lambda: SimpleNamespace(create_ipc_copy=lambda: {"kind": "dp"}), + ) + mapping = [[0, 1, 2]] + payloads = [] + storage = {} + + def serialize(value: tuple[object, object]) -> memoryview: + payload = pickle.dumps(value) + b"\0" * (3 * len(payloads)) + payloads.append(payload) + return memoryview(payload) + + def shared_memory(name: str, expected_size: int, force_mode: str) -> SimpleNamespace: + if force_mode == "create": + storage[name] = SimpleNamespace(buf=bytearray(expected_size), close=lambda: None) + return storage[name] + + monkeypatch.setattr(shared_att_resources_utils, "kv_trans_use_p2p", lambda: False) + monkeypatch.setattr(shared_att_resources_utils, "get_current_rank_in_node", lambda: 0) + monkeypatch.setattr(shared_att_resources_utils, "get_service_shm_name", lambda name: name) + monkeypatch.setattr(shared_att_resources_utils, "FileLock", lambda path: nullcontext()) + monkeypatch.setattr(shared_att_resources_utils, "create_or_link_shm", shared_memory) + monkeypatch.setattr(shared_att_resources_utils.ForkingPickler, "dumps", serialize) + + shared_att_resources_utils.publish_att_resources( + resources=resources, + request_mapping=mapping, + kind=kind, + reader_count=3, + ) + assert len(payloads) == 3 and len({len(payload) for payload in payloads}) == 3 + buffer = storage[f"rank_storage_{kind}_0"].buf + slot_bytes = int.from_bytes(buffer[4:8], "little") + assert slot_bytes == max(len(payload) for payload in payloads) + for index, payload in enumerate(payloads): + assert buffer[8 + index * slot_bytes : 8 + (index + 1) * slot_bytes] == payload.ljust(slot_bytes, b"\0") + for _ in payloads: + imported_resources, imported_mapping = shared_att_resources_utils.load_att_resources(rank_in_node=0, kind=kind) + assert imported_resources == {"kind": kind} + assert imported_mapping == mapping + with pytest.raises(RuntimeError, match="exhausted"): + shared_att_resources_utils.load_att_resources(rank_in_node=0, kind=kind) + + +def _ipc_worker(connection: Connection, ipc_kinds: tuple[str, ...]) -> None: + try: + torch.cuda.set_device(0) + for kind in ipc_kinds: + resources, mapping = shared_att_resources_utils.load_att_resources(0, kind=kind) + assert resources._state_registration_finished + assert resources.full_att_allocator is None + assert resources.bounded_att_big_page_allocator is None + assert resources.bounded_att_small_page_allocator is None + if kind == "dp": + for state in linear_states(resources): + assert state.big_pages_pool.conv_state.is_shared() + assert not state.big_pages_pool.conv_state.is_pinned() + resources.get_dp_shared_operator().register_page_pools_pin_memory() + resources.get_dp_shared_operator().register_page_pools_pin_memory() + for state in linear_states(resources): + for pool in (state.big_pages_pool, state.small_pages_pool): + if kind == "dp": + assert pool is not None + assert pool.conv_state.is_pinned() and pool.ssm_state.is_pinned() + copied = pool.conv_state.to(device="cuda", non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(copied.cpu(), pool.conv_state) + else: + assert pool is None + if kind == "dp": + main_linear_state(resources).small_pages_pool.conv_state.flatten()[0] = 37 + assert not hasattr(resources, "req_to_token_indexs") + assert mapping[1, 0].item() == 7 + resources.copy_full_att_tokens( + torch.tensor([0], dtype=torch.long, device="cuda"), torch.tensor([4], dtype=torch.long, device="cuda") + ) + for window in window_states(resources): + window.reset_bounded_att_state([1]) + torch.cuda.synchronize() + # 一类通道耗尽不会影响另一类通道的读取计数。 + with pytest.raises(RuntimeError, match="exhausted"): + shared_att_resources_utils.load_att_resources(0, kind=kind) + connection.send(None) + # Keep imported tensors alive until the parent has checked the writes. + connection.recv() + except BaseException: + connection.send(traceback.format_exc()) + finally: + AttentionStateManager._get_states.cache_clear() + connection.close() diff --git a/unit_tests/common/att_resources/test_state_creation.py b/unit_tests/common/att_resources/test_state_creation.py new file mode 100644 index 0000000000..dee0d9a836 --- /dev/null +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import bind_test_fp8_calibration + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from collections.abc import Callable + +import pytest +import torch + +from lightllm.common.attention_state.full_att import ( + GQAAttentionConfig, + MLAConfig, + FullAttentionState, + GQAFullAttentionState, + GroupQuantizedFullAttentionState, + StaticFP8FullAttentionState, + MLAFullAttentionState, + PackedMLAFullAttentionState, + PackedMLALogicalConfig, + create_gqa_state, + create_mla_state, + create_packed_mla_state, +) +from lightllm.common.attention_state.bounded_att import ( + BoundedAttentionState, + ConvSSMStateConfig, + create_gdn_state, + create_kda_state, +) +from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state + + +@pytest.mark.parametrize( + "kv_type, expected_type, storage_dtype, width", + [ + ("None", "gqa_bf16", torch.bfloat16, 128), + ("int4kv", "gqa_int4_group16_bf16_scale", torch.int8, 64), + ("int8kv", "gqa_int8_group16_bf16_scale", torch.int8, 128), + ("fp8kv_spt", "gqa_fp8_static_per_tensor", torch.uint8, 128), + ("fp8kv_sph", "gqa_fp8_static_per_head", torch.uint8, 128), + ], +) +def test_gqa_creation_preserves_logical_tp_config_and_deferred_capacity( + kv_type: str, + expected_type: str, + storage_dtype: torch.dtype, + width: int, +): + config = GQAAttentionConfig( + logical_layer_groups=((3,), (7,)), + global_head_num=8, + head_dim=128, + dtype=torch.bfloat16, + tp_rank=2, + tp_world_size=4, + ) + scales = torch.full((2, 4 if kv_type == "fp8kv_sph" else 2), 0.5, dtype=torch.float32) + q_scales = torch.ones((2, 2), dtype=torch.float32) if kv_type == "fp8kv_sph" else None + state = create_gqa_state( + config=config, + kv_type=kv_type, + page_size=4, + device="meta", + group_size=16, + ) + assert state.get_state_type() == expected_type + assert state.state_name is None + assert state.config.logical_layer_groups == ((3,), (7,)) + assert state.config.global_head_num == 8 and state.config.head_num == 2 + assert state.config.tp_rank == 2 and state.config.tp_world_size == 4 + assert state.config.head_dim == config.head_dim and config.dtype == torch.bfloat16 + assert state.capacity == 0 and read_kv_storage(state) is None + if kv_type.startswith("fp8"): + assert state.scales is None and state.q_scales is None + bind_test_fp8_calibration(state=state, scales=scales, q_scales=q_scales) + assert torch.equal(state.scales, scales) + if q_scales is not None: + assert torch.equal(state.q_scales, q_scales) + assert type(state.config).__bases__ == (object,) + if kv_type in ("int4kv", "int8kv"): + assert state.config.scale_dtype == config.dtype + elif kv_type.startswith("fp8"): + assert not hasattr(state.config, "dtype") + state.device = torch.device("cpu") + state.allocate_storage(8) + assert state.k_storage.shape == state.v_storage.shape == (2, 12, 2, width) + assert state.k_storage.untyped_storage().data_ptr() != state.v_storage.untyped_storage().data_ptr() + assert state.k_storage.dtype == state.v_storage.dtype == storage_dtype + if kv_type in ("int4kv", "int8kv"): + assert state.k_scales.shape == state.v_scales.shape == (2, 12, 2, 8) + assert state.k_scales.untyped_storage().data_ptr() != state.v_scales.untyped_storage().data_ptr() + + +@pytest.mark.parametrize( + "packed, kv_type, kv_dim, row_width, expected_type", + [ + (False, "None", 576, 576, "mla_bf16"), + (True, "None", 576, 648, "mla_packed_bf16"), + (True, "None", 512, 584, "mla_packed_bf16"), + (True, "fp8kv_dsa", 576, 800, "mla_packed_fp8_group128"), + ], +) +def test_mla_creation_computes_packed_width_from_logical_dimensions( + packed: bool, + kv_type: str, + kv_dim: int, + row_width: int, + expected_type: str, +): + if packed: + config = PackedMLALogicalConfig( + logical_layer_groups=((2,), (6,)), + latent_dim=kv_dim, + format=("packed_group128_fp8" if kv_type == "fp8kv_dsa" else "packed_bf16"), + tp_rank=1, + tp_world_size=4, + ) + state = create_packed_mla_state(config=config, page_size=4, device="meta") + else: + config = MLAConfig( + logical_layer_groups=((2,), (6,)), + latent_dim=kv_dim, + latent_dtype=torch.bfloat16, + tp_rank=1, + tp_world_size=4, + ) + state = create_mla_state(config=config, page_size=4, device="meta") + assert state.get_state_type() == expected_type + assert state.config.latent_dim == config.latent_dim == kv_dim + expected_dtype = torch.float8_e4m3fn if kv_type == "fp8kv_dsa" else torch.bfloat16 + if packed: + assert state.config.dtype == expected_dtype + else: + assert state.config.latent_dtype == expected_dtype + assert not hasattr(state.config, "storage_dtype") + assert not hasattr(state.config, "replicated") + assert state.capacity == 0 and read_kv_storage(state) is None + assert type(state.config).__bases__ == (object,) + assert not hasattr(state.config, "head_num") and not hasattr(state.config, "global_head_num") + state.device = torch.device("cpu") + state.allocate_storage(8) + assert read_kv_storage(state).shape == (2, 12, 1, row_width) + assert read_kv_storage(state).dtype == (torch.uint8 if kv_type == "fp8kv_dsa" else torch.bfloat16) + if packed: + # 保持旧算子读取的物理布局：FP8 indexer 紧接 KV，浮点格式保留中间的对齐间隙。 + expected_indexer_offset = 656 if kv_type == "fp8kv_dsa" else kv_dim * 2 + 12 + assert state.config.indexer_offset == expected_indexer_offset + assert state.get_layer_view(2).indexer_k.shape[0] == 12 + + +@pytest.mark.parametrize("create_state", [create_kda_state, create_gdn_state]) +def test_bounded_creation_allocates_runtime_but_defers_pools(create_state: Callable[..., BoundedAttentionState]): + config = ConvSSMStateConfig( + logical_layer_ids=(0, 2), + global_k_heads=2, + global_v_heads=2, + key_dim=8, + value_dim=8, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + state = create_state( + config=config, + request_capacity=3, + mtp_step=2, + device="cpu", + ) + assert state.config is config and state.state_name is None + assert state.conv_state.shape == (2, 3, 48, 5) + assert state.ssm_state.shape == (2, 9, 2, 8, 8) + assert state.mtp_state_index.shape == (3,) + assert state.big_pages_pool is None and state.small_pages_pool is None + + +def test_kpool_creation_preserves_pool_and_mtp_dimensions(): + config = WindowCompressorStateConfig( + logical_layer_ids=(3, 7), + window_size=4, + compress_stride=4, + target="indexer", + candidate_dim=128, + dtype=torch.bfloat16, + ) + state = create_window_compressor_state( + config=config, + request_capacity=3, + mtp_step=2, + device="cpu", + ) + assert state.config is config and state.state_name is None + assert state.buffer.shape == (2, 3, 6, 256) + assert state.get_state_type() == "compressor_indexer_w4_s4_bf16" + + +def test_create_rejects_encoding_from_another_logical_family(): + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=576, + dtype=torch.bfloat16, + ) + with pytest.raises(AssertionError, match="unsupported GQA encoding"): + create_gqa_state(config=config, kv_type="fp8kv_dsa", device="meta") + packed_config = PackedMLALogicalConfig( + logical_layer_groups=((0,),), + latent_dim=576, + format="packed_group128_fp8", + ) + with pytest.raises(AssertionError, match="normal MLA creation"): + create_mla_state(config=packed_config, device="meta") + with pytest.raises(AssertionError, match="packed MLA creation"): + create_packed_mla_state( + config=MLAConfig(logical_layer_groups=((0,),), latent_dim=576, latent_dtype=torch.bfloat16), + device="meta", + ) + + +@pytest.mark.parametrize( + "state_class", + [ + GQAFullAttentionState, + GroupQuantizedFullAttentionState, + StaticFP8FullAttentionState, + MLAFullAttentionState, + PackedMLAFullAttentionState, + ], +) +def test_full_attention_implementations_are_independent(state_class: type[FullAttentionState]): + # 具体方案只共享抽象契约，不继承其他具体方案的存储或操作实现。 + assert state_class.__bases__ == (FullAttentionState,) + assert not state_class.__abstractmethods__ + assert FullAttentionState.__abstractmethods__ <= state_class.__dict__.keys() diff --git a/unit_tests/common/att_resources/test_state_identity.py b/unit_tests/common/att_resources/test_state_identity.py new file mode 100644 index 0000000000..8fffc04927 --- /dev/null +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -0,0 +1,356 @@ +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups + +from typing import TYPE_CHECKING +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.full_att import StaticFP8FullAttentionState +from unit_tests.common.att_resources.model_helpers import create_full_attention_roles +from unit_tests.common.att_resources.test_capacity import definition + +if TYPE_CHECKING: + from easydict import EasyDict + + +KINDS = ["gqa", "int4", "int8", "static_head", "static_tensor", "mla", "packed", "packed_fp8", "packed_glm"] + + +@pytest.mark.parametrize("kind", [*KINDS, "compressed", "kda", "gdn", "swa", "compressor"]) +def test_pd_region_cache_returns_independent_snapshots(kind: str, pd_args: EasyDict) -> None: + from lightllm.common.attention_state.bounded_att import GDNState, KDAState + from unit_tests.common.att_resources.test_bounded_state import make_state + from unit_tests.common.att_resources.test_dsv4_state import make_atomic_state, make_full + + pd_args.pd_kv_page_size = 256 + if kind in KINDS: + state = definition(kind) + elif kind == "compressed": + state = make_full(ratio=4) + elif kind in ("kda", "gdn"): + state = make_state(cls=KDAState if kind == "kda" else GDNState, device="meta") + else: + state = make_atomic_state("swa" if kind == "swa" else "kv_w8") + + getter = state.get_pd_transfer_regions + before = getter.cache_info() + first, second = getter(), getter() + after = getter.cache_info() + assert after.misses == before.misses + 1 and after.hits == before.hits + 1 + assert first == second and first is not second + for first_region, second_region in zip(first, second): + assert first_region is not second_region + for first_mapping, second_mapping in zip(first_region.axis_mappings, second_region.axis_mappings): + assert first_mapping is not second_mapping + + +@pytest.mark.parametrize("kind", KINDS) +def test_full_state_prepares_layout_before_capacity_allocation(kind: str, pd_args: EasyDict): + from inspect import signature + + state = definition(kind) + parameters = signature(type(state)).parameters + assert not {"token_capacity", "buffer", "scale_buffer"} & parameters.keys() + assert state.capacity == 0 and read_kv_storage(state) is None + token_bytes = state.get_token_bytes() + cpu_regions = state.get_cpu_transfer_regions() + pd_regions = state.get_pd_transfer_regions() + assert read_kv_storage(state) is None + original_config = state.config + state.device = torch.device("cpu") + state.allocate_storage(8) + assert state.config is original_config + assert state.capacity == 8 and read_kv_storage(state).shape[1] == 8 + state.page_size + assert not read_kv_storage(state).is_meta + assert state.get_token_bytes() == token_bytes + current_cpu_regions = state.get_cpu_transfer_regions() + assert current_cpu_regions == cpu_regions + assert current_cpu_regions is not cpu_regions + assert current_cpu_regions[0] is not cpu_regions[0] + assert state.get_pd_transfer_regions() == pd_regions + actual_bytes = sum(buffer.numel() * buffer.element_size() for buffer in full_state_test_buffers(state).values()) + assert actual_bytes == token_bytes * (8 + state.page_size) + + +def test_numbered_drafts_keep_state_name_order_and_unique_members(): + manager = AttentionStateManager(request_capacity=1) + for role in ("draft35", "draft10", "draft1", "main", "draft0", "draft2"): + manager.register_state(role, definition("gqa")) + with pytest.raises(AssertionError, match="already bound"): + manager.register_state("draft1", definition("gqa")) + manager.finish_state_registration() + assert [state.state_name.split("_", 1)[0] for state in manager.get_full_attention_states()] == [ + "draft0", + "draft10", + "draft1", + "draft2", + "draft35", + "main", + ] + assert manager.get_full_attention_states("draft1") == (manager.draft1_gqa_bf16_state,) + assert manager.get_full_attention_states("draft10") == (manager.draft10_gqa_bf16_state,) + + +def test_custom_state_type_controls_registered_resource_name(): + from unit_tests.common.att_resources.test_full_state import full_state + from lightllm.common.attention_state.full_att import GQAFullAttentionState + + class CustomGQAState(GQAFullAttentionState): + def get_state_type(self) -> str: + return "gqa_custom_fp32" + + state = full_state(CustomGQAState) + manager = AttentionStateManager(request_capacity=1) + manager.register_state("draft0", state) + manager.finish_state_registration() + assert manager.draft0_gqa_custom_fp32_state is state + assert state.state_name == "draft0_gqa_custom_fp32_state" + assert manager.get_full_attention_states("draft0") == (state,) + assert not hasattr(state, "state_type") + + +@pytest.mark.parametrize("kind", ["full_meta", "full", "bounded", "tail"]) +@pytest.mark.parametrize("model_name", ["main", "draft1"]) +def test_register_state_preserves_object_and_storage(kind: str, model_name: str): + from unit_tests.common.att_resources.test_bounded_state import make_state + from unit_tests.common.att_resources.test_window_compressor_state import make_window + + if kind == "full_meta": + state = definition("gqa") + storage = tuple(full_state_test_buffers(state).values()) + elif kind == "full": + state = allocated_state("gqa") + storage = tuple(full_state_test_buffers(state).values()) + elif kind == "bounded": + state = make_state() + storage = state.conv_state + pool = state.big_pages_pool + else: + state = make_window() + storage = state.buffer + config = state.config + manager = AttentionStateManager(request_capacity=3) + manager.register_state(model_name, state) + name = f"{model_name}_{state.get_state_type()}_state" + assert getattr(manager, name) is state + assert state.state_name == name + assert state.config is config and not hasattr(config, "state_name") + if kind.startswith("full"): + assert all(actual is original for actual, original in zip(full_state_test_buffers(state).values(), storage)) + elif kind == "bounded": + assert state.conv_state is storage and state.big_pages_pool is pool + else: + assert state.buffer is storage + other_role = "draft0" if model_name == "main" else "main" + with pytest.raises(AssertionError, match="multiple model roles"): + manager.register_state(other_role, state) + assert state.state_name == name and getattr(manager, name) is state + manager.finish_state_registration() + with pytest.raises(AssertionError, match="registration is already finished"): + manager.register_state(other_role, definition("gqa")) + + +def test_model_resource_role_numbers_the_36th_draft_and_rejects_the_37th(): + from lightllm.models.llama.model import LlamaTpPartModel + + draft = LlamaTpPartModel.__new__(LlamaTpPartModel) + draft.is_mtp_draft_model = True + draft.mtp_previous_draft_models = [object()] * 35 + assert draft.att_state_role == "draft35" + draft.mtp_previous_draft_models.append(object()) + with pytest.raises(AssertionError, match="at most 36"): + _ = draft.att_state_role + + +@pytest.mark.parametrize("role", ["draft", "draft36", "draft-1", "draft01"]) +def test_register_state_rejects_invalid_draft_numbers(role: str): + manager = AttentionStateManager(request_capacity=1) + with pytest.raises(AssertionError, match="unknown model role"): + manager.register_state(role, definition("gqa")) + + +def allocated_state(kind: str, device: str = "cpu"): + template = definition(kind) + return rebuild_full_state( + template=template, + config=template.config, + token_capacity=8, + device=device, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", KINDS) +def test_independent_role_states_copy_and_transfer_distinct_layers(kind: str, cpu_args: EasyDict): + cpu_args.cpu_cache_token_page_size = 4 + physical = allocated_state(kind, device="cuda") + main, draft = create_full_attention_roles(physical, 1) + assert main.config.logical_layer_groups == ((0,),) and draft.config.logical_layer_groups == ((2,),) + for role, state, layer in (("main", main, 0), ("draft0", draft, 1)): + physical_buffers = full_state_test_buffers(physical) + for name, buffer in full_state_test_buffers(state).items(): + original = physical_buffers[name] + assert buffer.untyped_storage().data_ptr() != original.untyped_storage().data_ptr() + buffer.fill_(layer + 3) + if isinstance(state, StaticFP8FullAttentionState): + for buffer, original in ((state.scales, physical.scales), (state.q_scales, physical.q_scales)): + if buffer is not None: + assert buffer.untyped_storage().data_ptr() != original.untyped_storage().data_ptr() + + source = AttentionStateManager(request_capacity=2) + source.register_state("main", main) + source.register_state("draft0", draft) + source.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + source.finish_state_registration() + assert getattr(source, main.state_name) is main + assert getattr(source, draft.state_name) is draft + source.copy_full_att_tokens( + torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda") + ) + + target_physical = allocated_state(kind, device="cuda") + target_main, target_draft = create_full_attention_roles(target_physical, 1) + target = AttentionStateManager(request_capacity=2) + target.register_state("main", target_main) + target.register_state("draft0", target_draft) + target.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + target.finish_state_registration() + page = torch.zeros((1, source.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + source.get_cpu_cache_operator().write_req( + token_indices=torch.tensor([0, 1], dtype=torch.long, device="cuda"), + cpu_page_buffers=page, + cpu_page_indices=[0], + layouts=source.get_cpu_cache_operator().get_page_layouts(), + ) + target.get_cpu_cache_operator().read_req( + cpu_page_buffers=page, + cpu_page_indices=[0], + token_indices=torch.tensor([6, 7], dtype=torch.long, device="cuda"), + layouts=source.get_cpu_cache_operator().get_page_layouts(), + ) + for src, dst in ((main, target_main), (draft, target_draft)): + fields = { + name: field.tensor for name, field in full_state_test_fields(dst, full_state_test_buffers(dst)).items() + } + for name, field in full_state_test_fields(src, full_state_test_buffers(src)).items(): + if "token" in field.axes: + assert torch.equal(field.tensor[:, :2], fields[name][:, 6:8]) + assert set(source.get_cpu_cache_operator().get_page_layouts()) == {main.state_name, draft.state_name} + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_explicit_members_keep_role_and_scheme_distinct_even_with_overlapping_layers(overlap: bool): + main = allocated_state("gqa") + draft = allocated_state("gqa") + manager = AttentionStateManager(request_capacity=2) + manager.register_state("main", main) + manager.register_state("draft0", draft) + manager.full_att_allocator = FullAttKVAllocator(main.capacity, main.page_size, pin_memory=False) + template = allocated_state("int8") + quant = rebuild_full_state(template, config=replace_logical_layer_groups(template.config, ((4,), (6,)))) + manager.register_state("main", quant) + if overlap: + overlapping = allocated_state("static_head") + manager.register_state("main", overlapping) + manager.finish_state_registration() + assert manager.main_gqa_bf16_state is main + assert manager.draft0_gqa_bf16_state is draft + assert manager.main_gqa_int8_group8_bf16_scale_state is quant + assert manager.get_full_attention_states("draft0") == (draft,) + expected = [main, quant] + if overlap: + assert manager.main_gqa_fp8_static_per_head_state is overlapping + expected.append(overlapping) + assert manager.get_full_attention_states("main") == tuple(sorted(expected, key=lambda state: state.state_name)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mixed_quantization_in_one_role_remains_distinct_in_copy_and_cpu_pages(cpu_args: EasyDict): + cpu_args.cpu_cache_token_page_size = 4 + full = allocated_state("gqa", device="cuda") + quant = allocated_state("int8", device="cuda") + assert full.get_state_type() == "gqa_bf16" + assert quant.get_state_type() == "gqa_int8_group8_bf16_scale" + for state in (full, quant): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + field.tensor.fill_(5) + field.tensor[:, 4:6].zero_() + manager = AttentionStateManager(request_capacity=2) + for state in (full, quant): + manager.register_state("main", state) + manager.finish_state_registration() + manager.copy_full_att_tokens( + torch.tensor([0, 1], dtype=torch.long, device="cuda"), torch.tensor([4, 5], dtype=torch.long, device="cuda") + ) + for state in (full, quant): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + assert (field.tensor[:, 4:6] == 5).all() + layouts = manager.get_cpu_cache_operator().get_page_layouts() + page = torch.zeros((1, max(layout.byte_end for layout in layouts.values())), dtype=torch.uint8, pin_memory=True) + manager.get_cpu_cache_operator().write_req( + token_indices=torch.tensor([0, 1], dtype=torch.long, device="cuda"), + cpu_page_buffers=page, + cpu_page_indices=[0], + layouts=layouts, + ) + for state in (full, quant): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + field.tensor[:, 6:8].zero_() + manager.get_cpu_cache_operator().read_req( + cpu_page_buffers=page, + cpu_page_indices=[0], + token_indices=torch.tensor([6, 7], dtype=torch.long, device="cuda"), + layouts=layouts, + ) + for state in (full, quant): + for field in full_state_test_fields(state, full_state_test_buffers(state)).values(): + assert (field.tensor[:, 6:8] == 5).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kind", KINDS) +def test_token_copy_captures_large_index_sets_and_preserves_complete_buffers(kind: str): + template = definition(kind) + state = rebuild_full_state(template, token_capacity=512, device="cuda") + # 覆盖 int32/int64 和非连续索引；128 个元素足以触发旧 isin 实现的同步路径。 + source = torch.arange(256, dtype=torch.int32, device="cuda")[::2] + target = torch.arange( + start=256, + end=512, + dtype=torch.int64, + device="cuda", + )[::2] + buffers = tuple(full_state_test_buffers(state).values()) + for buffer in buffers: + buffer.view(torch.uint8).random_(0, 256) + state.copy_full_att_tokens(source, target) # 先完成 Triton 编译。 + torch.cuda.synchronize() + + expected = [] + for buffer in buffers: + buffer.view(torch.uint8).random_(0, 256) + snapshot = buffer.view(torch.uint8).clone() + snapshot[:, target] = snapshot[:, source] + expected.append(snapshot) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + state.copy_full_att_tokens(source, target) + for buffer in buffers: + buffer.index_fill_(1, target, 0) + graph.replay() + torch.cuda.synchronize() + for buffer, snapshot in zip(buffers, expected): + # 逐字节核对整块存储，包含量化 scale、packed padding 和未参与搬运的槽位。 + assert torch.equal(buffer.view(torch.uint8), snapshot) diff --git a/unit_tests/common/att_resources/test_state_queries.py b/unit_tests/common/att_resources/test_state_queries.py new file mode 100644 index 0000000000..59d5751f69 --- /dev/null +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -0,0 +1,189 @@ +"""State 对象创建完成后即可查询，容量分配保持对象身份和查询缓存不变。""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import read_kv_storage + + +import gc +import pickle +import weakref +from functools import partial + +import pytest +from easydict import EasyDict + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.full_att import FullAttentionState +from lightllm.server.core.objs.start_args_type import StartArgs +from unit_tests.common.att_resources.manager_fixtures import MainDraftTestResourceManager +from unit_tests.common.att_resources.test_capacity import definition + + +def test_created_state_queries_keep_the_same_objects_through_storage_initialization(monkeypatch: pytest.MonkeyPatch): + args = StartArgs(page_size=4) + definitions = [] + manager = AttentionStateManager(request_capacity=1) + for role, kind in (("main", "gqa"), ("draft0", "int8")): + manager.register_state(role, definition(kind)) + definitions.append( + next( + state + for state in vars(manager).values() + if isinstance(state, FullAttentionState) and state.state_name.startswith(role + "_") + ) + ) + for query in ( + manager.get_full_attention_states, + manager.get_bounded_attention_states, + ): + with pytest.raises(AssertionError, match="State registration"): + query() + assert manager._get_states.cache_info().currsize == 0 + + cached = [] + allocate_storage = AttentionStateManager.allocate_storage + + def check_created_before_allocation(manager: AttentionStateManager, *args: object, **kwargs: object): + states = manager.get_full_attention_states() + assert all(read_kv_storage(state) is None for state in states) + assert manager.full_att_allocator is None + cached.append(states) + allocate_storage(manager, *args, **kwargs) + + monkeypatch.setattr(AttentionStateManager, "allocate_storage", check_created_before_allocation) + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.finish_state_registration() + manager.initialize_storage( + args=args, + token_capacity=8, + mem_fraction=1.0, + device="cpu", + ) + assert manager.get_full_attention_states() is cached[0] + for role, meta in zip(("main", "draft0"), definitions): + (state,) = manager.get_full_attention_states(role) + assert state is meta and read_kv_storage(state).device.type == "cpu" + assert manager.get_full_attention_states(role) == (state,) + + +def test_cached_role_queries_reuse_instance_references_and_keep_stable_order(): + manager = MainDraftTestResourceManager() + queries = (manager.get_full_attention_states, manager.get_bounded_attention_states) + results = {(query.__name__, role): query(role) for query in queries for role in (None, "main", "draft0")} + assert manager._get_states.cache_info().currsize == 6 + assert manager.get_full_attention_states() == ( + manager.draft0_mla_fp32_state, + manager.main_gqa_fp32_state, + ) + assert manager.get_bounded_attention_states("main") == (manager.main_kda_conv_fp32_ssm_fp32_state,) + assert manager.get_bounded_attention_states("draft0") == (manager.draft0_gdn_conv_fp32_ssm_fp32_state,) + + hits = manager._get_states.cache_info().hits + for query in queries: + for role in (None, "main", "draft0"): + assert query(role) is results[(query.__name__, role)] + with pytest.raises(AssertionError, match="unknown model role"): + query("target") + assert manager._get_states.cache_info().hits == hits + 6 + manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.fill_(7) + assert (manager.get_bounded_attention_states("main")[0].conv_state == 7).all() + + +def test_finished_state_creation_allows_unallocated_storage(): + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", definition("gqa")) + manager.full_att_allocator = FullAttKVAllocator(0, 4, pin_memory=False) + manager.finish_state_registration() + assert read_kv_storage(manager.get_full_attention_states()[0]) is None + + +def test_ipc_copy_preserves_owner_storage_and_shares_runtime_tensors(): + source = MainDraftTestResourceManager() + source_states = source.get_bounded_attention_states() + owner_pools = [(state.big_pages_pool, state.small_pages_pool) for state in source_states] + + exported = source.get_pd_operator().create_ipc_copy() + + assert type(exported) is type(source) and exported is not source + for name in ("full_att_allocator", "bounded_att_big_page_allocator", "bounded_att_small_page_allocator"): + assert getattr(source, name) is not None + assert getattr(exported, name) is None + for state, (big_pool, small_pool) in zip(source_states, owner_pools): + copied = getattr(exported, state.state_name) + assert copied is not state + assert copied.conv_state is state.conv_state + assert copied.ssm_state is state.ssm_state + assert copied.mtp_state_index is state.mtp_state_index + assert copied.big_pages_pool is None and copied.small_pages_pool is None + assert copied.request_capacity == state.request_capacity + assert copied.mtp_step == state.mtp_step + assert state.big_pages_pool is big_pool and state.small_pages_pool is small_pool + assert big_pool.size == source.bounded_att_big_page_allocator.size + assert small_pool.size == source.bounded_att_small_page_allocator.size + assert exported.get_bounded_attention_states() == tuple( + getattr(exported, state.state_name) for state in source_states + ) + for state in source.get_full_attention_states(): + assert getattr(exported, state.state_name) is state + + +@pytest.mark.parametrize("share_pools", [False, True]) +def test_pd_and_dp_export_have_distinct_pool_ownership(cpu_args: EasyDict, share_pools: bool) -> None: + cpu_args.enable_dp_prompt_cache_fetch = share_pools + source = MainDraftTestResourceManager() + pd_copy = source.get_pd_operator().create_ipc_copy() + for state in source.get_bounded_attention_states(): + exported = getattr(pd_copy, state.state_name) + assert exported.big_pages_pool is None and exported.small_pages_pool is None + assert state.big_pages_pool is not None and state.small_pages_pool is not None + if not share_pools: + with pytest.raises(AssertionError, match="shared CPU page pools"): + source.get_dp_shared_operator().create_ipc_copy() + return + dp_copy = source.get_dp_shared_operator().create_ipc_copy() + for name in ("full_att_allocator", "bounded_att_big_page_allocator", "bounded_att_small_page_allocator"): + assert getattr(dp_copy, name) is None + assert getattr(source, name) is not None + for state in source.get_bounded_attention_states(): + exported = getattr(dp_copy, state.state_name) + assert exported.big_pages_pool is state.big_pages_pool + assert exported.small_pages_pool is state.small_pages_pool + assert exported.conv_state is state.conv_state + + +def test_ipc_rebuilds_queries_from_imported_members_without_allocators(): + source = MainDraftTestResourceManager() + source.get_full_attention_states() + source.get_bounded_attention_states() + imported = pickle.loads(pickle.dumps(source.get_pd_operator().create_ipc_copy())) + assert imported._state_registration_finished + assert imported.full_att_allocator is None + assert imported.bounded_att_big_page_allocator is None + assert imported.bounded_att_small_page_allocator is None + misses = source._get_states.cache_info().misses + for query in ( + imported.get_full_attention_states, + imported.get_bounded_attention_states, + ): + for state in query(): + name = state.state_name + assert state is getattr(imported, name) and state is not getattr(source, name) + assert imported._get_states.cache_info().misses == misses + 2 + + +def test_clearing_query_cache_releases_manager_and_states(): + manager = MainDraftTestResourceManager() + manager.get_full_attention_states() + manager.get_bounded_attention_states() + manager_ref = weakref.ref(manager) + state_ref = weakref.ref(manager.main_kda_conv_fp32_ssm_fp32_state) + del manager + gc.collect() + assert manager_ref() is not None and state_ref() is not None + AttentionStateManager._get_states.cache_clear() + gc.collect() + assert manager_ref() is None and state_ref() is None diff --git a/unit_tests/common/att_resources/test_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py new file mode 100644 index 0000000000..1fd660fa8d --- /dev/null +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -0,0 +1,223 @@ +"""State 配置独立保留全局布局，改变 TP 后仍可恢复同一份缓存。""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups + +from typing import TYPE_CHECKING +import pickle +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.attention_state.full_att import ( + GQAAttentionConfig, + MLAConfig, + GQAFullAttentionState, + MLAFullAttentionState, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + from easydict import EasyDict + from lightllm.common.attention_state.bounded_att.base_state import BoundedAttentionState + from lightllm.common.attention_state.full_att.base_state import FullAttentionState + + +def cpu_layout(state: FullAttentionState | BoundedAttentionState): + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + return manager.get_cpu_cache_operator().get_page_layouts()[state.state_name] + + +@pytest.mark.parametrize( + "state_cls,heads,replicated,source_ranges,target_ranges", + [ + (GQAFullAttentionState, 4, False, [(0, 2), (2, 4)], [(i, i + 1) for i in range(4) for _ in range(2)]), + (GQAFullAttentionState, 2, False, [(0, 1)] * 4 + [(1, 2)] * 4, [(0, 2)]), + (GQAFullAttentionState, 2, True, [(0, 2)] * 4, [(0, 2)] * 2), + (MLAFullAttentionState, 1, True, [(0, 1)] * 4, [(0, 1)]), + ], +) +def test_serialized_full_config_restores_cpu_pages_with_different_tp( + state_cls: type[FullAttentionState], + heads: int, + replicated: bool, + source_ranges: Sequence[tuple[int, int]], + target_ranges: Sequence[tuple[int, int]], + cpu_args: EasyDict, +): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + cpu_args.cpu_cache_token_page_size = 4 + if state_cls is MLAFullAttentionState: + config = MLAConfig( + logical_layer_groups=((0,),), + latent_dim=4, + latent_dtype=torch.float32, + ) + else: + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=heads, + head_dim=4, + dtype=torch.float32, + replicated=replicated, + ) + reference = state_cls(config, page_size=4, device="cuda") + reference.allocate_storage(4) + for storage_buffer in full_state_test_buffers(reference).values(): + storage_buffer.copy_(torch.arange(storage_buffer.numel()).reshape(storage_buffer.shape)) + fields = { + name: field.tensor + for name, field in full_state_test_fields(reference, full_state_test_buffers(reference)).items() + } + reference_layout = cpu_layout(reference) + buffer = torch.zeros((1, reference_layout.byte_end), dtype=torch.uint8, pin_memory=True) + + # 接收方只需要配置和缓存页；config 中不保存本地 head 数或 manager 映射。 + restored_config = pickle.loads(pickle.dumps(config)) + assert "head_num" not in vars(restored_config) + for rank, (start, end) in enumerate(source_ranges): + source = state_cls(replace(config, tp_rank=rank, tp_world_size=len(source_ranges)), page_size=4, device="cuda") + source.allocate_storage(4) + layout = cpu_layout(source) + assert layout.byte_end == reference_layout.byte_end + for name, field in full_state_test_fields(source, full_state_test_buffers(source)).items(): + field.tensor.copy_(fields[name][:, :, start:end]) + source.write_to_cpu( + token_indices=torch.arange(4, dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[0], + layout=layout, + ) + + for rank, (start, end) in enumerate(target_ranges): + target = state_cls( + replace(restored_config, tp_rank=rank, tp_world_size=len(target_ranges)), page_size=4, device="cuda" + ) + target.allocate_storage(4) + layout = cpu_layout(target) + assert layout.byte_end == reference_layout.byte_end + target.read_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[0], + token_indices=torch.arange(4, dtype=torch.long, device="cuda"), + layout=layout, + ) + for name, field in full_state_test_fields(target, full_state_test_buffers(target)).items(): + torch.testing.assert_close(field.tensor[:, :4], fields[name][:, :4, start:end]) + + +def test_bounded_config_keeps_distinct_global_heads_when_local_heads_are_replicated(cpu_args: EasyDict): + cpu_args.cpu_cache_token_page_size = 4 + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=2, + global_v_heads=4, + key_dim=2, + value_dim=3, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + tp_rank=7, + tp_world_size=8, + ) + source = GDNState(config, request_capacity=1, device="meta") + assert config.num_k_heads == config.num_v_heads == 1 + restored_config = replace(pickle.loads(pickle.dumps(config)), tp_rank=0, tp_world_size=1) + target = GDNState(restored_config, request_capacity=1, device="meta") + assert restored_config.num_k_heads == 2 and restored_config.num_v_heads == 4 + assert "num_k_heads" not in vars(config) and "num_v_heads" not in vars(config) + assert target.conv_state.shape == (1, 1, 20, 2) + assert target.ssm_state.shape == (1, 1, 4, 2, 3) + source_layout = cpu_layout(source) + target_layout = cpu_layout(target) + assert source_layout.byte_end == target_layout.byte_end + assert [(r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in source_layout.regions] == [ + (r.field, r.byte_start, r.byte_end, r.dtype, r.axes, r.shape) for r in target_layout.regions + ] + + +@pytest.mark.parametrize("bounded", [False, True]) +@pytest.mark.parametrize("page_tokens", [4, 8]) +def test_drafts_keep_independent_global_layouts_with_equal_local_shapes( + bounded: bool, + page_tokens: int, + pd_args: EasyDict, +): + pd_args.pd_kv_page_size = page_tokens + if bounded: + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=2, + global_v_heads=4, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + tp_world_size=8, + ) + source = GDNState(config, request_capacity=1, device="meta") + other_config = replace(config, logical_layer_ids=(1,), global_k_heads=4) + else: + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=2, + head_dim=4, + dtype=torch.float32, + tp_world_size=8, + ) + source = GQAFullAttentionState(config, page_size=4, device="meta") + other_config = replace(config, logical_layer_groups=((1,),), global_head_num=4) + if bounded: + other = type(source)( + config=other_config, + request_capacity=source.request_capacity, + mtp_step=source.mtp_step, + device="meta", + ) + else: + other = rebuild_full_state(source, config=other_config) + manager = AttentionStateManager(request_capacity=1) + manager.register_state("draft0", source) + manager.register_state("draft1", other) + manager.finish_state_registration() + query = manager.get_bounded_attention_states if bounded else manager.get_full_attention_states + assert query("draft0") == (source,) + assert query("draft1") == (other,) + assert query() == (source, other) + + page_bytes = manager.get_pd_operator().get_page_bytes() + assert not hasattr(manager, "pd_page_bytes") + assert not hasattr(manager, "pd_transfer_layouts") + if not bounded: + # 两个 draft 共 6 个全局 KV head：K/V、head_dim=4、FP32，无额外页头。 + assert page_bytes == page_tokens * 192 + assert manager.get_pd_operator().get_page_bytes() == page_bytes + getter = ( + manager.get_pd_operator().get_pd_bounded_att_page_layouts + if bounded + else manager.get_pd_operator().get_pd_full_att_page_layouts + ) + layouts = getter() + assert getter() is layouts + assert set(layouts) == {source.state_name, other.state_name} + assert all(layout.tokens_per_page == (None if bounded else page_tokens) for layout in layouts.values()) + unused_getter = ( + manager.get_pd_operator().get_pd_full_att_page_layouts + if bounded + else manager.get_pd_operator().get_pd_bounded_att_page_layouts + ) + assert unused_getter() == {} + assert manager.get_pd_operator().get_page_bytes() == page_bytes diff --git a/unit_tests/common/att_resources/test_transfer_views.py b/unit_tests/common/att_resources/test_transfer_views.py new file mode 100644 index 0000000000..e3fa612818 --- /dev/null +++ b/unit_tests/common/att_resources/test_transfer_views.py @@ -0,0 +1,233 @@ +"""传输视图整理必须保留存储别名、位模式、未参与区域和独立寻址轴。""" + +import pytest +import torch + +from lightllm.common.attention_state.utils import merge_contiguous_tail_dims, view_max_copy_dtype + + +@pytest.mark.parametrize( + "layout, expected_shape, start_dim", + [ + ("dense", (48,), 0), + ("source_gap", (2, 3, 8), 0), + ("destination_gap", (2, 3, 8), 0), + ("strided_tail", (48,), 0), + ("singleton_gap", (2, 8), 0), + ("protected_window", (2, 5, 24), 2), + ], +) +def test_merge_tail_preserves_copy_regions( + layout: str, + expected_shape: tuple[int, ...], + start_dim: int, +) -> None: + source_shape = (2, 5, 3, 8) if layout == "protected_window" else (2, 3, 8) + destination_shape = source_shape + if layout == "source_gap": + source_shape = (2, 3, 10) + elif layout == "destination_gap": + destination_shape = (2, 3, 10) + elif layout == "strided_tail": + source_shape = (2, 3, 16) + elif layout == "singleton_gap": + source_shape, destination_shape = (2, 1, 10), (2, 1, 12) + + source_storage = torch.arange(torch.Size(source_shape).numel(), dtype=torch.float32).view(source_shape) + destination_storage = torch.full(size=destination_shape, fill_value=-1, dtype=torch.float32) + source = source_storage[..., ::2] if layout == "strided_tail" else source_storage[..., :8] + destination = destination_storage[..., :8] + # 直接按原始视图执行参考复制，保留目标 padding，随后比较整个底层存储。 + expected = destination_storage.clone() + expected.as_strided( + size=destination.shape, + stride=destination.stride(), + storage_offset=destination.storage_offset(), + ).copy_(source) + + copies = [(source, destination)] + merged = merge_contiguous_tail_dims(copies=copies, start_dim=start_dim) + source_view, destination_view = merged[0] + assert source_view.shape == destination_view.shape == expected_shape + assert source_view.data_ptr() == source.data_ptr() + assert destination_view.data_ptr() == destination.data_ptr() + if expected_shape == source.shape: + assert merged is copies + destination_view.copy_(source_view) + assert torch.equal(destination_storage, expected) + + +def test_merge_request_slices_preserves_batch_offsets() -> None: + source = torch.arange(2 * 4 * 3 * 8, dtype=torch.float32).view(2, 4, 3, 8) + destination = torch.full(size=(4, 2, 3, 8), fill_value=-1, dtype=torch.float32) + expected = destination.clone() + copies = [] + for req_idx in (0, 2): + copies.append((source[:, req_idx], destination[req_idx])) + expected[req_idx].copy_(source[:, req_idx]) + + for source_view, destination_view in merge_contiguous_tail_dims(copies=copies): + assert source_view.shape == destination_view.shape == (2, 24) + destination_view.copy_(source_view) + assert torch.equal(destination, expected) + + +@pytest.mark.parametrize("destination_prefix", [(2, 11), (3, 2, 5)]) +@pytest.mark.parametrize("keepdim", [False, True]) +def test_merge_tail_preserves_different_prefix_layouts(destination_prefix: tuple[int, ...], keepdim: bool) -> None: + source = torch.arange(2 * 7 * 4 * 2, dtype=torch.float32).view(2, 7, 4, 2) + destination_shape = destination_prefix + (4, 2) + destination = torch.arange(torch.Size(destination_shape).numel(), dtype=torch.float32).view(destination_shape) + source_view, destination_view = merge_contiguous_tail_dims( + copies=[(source, destination)], + start_dim=-2, + keepdim=keepdim, + )[0] + tail_shape = (1, 8) if keepdim else (8,) + assert source_view.shape == (2, 7) + tail_shape + assert destination_view.shape == destination_prefix + tail_shape + assert source_view.data_ptr() == source.data_ptr() + assert destination_view.data_ptr() == destination.data_ptr() + assert torch.equal(source_view.flatten(), source.flatten()) + assert torch.equal(destination_view.flatten(), destination.flatten()) + + +@pytest.mark.parametrize("gap_side", ["source", "destination"]) +def test_merge_tail_keeps_head_padding(gap_side: str) -> None: + source = torch.empty(size=(2, 7, 4, 3 if gap_side == "source" else 2))[..., :2] + destination = torch.empty(size=(3, 2, 5, 4, 3 if gap_side == "destination" else 2))[..., :2] + copies = [(source, destination)] + assert merge_contiguous_tail_dims(copies=copies, start_dim=-2, keepdim=True) is copies + + +@pytest.mark.parametrize( + "shape, source_strides, destination_strides, flat_shape, kept_shape", + [ + ((2, 1, 8), (10, 255, 1), (12, 7, 1), (2, 8), (2, 1, 8)), + ((2, 3, 1), (30, 2, 71), (11, 1, 59), (2, 3), (2, 1, 3)), + ((2, 1, 3, 1), (30, 500, 2, 99), (16, 1, 1, 32), (2, 3), (2, 1, 1, 3)), + ((2, 1, 1), (7, 100, 300), (11, 90, 20), (2,), (1, 1, 2)), + ], +) +@pytest.mark.parametrize("keepdim", [False, True]) +def test_merge_tail_singleton_axes_preserve_addressing( + shape: tuple[int, ...], + source_strides: tuple[int, ...], + destination_strides: tuple[int, ...], + flat_shape: tuple[int, ...], + kept_shape: tuple[int, ...], + keepdim: bool, +) -> None: + source_size = 1 + sum((size - 1) * stride for size, stride in zip(shape, source_strides)) + destination_size = 1 + sum((size - 1) * stride for size, stride in zip(shape, destination_strides)) + source = torch.arange(source_size).as_strided(size=shape, stride=source_strides) + destination = torch.arange(destination_size).as_strided(size=shape, stride=destination_strides) + source_view, destination_view = merge_contiguous_tail_dims( + copies=[(source, destination)], + keepdim=keepdim, + )[0] + assert source_view.shape == destination_view.shape == (kept_shape if keepdim else flat_shape) + assert source_view.data_ptr() == source.data_ptr() + assert destination_view.data_ptr() == destination.data_ptr() + assert torch.equal(source_view.flatten(), source.flatten()) + assert torch.equal(destination_view.flatten(), destination.flatten()) + + +@pytest.mark.parametrize( + "dtype, width, src_offset, dst_offset, expected_dtype", + [ + (torch.bfloat16, 8, 0, 0, torch.int64), + (torch.bfloat16, 6, 0, 0, torch.int32), + (torch.int8, 2, 0, 0, torch.int16), + (torch.float32, 3, 0, 0, torch.float32), + (torch.int64, 8, 0, 0, torch.int64), + (torch.bfloat16, 8, 1, 0, torch.bfloat16), + (torch.bfloat16, 8, 0, 1, torch.bfloat16), + (torch.bfloat16, 8, 2, 0, torch.int32), + ], +) +def test_max_copy_dtype_preserves_data_and_aliases( + dtype: torch.dtype, + width: int, + src_offset: int, + dst_offset: int, + expected_dtype: torch.dtype, +) -> None: + source = torch.arange(32).to(dtype=dtype) + destination = torch.full(size=(32,), fill_value=-1, dtype=dtype) + expected = destination.clone() + source_view = source[src_offset : src_offset + width] + destination_view = destination[dst_offset : dst_offset + width] + expected[dst_offset : dst_offset + width].copy_(source_view) + + copies = [(source_view, destination_view)] + converted = view_max_copy_dtype(copies=copies) + copy_source, copy_destination = converted[0] + assert copy_source.dtype == copy_destination.dtype == expected_dtype + assert copy_source.data_ptr() == source_view.data_ptr() + assert copy_destination.data_ptr() == destination_view.data_ptr() + assert copy_source.ndim == source_view.ndim + if expected_dtype == dtype: + assert converted is copies + copy_destination.copy_(copy_source) + assert torch.equal(destination, expected) + + +@pytest.mark.parametrize( + "src_padding, dst_padding, expected_dtype", + [(1, 0, torch.bfloat16), (0, 2, torch.int32), (0, 0, torch.int64)], +) +def test_max_copy_dtype_preserves_row_padding( + src_padding: int, + dst_padding: int, + expected_dtype: torch.dtype, +) -> None: + source = torch.arange(2 * (8 + src_padding)).to(dtype=torch.bfloat16).view(2, 8 + src_padding) + destination = torch.full(size=(2, 8 + dst_padding), fill_value=-1, dtype=torch.bfloat16) + expected = destination.clone() + expected[:, :8].copy_(source[:, :8]) + + source_view, destination_view = view_max_copy_dtype(copies=[(source[:, :8], destination[:, :8])])[0] + assert source_view.dtype == destination_view.dtype == expected_dtype + assert source_view.shape[0] == destination_view.shape[0] == 2 + destination_view.copy_(source_view) + assert torch.equal(destination, expected) + + +def test_max_copy_dtype_falls_back_for_middle_batch_offset() -> None: + source = torch.arange(64).to(dtype=torch.bfloat16) + destination = torch.full(size=(64,), fill_value=-1, dtype=torch.bfloat16) + copies = [(source[start : start + 8], destination[end : end + 8]) for start, end in ((0, 0), (9, 16), (24, 32))] + assert view_max_copy_dtype(copies=copies) is copies + + +def test_max_copy_dtype_keeps_strided_last_axis() -> None: + source = torch.arange(32, dtype=torch.float32)[::2] + destination = torch.empty(size=(16,), dtype=torch.float32) + copies = [(source, destination)] + assert view_max_copy_dtype(copies=copies) is copies + + +def test_deepcopy_result_protects_cached_values_and_keeps_cache_controls() -> None: + from functools import lru_cache + from lightllm.common.attention_state.utils import deepcopy_result + + calls = 0 + + @deepcopy_result + @lru_cache(maxsize=None) + def describe() -> dict[str, list[int]]: + nonlocal calls + calls += 1 + return {"shape": [2, 128]} + + first = describe() + first["shape"].append(4) + second = describe() + assert second == {"shape": [2, 128]} + assert first is not second and first["shape"] is not second["shape"] + assert calls == 1 and describe.cache_info().hits == 1 + assert describe.cache_parameters() == {"maxsize": None, "typed": False} + describe.cache_clear() + assert describe.cache_info().currsize == 0 + assert describe() == second and calls == 2 diff --git a/unit_tests/common/att_resources/test_window_compressor_state.py b/unit_tests/common/att_resources/test_window_compressor_state.py new file mode 100644 index 0000000000..106743b126 --- /dev/null +++ b/unit_tests/common/att_resources/test_window_compressor_state.py @@ -0,0 +1,195 @@ +"""统一窗口状态：K-pool 的组边界及不同 MTP 容量间的 PD 恢复。""" + +from dataclasses import replace +from functools import partial +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.server.core.objs.start_args_type import StartArgs + + +def make_window( + device: str = "cpu", + mtp_step: int = 2, + dtype: torch.dtype = torch.float32, +) -> WindowCompressorState: + return WindowCompressorState( + config=WindowCompressorStateConfig( + logical_layer_ids=(3, 7), + window_size=4, + compress_stride=4, + candidate_dim=2, + target="indexer", + dtype=dtype, + ), + request_capacity=3, + mtp_step=mtp_step, + device=device, + ) + + +def window_layout(state: WindowCompressorState) -> StateTransferLayout: + regions = tuple( + replace(region, byte_start=region.byte_start + 16, byte_end=region.byte_end + 16) + for region in state.get_pd_transfer_regions() + ) + return StateTransferLayout( + state_name=state.state_name, + tokens_per_page=None, + byte_start=16, + byte_end=regions[-1].byte_end, + regions=regions, + ) + + +def test_nonoverlapping_window_needs_no_cpu_snapshot() -> None: + state = make_window() + assert state.buffer.shape == (2, 3, 6, 4) + state.buffer.fill_(17) + state.reset_bounded_att_state([0, 1]) + assert not state.buffer[:, :2, :, :2].any() + assert torch.isneginf(state.buffer[:, :2, :, 2:]).all() + assert (state.buffer[:, 2] == 17).all() + state.initialize_page_pools(big_page_capacity=5, small_page_capacity=2) + assert state.big_pages_pool is state.small_pages_pool is None + assert state.get_cpu_transfer_regions() == () + # 组边界不依赖此前窗口；下一组四个输入均由新数据覆盖，无需读取快照。 + state.restore_bounded_att_runtime_from_pages( + req_indices=[2], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[0], + ) + with pytest.raises(AssertionError, match="stride"): + state.restore_bounded_att_runtime_from_pages( + req_indices=[2], + state_token_lens=[7], + page_kind="big", + page_slot_ids=[0], + ) + + +@pytest.mark.parametrize("window_size,compress_stride", [(4, 4), (8, 4), (128, 128)]) +@pytest.mark.parametrize("page_kind", ["small", "cpu"]) +def test_window_cache_boundaries_must_close_compressor_groups( + window_size: int, + compress_stride: int, + page_kind: str, + cpu_args: StartArgs, +) -> None: + cpu_args.linear_att_hash_page_size = compress_stride + cpu_args.cpu_cache_token_page_size = 2 * compress_stride + if page_kind == "small": + cpu_args.linear_att_hash_page_size += 1 + else: + cpu_args.cpu_cache_token_page_size += 1 + state = WindowCompressorState( + config=replace(make_window().config, window_size=window_size, compress_stride=compress_stride), + request_capacity=3, + device="meta", + ) + # 即使不占 CPU 页，也不能在未完成的压缩组内恢复，否则缺少该组历史。 + with pytest.raises(AssertionError, match="pages must align to the compressor stride"): + state.get_cpu_transfer_regions() + + +@pytest.mark.parametrize("window_size,compress_stride", [(4, 4), (8, 4), (128, 128)]) +def test_manager_allocates_page_tensors_only_for_snapshot_states( + window_size: int, + compress_stride: int, + cpu_args: StartArgs, + monkeypatch: pytest.MonkeyPatch, +) -> None: + cpu_args.linear_att_hash_page_size = 128 + cpu_args.linear_att_page_block_num = 2 + cpu_args.cpu_cache_token_page_size = 256 + state = WindowCompressorState( + config=replace(make_window().config, window_size=window_size, compress_stride=compress_stride), + request_capacity=3, + device="cpu", + ) + initialize_pools = Mock(wraps=state.initialize_page_pools) + monkeypatch.setattr(state, "initialize_page_pools", initialize_pools) + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", + partial(FullAttKVAllocator, pin_memory=False), + ) + manager = AttentionStateManager(request_capacity=3) + manager.register_state( + "main", + GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=((3,),), + global_head_num=1, + head_dim=4, + dtype=torch.bfloat16, + ), + device="meta", + ), + ) + manager.register_state("main", state) + manager.finish_state_registration() + manager.allocate_storage( + token_capacity=512, + page_size=4, + big_page_tokens=256, + small_page_capacity=3, + device="cpu", + ) + if window_size == compress_stride: + initialize_pools.assert_not_called() + assert state.big_pages_pool is state.small_pages_pool is None + assert state.state_name not in manager.get_cpu_cache_operator().get_page_layouts() + else: + initialize_pools.assert_called_once() + assert state.big_pages_pool.buffer.shape == (4, 2, window_size - compress_stride, 4) + assert state.small_pages_pool.buffer.shape == (3, 2, window_size - compress_stride, 4) + assert state.state_name in manager.get_cpu_cache_operator().get_page_layouts() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("source_step,target_step", [(0, 0), (0, 5), (2, 0), (1, 3)]) +@pytest.mark.parametrize("length", [0, 1, 3, 4, 5, 11, 12]) +def test_pd_window_restores_different_mtp_rings( + dtype: torch.dtype, + source_step: int, + target_step: int, + length: int, +) -> None: + source = make_window(device="cuda", mtp_step=source_step, dtype=dtype) + target = make_window(device="cuda", mtp_step=target_step, dtype=dtype) + source.buffer.copy_(torch.arange(source.buffer.numel(), device="cuda").reshape_as(source.buffer)) + target.buffer.fill_(37) + layout = window_layout(source) + assert layout == window_layout(target) + pages = torch.full((3, layout.byte_end + 16), 253, dtype=torch.uint8, device="cuda") + source.write_to_pd_page( + req_idx=1, + pd_page_buffers=pages, + pd_page_index=1, + state_token_len=length, + layout=layout, + ) + target.read_from_pd_page( + pd_page_buffers=pages, + req_idx=2, + pd_page_index=1, + state_token_len=length, + layout=layout, + ) + expected = torch.full_like(target.buffer[:, 2], 37) + for position in range(length - 4, length): + slot = position % (4 + target_step) + expected[:, slot] = 0 if position < 0 else source.buffer[:, 1, position % (4 + source_step)] + torch.testing.assert_close(target.buffer[:, 2], expected, rtol=0, atol=0) + assert (target.buffer[:, :2] == 37).all() + assert (pages[0] == 253).all() and (pages[2] == 253).all() + assert (pages[1, :16] == 253).all() and (pages[1, layout.byte_end :] == 253).all() diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index eeb2fb1438..8ec77c6f4d 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -6,13 +6,15 @@ import pytest import torch +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.basemodel.infer_struct import InferStateInfo import lightllm.common.basemodel.attention.linear.gdn as gdn import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.common.state_cache_manager import LinearAttCacheConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.server.api_cli import make_argument_parser import lightllm.utils.backend_validator as backend_validator @@ -62,29 +64,6 @@ def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): assert torch.equal(ssm_states, final_state.to(cache_dtype)) -def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_state_index=None): - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), - ) - infer_state = InferStateInfo() - infer_state.batch_size = b_req_idx.shape[0] - infer_state.b_req_idx = b_req_idx - 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, - ) - state = gdn.LinearAttDecodeAttState( - backend=SimpleNamespace( - model=model, - uses_dynamic_spec_verify_layout=lambda: dynamic_layout, - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - return state, infer_state - - def test_decode_state_initializes_normal_layout(): b_req_idx = torch.tensor([2, 4], dtype=torch.int32) state, infer_state = _create_decode_state(draft_step=0, dynamic_layout=False, b_req_idx=b_req_idx) @@ -97,7 +76,8 @@ def test_decode_state_initializes_normal_layout(): assert state.b_num_accepted_tokens is None -def test_decode_state_initializes_fixed_mtp_layout(): +@pytest.mark.parametrize("role", ["main", "draft0"]) +def test_decode_state_initializes_fixed_mtp_layout(role: str): b_req_idx = torch.tensor([2, 2, 2, 4, 4, 4], dtype=torch.int32) req_to_mtp_state_index = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) state, infer_state = _create_decode_state( @@ -105,6 +85,7 @@ def test_decode_state_initializes_fixed_mtp_layout(): dynamic_layout=False, b_req_idx=b_req_idx, req_to_mtp_state_index=req_to_mtp_state_index, + role=role, ) state.init_state() @@ -118,7 +99,8 @@ def test_decode_state_initializes_fixed_mtp_layout(): ) -def test_decode_state_initializes_dynamic_mtp_layout(monkeypatch): +@pytest.mark.parametrize("role", ["main", "draft0"]) +def test_decode_state_initializes_dynamic_mtp_layout(monkeypatch, role: str): b_req_idx = torch.tensor([2, 2, 4, 4, 4], dtype=torch.int32) req_to_mtp_state_index = torch.tensor([0, 0, 1, 0, 2, 0], dtype=torch.int32) expected_cu_q_seq_len = torch.tensor([0, 2, 5, 5, 5, 5], dtype=torch.int32) @@ -136,6 +118,7 @@ def build_params(**kwargs): dynamic_layout=True, b_req_idx=b_req_idx, req_to_mtp_state_index=req_to_mtp_state_index, + role=role, ) state.init_state() @@ -196,8 +179,9 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear def test_flashqla_validation_loads_linear_cache_config(monkeypatch): linear_config = SimpleNamespace( - num_linear_k_heads=2, - num_linear_v_heads=4, + tp_world_size=2, + global_linear_k_heads=4, + global_linear_v_heads=8, head_linear_k_dim=8, head_linear_v_dim=16, conv_state_dtype=torch.bfloat16, @@ -205,7 +189,7 @@ def test_flashqla_validation_loads_linear_cache_config(monkeypatch): ) load_calls = [] monkeypatch.setattr( - LinearAttCacheConfig, + Qwen3NextResourceConfig, "load_from_args", staticmethod(lambda: load_calls.append(True) or linear_config), ) @@ -359,3 +343,47 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ assert backend_class is FlashQlaLinearAttBackend assert validate_calls == ["flashqla2", "flashqla3"] + + +def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_state_index=None, role="main"): + model = SimpleNamespace( + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), + ) + infer_state = InferStateInfo() + infer_state.batch_size = b_req_idx.shape[0] + infer_state.b_req_idx = b_req_idx + 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, + ) + infer_state.att_state_role = role + if req_to_mtp_state_index is not None: + request_capacity = req_to_mtp_state_index.numel() + resources = AttentionStateManager(request_capacity=request_capacity) + config = ConvSSMStateConfig( + logical_layer_ids=(3,), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=2, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + for model_role in ("main", "draft0"): + linear = GDNState(config, request_capacity=request_capacity, mtp_step=draft_step, device="meta") + linear.mtp_state_index = ( + req_to_mtp_state_index if model_role == role else torch.zeros_like(req_to_mtp_state_index) + ) + resources.register_state(model_role, linear) + resources.finish_state_registration() + infer_state.att_state_manager = resources + state = gdn.LinearAttDecodeAttState( + backend=SimpleNamespace( + model=model, + uses_dynamic_spec_verify_layout=lambda: dynamic_layout, + ), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + return state, infer_state diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py index 72fb73a9e5..ae848f13de 100644 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -45,6 +45,41 @@ def find_spec(self, fullname, path=None, target=None): assert result.returncode == 0, result.stdout + result.stderr +@pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode"]) +def test_swa_workspace_includes_decode_padding_prefill(monkeypatch: pytest.MonkeyPatch, run_mode: str) -> None: + backend_class = attention.DeepseekV4FlashMlaFp8SparseAttBackend + model = SimpleNamespace( + args=SimpleNamespace(run_mode=run_mode, enable_prefill_microbatch_overlap=False), + config=dict( + compress_ratios=[0, 4, 128], + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + sliding_window=128, + ), + layers_infer=[SimpleNamespace(layer_num_=layer) for layer in range(3)], + tp_rank_=0, + tp_world_size_=4, + batch_max_tokens=4096, + max_req_num=8, + att_state_role="main", + ) + monkeypatch.setattr( + backend_class, + "get_gpu_workspace_buffer", + lambda key_name, workspace_size, dtype: torch.empty(workspace_size, dtype=dtype), + ) + workspace, capacity = backend_class.get_swa_prefill_workspace(model=model) + # D 节点仍执行 HOLD 请求的单 token prefill，但无需保留普通请求的长 prefill 临时区。 + assert workspace is not None + if run_mode == "decode": + assert capacity == 0 + assert workspace.shape == (3, 128, 585) + else: + assert capacity >= model.batch_max_tokens + assert workspace.shape == (3, capacity + 128, 585) + + def _pool(page_size, pages): stride = ((page_size * 584 + 575) // 576) * 576 buffer = torch.zeros((pages, stride), dtype=torch.uint8, device="cuda") @@ -63,21 +98,9 @@ def _case(heads, ratio): tokens = 5 model = type("TestModel", (), {})() model.config = {"num_attention_heads": heads, "compress_ratios": [0, 4, 128]} - from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4CacheConfig - - model.req_manager = SimpleNamespace( - cache_config=DeepseekV4CacheConfig.from_model_config( - dict( - num_hidden_layers=3, - compress_ratios=[0, 4, 128], - head_dim=512, - qk_rope_head_dim=64, - index_head_dim=128, - sliding_window=128, - ), - SimpleNamespace(mtp_mode=None), - ) - ) + model.config.update(head_dim=512, qk_rope_head_dim=64, index_head_dim=128, sliding_window=128) + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(3)] + model.tp_rank_ = 0 model.tp_world_size_ = 1 backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend(model=model) q = torch.randn((tokens, backend.padded_q_head_num, 512), dtype=torch.bfloat16, device="cuda") * 0.1 @@ -278,11 +301,12 @@ def test_prefill_cuda_graph_reads_new_request_state(monkeypatch, heads, ratio): packed_kv = (packed_kv[0], extra) params["extra_indices"].copy_(indices) + monkeypatch.setattr(backend, "get_att_input_params", lambda infer_state, logical_layer_id: packed_kv) + def new_request(length): infer_state = DeepseekV4InferStateInfo() infer_state.input_ids = torch.zeros(q.shape[0], dtype=torch.int64, device=q.device) infer_state.hidden_collector = NoopHiddenCollector() - infer_state.req_manager = SimpleNamespace(get_att_input_params=lambda layer_index: packed_kv) infer_state.attn_params = dict(params, swa_lengths=torch.full_like(params["swa_lengths"], length)) if ratio: infer_state.attn_params["extra_lengths"] = torch.full_like(params["extra_lengths"], length) diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py new file mode 100644 index 0000000000..8a981bc726 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -0,0 +1,119 @@ +"""FP8 后端按推理层传入的逻辑层号选择标定，覆盖稀疏层号、共享层和 draft。""" + +from dataclasses import replace +from types import SimpleNamespace +import weakref + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 +from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 +from lightllm.models.llama.infer_struct import LlamaInferStateInfo +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer +from unit_tests.common.att_resources.model_helpers import rebuild_full_state +from unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups +from unit_tests.common.att_resources.test_capacity import definition + + +@pytest.mark.parametrize("backend_name", ["fa3", "flashinfer"]) +@pytest.mark.parametrize("operation", ["prefill", "decode"]) +@pytest.mark.parametrize("role", ["main", "draft0"]) +@pytest.mark.parametrize("logical_layer_id,storage_layer_index", [(3, 0), (7, 0), (8, 1), (11, 1)]) +def test_fp8_backend_selects_calibration_from_logical_layer( + monkeypatch: pytest.MonkeyPatch, + backend_name: str, + operation: str, + role: str, + logical_layer_id: int, + storage_layer_index: int, +) -> None: + template = definition("static_head" if backend_name == "fa3" else "static_tensor") + state = rebuild_full_state( + template=template, + config=replace_logical_layer_groups(template.config, ((3, 7), (8, 11))), + token_capacity=8, + device="cpu", + ) + state.scales[0].fill_(2) + state.scales[1].fill_(5) + expected_scale = 2 if storage_layer_index == 0 else 5 + manager = AttentionStateManager(request_capacity=1) + manager.register_state(role, state) + manager.finish_state_registration() + infer = LlamaInferStateInfo() + infer.att_state_manager = manager + infer.att_state_role = role + infer.b_seq_len = torch.tensor([1], dtype=torch.int32) + infer.max_q_seq_len = 1 + calls = [] + + # 仅替换外部计算内核，保留推理层、公开 attention 接口及 State 标定查询链路。 + if backend_name == "fa3": + state.q_scales[0].fill_(3) + state.q_scales[1].fill_(7) + expected_q_scale = 3 if storage_layer_index == 0 else 7 + + def quantize(q: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + assert torch.all(scales == expected_q_scale) + return q + + def attend(**kwargs: object) -> torch.Tensor: + assert torch.all(kwargs["k_descale"] == expected_scale) + assert torch.all(kwargs["v_descale"] == expected_scale) + assert torch.all(kwargs["q_descale"] == expected_q_scale) + calls.append(operation) + return kwargs["q"] + + monkeypatch.setattr(fa3_fp8, "q_per_head_static_fp8_quant", quantize) + monkeypatch.setattr(fa3_fp8, "flash_attn_with_kvcache", attend) + state_class = fa3_fp8.Fp8Fa3PrefillAttState if operation == "prefill" else fa3_fp8.Fp8Fa3DecodeAttState + else: + state_class = ( + flashinfer_fp8.Fp8FlashInferPrefillAttState + if operation == "prefill" + else flashinfer_fp8.Fp8FlashInferDecodeAttState + ) + + attention = state_class(backend=SimpleNamespace(), _infer_state_weak_ref=weakref.ref(infer)) + if backend_name == "fa3": + attention.k_descale = state.scales[:, :2].unsqueeze(1) + attention.v_descale = state.scales[:, 2:].unsqueeze(1) + attention.page_table = attention.cu_seqlens_q = attention.cu_seqlens_k = None + attention.b_att_seq_len = infer.b_seq_len + attention.decode_max_q_seq_len = 1 + attention.causal = True + else: + + def run( + q: torch.Tensor, + paged_kv_cache: tuple[torch.Tensor, torch.Tensor], + k_scale: float, + v_scale: float, + out: torch.Tensor, + ) -> None: + assert k_scale == v_scale == expected_scale + calls.append(operation) + out.copy_(q) + + attention.prefill_wrapper = attention.decode_wrapper = SimpleNamespace(run=run) + + infer.prefill_att_state = infer.decode_att_state = attention + layer = LlamaTransformerLayerInfer.__new__(LlamaTransformerLayerInfer) + layer.layer_num_ = logical_layer_id + layer.tp_q_head_num_ = layer.tp_k_head_num_ = state.config.head_num + layer.head_dim_ = state.config.head_dim + layer.alloc_tensor = lambda shape, dtype, device: torch.empty(shape, dtype=dtype, device="cpu") + q = torch.ones((1, state.config.head_num, state.config.head_dim)) + if operation == "prefill": + output = layer._context_attention_kernel( + q=q, + kv=None, + infer_state=infer, + layer_weight=None, + ) + else: + output = layer._token_attention_kernel(q=q, infer_state=infer, layer_weight=None) + assert torch.equal(output, q) + assert calls == [operation] diff --git a/unit_tests/common/basemodel/attention/test_state_lifetime.py b/unit_tests/common/basemodel/attention/test_state_lifetime.py index aa1795058c..9e049d54b4 100644 --- a/unit_tests/common/basemodel/attention/test_state_lifetime.py +++ b/unit_tests/common/basemodel/attention/test_state_lifetime.py @@ -152,3 +152,23 @@ def replay(self): assert runtime_ref() is None graph.graph.clear() assert captured_ref() is None + + +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_secondary_state_without_primary(is_prefill): + owner = InferStateInfo() + owner.is_prefill = is_prefill + owner.position_ids = torch.arange(4) + slot = "prefill_att_state1" if is_prefill else "decode_att_state1" + state_cls = _PrefillState if is_prefill else _DecodeState + secondary = state_cls(_infer_state_weak_ref=weakref.ref(owner)) + setattr(owner, slot, secondary) + owner.init_att_state() + torch.testing.assert_close(secondary.positions, owner.position_ids) + if not is_prefill: + runtime = InferStateInfo() + runtime.position_ids = owner.position_ids + 7 + runtime.decode_att_state1 = _DecodeState(_infer_state_weak_ref=weakref.ref(runtime)) + runtime.decode_att_state1.init_state() + owner.copy_for_cuda_graph(runtime) + torch.testing.assert_close(secondary.positions, runtime.position_ids) diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index cb0d9fef54..123f603ff2 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -80,24 +80,6 @@ def test_prefill_unpad_restores_empty_output(): assert unpadded.prompt_logics.shape == (0, 4) -def _create_empty_decode_input(): - return ModelInput( - batch_size=0, - total_token_num=0, - max_q_seq_len=1, - max_kv_seq_len=0, - input_ids=torch.empty((0,), dtype=torch.int64), - b_req_idx=torch.empty((0,), dtype=torch.int32), - b_mtp_index=torch.empty((0,), dtype=torch.int32), - b_seq_len=torch.empty((0,), dtype=torch.int32), - b_position_delta=torch.empty((0,), dtype=torch.int32), - b_shared_seq_len=torch.empty((0,), dtype=torch.int32), - b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64), - is_prefill=False, - multimodal_params=[], - ) - - @torch.no_grad() def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): execution_configs = ( @@ -122,7 +104,6 @@ def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): model.tp_world_size_ = tp_world_size model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) - model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1) model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) pad_batch_sizes = [] @@ -195,7 +176,6 @@ def test_overlap_decode_initializes_attention_for_the_selected_execution_mode(le model.tp_world_size_ = 2 model.is_mtp_draft_model = False model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) - model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1) model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88) initializations = [] executions = [] @@ -344,3 +324,21 @@ def forward(state): assert seen_sizes == [expected_size] * (2 if overlap else 1) torch.testing.assert_close(output.logits, expected) torch.testing.assert_close(torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1), expected) + + +def _create_empty_decode_input(): + return ModelInput( + batch_size=0, + total_token_num=0, + max_q_seq_len=1, + max_kv_seq_len=0, + input_ids=torch.empty((0,), dtype=torch.int64), + b_req_idx=torch.empty((0,), dtype=torch.int32), + b_mtp_index=torch.empty((0,), dtype=torch.int32), + b_seq_len=torch.empty((0,), dtype=torch.int32), + b_position_delta=torch.empty((0,), dtype=torch.int32), + b_shared_seq_len=torch.empty((0,), dtype=torch.int32), + b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64), + is_prefill=False, + multimodal_params=[], + ) diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py index 54f2ccacf1..ef7c18a061 100644 --- a/unit_tests/common/basemodel/test_overlap_utils.py +++ b/unit_tests/common/basemodel/test_overlap_utils.py @@ -8,68 +8,6 @@ from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -def _empty_multimodal_params(batch_size: int): - return [{"images": [], "audios": []} for _ in range(batch_size)] - - -def _make_prefill_input(input_ids: list[int], req_idx: int, is_decode_req: bool) -> ModelInput: - token_num = len(input_ids) - return ModelInput( - batch_size=1, - total_token_num=token_num, - max_q_seq_len=token_num, - max_kv_seq_len=token_num, - max_cache_len=0, - input_ids=torch.tensor(input_ids, dtype=torch.int64), - b_req_idx=torch.tensor([req_idx], dtype=torch.int32), - b_mtp_index=torch.zeros(1, dtype=torch.int32), - b_seq_len=torch.tensor([token_num], dtype=torch.int32), - b_is_decode_req=torch.tensor([is_decode_req]), - b_ready_cache_len=torch.zeros(1, dtype=torch.int32), - b_prefill_start_loc=torch.zeros(1, dtype=torch.int32), - is_prefill=True, - b_prefill_has_output_cpu=[True], - multimodal_params=_empty_multimodal_params(1), - ) - - -def _make_decode_input() -> ModelInput: - batch_size = 6 - return ModelInput( - batch_size=batch_size, - total_token_num=39, - max_q_seq_len=1, - max_kv_seq_len=9, - input_ids=torch.arange(batch_size, dtype=torch.int64), - b_req_idx=torch.tensor([10, 11, 12, 12, 12, 12], dtype=torch.int32), - b_mtp_index=torch.tensor([0, 0, 0, 1, 2, 3], dtype=torch.int32), - b_seq_len=torch.tensor([4, 5, 6, 7, 8, 9], dtype=torch.int32), - b_position_delta=torch.arange(batch_size, dtype=torch.int32), - b_shared_seq_len=torch.arange(10, 16, dtype=torch.int32), - b_shared_radix_node_id=torch.arange(20, 26, dtype=torch.int64), - is_prefill=False, - multimodal_params=_empty_multimodal_params(batch_size), - ) - - -def _make_empty_decode_input() -> ModelInput: - return ModelInput( - batch_size=0, - total_token_num=0, - max_q_seq_len=1, - max_kv_seq_len=0, - input_ids=torch.empty((0,), dtype=torch.int64), - b_req_idx=torch.empty((0,), dtype=torch.int32), - b_mtp_index=torch.empty((0,), dtype=torch.int32), - b_seq_len=torch.empty((0,), dtype=torch.int32), - b_position_delta=torch.empty((0,), dtype=torch.int32), - b_shared_seq_len=torch.empty((0,), dtype=torch.int32), - b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64), - is_prefill=False, - multimodal_params=[], - ) - - def test_base_model_prefill_accepts_two_prebuilt_inputs(monkeypatch): model_input0 = _make_prefill_input([10, 11], req_idx=20, is_decode_req=False) model_input1 = _make_prefill_input([12], req_idx=21, is_decode_req=True) @@ -130,7 +68,12 @@ def record_to_cuda(self): def fake_gather(**kwargs): events.append("gather") - return torch.arange(40, 46, dtype=torch.int64, device="cuda") + return torch.arange( + start=40, + end=46, + dtype=torch.int64, + device="cuda", + ) monkeypatch.setattr(basemodel, "gather_token", fake_gather) @@ -184,7 +127,6 @@ def test_overlap_decode_cuda_pads_empty_side_and_unpads_outputs(monkeypatch): model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) model.graph = None model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) - model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(77,), page_size=1) infer_batch_sizes = [] def fake_create_inferstate(model_input, microbatch_index): @@ -207,3 +149,65 @@ def fake_create_inferstate(model_input, microbatch_index): assert infer_batch_sizes == [2, 2] assert output0.logits.shape == (1, 1) assert output1.logits.shape == (0, 1) + + +def _empty_multimodal_params(batch_size: int): + return [{"images": [], "audios": []} for _ in range(batch_size)] + + +def _make_prefill_input(input_ids: list[int], req_idx: int, is_decode_req: bool) -> ModelInput: + token_num = len(input_ids) + return ModelInput( + batch_size=1, + total_token_num=token_num, + max_q_seq_len=token_num, + max_kv_seq_len=token_num, + max_cache_len=0, + input_ids=torch.tensor(input_ids, dtype=torch.int64), + b_req_idx=torch.tensor([req_idx], dtype=torch.int32), + b_mtp_index=torch.zeros(1, dtype=torch.int32), + b_seq_len=torch.tensor([token_num], dtype=torch.int32), + b_is_decode_req=torch.tensor([is_decode_req]), + b_ready_cache_len=torch.zeros(1, dtype=torch.int32), + b_prefill_start_loc=torch.zeros(1, dtype=torch.int32), + is_prefill=True, + b_prefill_has_output_cpu=[True], + multimodal_params=_empty_multimodal_params(1), + ) + + +def _make_decode_input() -> ModelInput: + batch_size = 6 + return ModelInput( + batch_size=batch_size, + total_token_num=39, + max_q_seq_len=1, + max_kv_seq_len=9, + input_ids=torch.arange(batch_size, dtype=torch.int64), + b_req_idx=torch.tensor([10, 11, 12, 12, 12, 12], dtype=torch.int32), + b_mtp_index=torch.tensor([0, 0, 0, 1, 2, 3], dtype=torch.int32), + b_seq_len=torch.tensor([4, 5, 6, 7, 8, 9], dtype=torch.int32), + b_position_delta=torch.arange(batch_size, dtype=torch.int32), + b_shared_seq_len=torch.arange(10, 16, dtype=torch.int32), + b_shared_radix_node_id=torch.arange(20, 26, dtype=torch.int64), + is_prefill=False, + multimodal_params=_empty_multimodal_params(batch_size), + ) + + +def _make_empty_decode_input() -> ModelInput: + return ModelInput( + batch_size=0, + total_token_num=0, + max_q_seq_len=1, + max_kv_seq_len=0, + input_ids=torch.empty((0,), dtype=torch.int64), + b_req_idx=torch.empty((0,), dtype=torch.int32), + b_mtp_index=torch.empty((0,), dtype=torch.int32), + b_seq_len=torch.empty((0,), dtype=torch.int32), + b_position_delta=torch.empty((0,), dtype=torch.int32), + b_shared_seq_len=torch.empty((0,), dtype=torch.int32), + b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64), + is_prefill=False, + multimodal_params=[], + ) diff --git a/unit_tests/common/dsv4_test_utils.py b/unit_tests/common/dsv4_test_utils.py new file mode 100644 index 0000000000..35caf326e3 --- /dev/null +++ b/unit_tests/common/dsv4_test_utils.py @@ -0,0 +1,473 @@ +"""旧 V4 行为用例的测试夹具；数据操作委托新的 State，不向生产代码添加兼容接口。""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +import torch +import triton + +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.common.attention_state import AttentionStateManager +from lightllm.models.deepseek_v4.config import DSV4AttentionConfig +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel +from lightllm.common.attention_state.bounded_att import FP8E8M0SlidingWindowState, WindowCompressorState +from lightllm.common.attention_state.bounded_att.window_pages_pool import WindowPagesPool +from lightllm.common.attention_state.bounded_att.utils import create_page_pool_tensor +from lightllm.common.attention_state.full_att.compressed_att.compressed_att_state import ( + CompressedAttentionConfig, + CompressedAttentionState, +) +from lightllm.common.attention_state.allocators import FullAttKVAllocator, StatePageAllocator +from lightllm.common.req_manager import ReqManager + +DSV4_COMPRESSED_KV_TOKEN_ALIGNMENT = 256 + + +@dataclass +class DeepseekV4CacheConfig: + """V4 KV 布局和大小页状态缓存配置。""" + + compress_ratios: tuple[int, ...] + head_dim: int + rope_head_dim: int + indexer_head_dim: int + sliding_window: int + + # Packed FP8 KV format shared by Flash and Pro. + quant_group_dim = 64 + scale_alignment_bytes = 8 + # 每个 SWA 物理页的 token 数，用于 KV/scale 打包和 ring 容量对齐。 + swa_page_token_num = 128 + compressed_kv_token_alignment = DSV4_COMPRESSED_KV_TOKEN_ALIGNMENT + short_compress_ratio = 4 + long_compress_ratio = 128 + + def __post_init__(self): + self.compress_ratios = tuple(self.compress_ratios) + self.layer_num = len(self.compress_ratios) + c4_layers = [layer for layer, rate in enumerate(self.compress_ratios) if rate == self.short_compress_ratio] + c128_layers = [layer for layer, rate in enumerate(self.compress_ratios) if rate == self.long_compress_ratio] + self.layer_to_c4_idx = {layer: slot for slot, layer in enumerate(c4_layers)} + self.layer_to_c128_idx = {layer: slot for slot, layer in enumerate(c128_layers)} + self.n_c4, self.n_c128 = len(self.layer_to_c4_idx), len(self.layer_to_c128_idx) + compressed_kv_layers = [layer for layer, rate in enumerate(self.compress_ratios) if rate] + self.layer_to_compressed_kv_idx = {layer: slot for slot, layer in enumerate(compressed_kv_layers)} + self.kv_data_bytes = self.head_dim + self.rope_head_dim + self.kv_scale_bytes = ( + triton.cdiv((self.head_dim - self.rope_head_dim) // self.quant_group_dim, self.scale_alignment_bytes) + * self.scale_alignment_bytes + ) + self.kv_bytes = self.kv_data_bytes + self.kv_scale_bytes + self.indexer_bytes = self.indexer_head_dim + torch.float32.itemsize + self.swa_slot_bytes = self.get_packed_kv_bytes(self.swa_page_token_num) // self.swa_page_token_num + self.compressed_kv_slot_bytes = self._get_compressed_kv_slot_bytes() + + @classmethod + def from_model_config(cls, config: Any, args: Any): + from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num + + layer_num = config["num_hidden_layers"] + compress_ratios = list(config["compress_ratios"]) + if args.mtp_mode is not None: + layer_num += get_added_mtp_kv_layer_num() + # Runtime-added draft layers use sliding-window attention. + compress_ratios.extend([0] * (layer_num - len(compress_ratios))) + return cls( + compress_ratios=tuple(compress_ratios[:layer_num]), + head_dim=config["head_dim"], + rope_head_dim=config["qk_rope_head_dim"], + indexer_head_dim=config["index_head_dim"], + sliding_window=config["sliding_window"], + ) + + def get_packed_kv_bytes(self, token_count: int): + return triton.cdiv(token_count * self.kv_bytes, self.kv_data_bytes) * self.kv_data_bytes + + def get_swa_state_shape(self): + return (self.layer_num, self.sliding_window, self.kv_bytes) + + def get_c4_state_shape(self): + return (self.n_c4, self.short_compress_ratio, 4 * self.head_dim) + + def get_indexer_state_shape(self): + return (self.n_c4, self.short_compress_ratio, 4 * self.indexer_head_dim) + + def get_runtime_state_shapes(self, mtp_step: int): + """每个请求的完整 ring；P/D 使用相同布局，可直接复制槽位。""" + swa_tokens = triton.cdiv(self.sliding_window + 2 * mtp_step, self.swa_page_token_num) * self.swa_page_token_num + c4_tokens = 2 * self.short_compress_ratio + mtp_step + c128_tokens = self.long_compress_ratio + mtp_step + return ( + (self.layer_num, swa_tokens, self.swa_slot_bytes), + (self.n_c4, c4_tokens, 4 * self.head_dim), + (self.n_c4, c4_tokens, 4 * self.indexer_head_dim), + (self.n_c128, c128_tokens, 2 * self.head_dim), + ) + + def get_cpu_cache_full_att_bytes(self): + args = get_env_start_args() + return args.cpu_cache_token_page_size * sum(self.compressed_kv_slot_bytes) + + def get_cpu_cache_big_page_bytes(self): + # MLA KV and compressor state are replicated across TP ranks. + swa_bytes = math.prod(self.get_swa_state_shape()) + c4_bytes = math.prod(self.get_c4_state_shape()) * torch.float32.itemsize + indexer_bytes = math.prod(self.get_indexer_state_shape()) * torch.float32.itemsize + return self.get_cpu_cache_full_att_bytes() + swa_bytes + c4_bytes + indexer_bytes + + def _get_compressed_kv_slot_bytes(self): + block_token_num = self.compressed_kv_token_alignment + # 整块按 KV 数据行对齐，同时能均分到 token 槽位。 + block_bytes_alignment = math.lcm(block_token_num, self.kv_data_bytes) + slot_bytes_list = [] + for ratio in self.compress_ratios: + if ratio == 0: + continue + compressed_token_num = block_token_num // ratio + block_bytes = self.get_packed_kv_bytes(compressed_token_num) + if ratio == self.short_compress_ratio: + block_bytes += compressed_token_num * self.indexer_bytes + block_bytes = triton.cdiv(block_bytes, block_bytes_alignment) * block_bytes_alignment + slot_bytes_list.append(block_bytes // block_token_num) + return tuple(slot_bytes_list) + + +def state_config(config: Any): + return DSV4AttentionConfig( + logical_layer_ids=tuple(range(config.layer_num)), + compress_ratios=config.compress_ratios, + head_dim=config.head_dim, + rope_head_dim=config.rope_head_dim, + indexer_head_dim=config.indexer_head_dim, + sliding_window=config.sliding_window, + ) + + +class DeepseekV4StateCacheManager: + def __init__(self, size: int, cache_config: DeepseekV4CacheConfig, keep_num: int = 0): + # 仅供旧用例保留原联合页视图；生产原子 State 分别拥有自己的页池。 + shapes = ( + cache_config.get_swa_state_shape(), + cache_config.get_c4_state_shape(), + cache_config.get_indexer_state_shape(), + ) + lengths = [ + math.prod(shape) * dtype.itemsize + for shape, dtype in zip(shapes, (torch.uint8, torch.float32, torch.float32)) + ] + self.buffer = create_page_pool_tensor(size=(size, sum(lengths)), dtype=torch.uint8) + offset = 0 + # manager 按规范 State 名排序，联合页依次为 indexer、KV compressor、SWA。 + for name, index, dtype in ( + ("indexer_state", 2, torch.float32), + ("c4_state", 1, torch.float32), + ("swa_state", 0, torch.uint8), + ): + length = lengths[index] + setattr(self, name, self.buffer[:, offset : offset + length].view(dtype).view(size, *shapes[index])) + offset += length + self.size = size + self.allocator = StatePageAllocator(size=size, reserved_slots=keep_num) + self.keep_num = keep_num + + def alloc_one_state_cache(self): + slots = self.allocator.alloc(1) + return slots[0] if slots is not None else None + + def alloc_state_cache(self, count: int): + return self.allocator.alloc(count) + + def free_state_cache(self, slots: list[int]): + self.allocator.free(slots) + + def get_free_cache_num(self): + return self.allocator.free_count + + def get_state_cache(self, slot: int): + return self.swa_state[slot], self.c4_state[slot], self.indexer_state[slot] + + +class DeepseekV4MemoryManager: + def __init__( + self, + size: int, + dtype: torch.dtype, + head_num: int, + head_dim: int, + layer_num: int, + cache_config: DeepseekV4CacheConfig, + ): + self.cache_config = cache_config + self.page_size = get_env_start_args().page_size + self.size = size = size // self.page_size * self.page_size + self.allocator = FullAttKVAllocator(size=size, page_size=self.page_size) + self.manager = AttentionStateManager(request_capacity=3) + self.full_states = {} + self.kv_buffer = [] + config = state_config(cache_config) + for layer, ratio in enumerate(config.compress_ratios): + if not ratio: + continue + state = CompressedAttentionState( + layout_config=CompressedAttentionConfig( + logical_layer_groups=((layer,),), + compress_ratio=ratio, + kv_dim=config.head_dim, + rope_dim=config.rope_head_dim, + block_token_size=config.compressed_kv_token_alignment, + indexer_dim=config.indexer_head_dim if ratio == 4 else 0, + tp_rank=config.tp_rank, + tp_world_size=config.tp_world_size, + ), + page_size=self.page_size, + ) + state.allocate_storage(size) + self.manager.register_state("main", state) + self.full_states[layer] = state + self.kv_buffer.append(state.storage[0]) + big_tokens = get_env_start_args().linear_att_hash_page_size * get_env_start_args().linear_att_page_block_num + count = triton.cdiv(size, big_tokens) + self.big_page_buffers = DeepseekV4StateCacheManager(count + 2, cache_config, keep_num=2) + self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = count + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = count + 1 + self.operator = SimpleNamespace(copy_mem_to_mem=self.copy_mem_to_mem) + + def alloc(self, count: int): + return self.allocator.alloc(count) + + def free(self, indices: torch.Tensor): + self.allocator.free(indices) + + def get_cell_size(self): + return sum(state.get_token_bytes() for state in self.full_states.values()) + + def get_att_input_params(self, layer: int): + index = self.cache_config.layer_to_compressed_kv_idx.get(layer) + return self.kv_buffer[index] if index is not None else None + + def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): + for state in self.full_states.values(): + state.copy_full_att_tokens(src_mem_index, dst_mem_index) + + def alloc_paged_kv_move_buffer(self, count: int, page_size: int): + from copy import copy + + self.manager = copy(self.manager) + for original in self.manager.get_bounded_attention_states(): + state = copy(original) + if isinstance(state, FP8E8M0SlidingWindowState): + if hasattr(self, "req_to_swa"): + state.swa_buffer = self.req_to_swa + else: + field = ( + "req_to_c4_indexer_state" + if state.config.target == "indexer" + else ("req_to_c4_state" if state.config.window_size == 8 else "req_to_c128_state") + ) + if hasattr(self, field): + state.buffer = getattr(self, field) + setattr(self.manager, state.state_name, state) + self.manager._get_states.cache_clear() + get_env_start_args().pd_kv_page_size = page_size + self.manager.get_pd_operator().get_pd_full_att_page_layouts.cache_clear() + self.manager.get_pd_operator().get_pd_bounded_att_page_layouts.cache_clear() + self.manager.get_pd_operator().get_page_bytes.cache_clear() + self.kv_move_buffer = torch.empty( + (count, self.manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda" + ) + return self.kv_move_buffer + + def get_pd_kv_move_buffer_size(self): + return self.manager.get_pd_operator().get_page_bytes() * get_env_start_args().pd_kv_page_num + + def get_paged_kv_move_buffer_shape(self, count: int, page_size: int): + return (count, self.manager.get_pd_operator().get_page_bytes()) + + def write_mem_to_page_kv_move_buffer( + self, + indices: torch.Tensor, + page: int, + dp_index: int, + managers: list[DeepseekV4MemoryManager], + tp_size: int, + kind: str = "kv", + req_idx: int | None = None, + ): + for manager in managers[dp_index * tp_size : dp_index * tp_size + 1]: + manager.manager.get_pd_operator().write_page( + pd_page_buffers=self.kv_move_buffer, + pd_page_index=page, + kind=kind, + layouts=manager.manager.get_pd_operator().get_pd_full_att_page_layouts() + if kind == "kv" + else manager.manager.get_pd_operator().get_pd_bounded_att_page_layouts(), + token_indices=torch.tensor(indices, device="cuda", dtype=torch.int64), + req_idx=req_idx, + state_token_len=getattr(manager, "state_token_len", 0), + ) + + def read_page_kv_move_buffer_to_mem( + self, + indices: torch.Tensor, + page: int, + dp_index: int, + managers: list[DeepseekV4MemoryManager], + tp_size: int, + kind: str = "kv", + req_idx: int | None = None, + ): + for manager in managers[dp_index * tp_size : (dp_index + 1) * tp_size]: + manager.manager.get_pd_operator().read_page( + pd_page_buffers=self.kv_move_buffer, + pd_page_index=page, + kind=kind, + layouts=manager.manager.get_pd_operator().get_pd_full_att_page_layouts() + if kind == "kv" + else manager.manager.get_pd_operator().get_pd_bounded_att_page_layouts(), + token_indices=torch.tensor(indices, device="cuda", dtype=torch.int64), + req_idx=req_idx, + state_token_len=getattr(manager, "state_token_len", 0), + ) + + +class DeepseekV4ReqManager(ReqManager): + def __init__(self, count: int, length: int, memory: DeepseekV4MemoryManager | None, config: Any, mtp_step: int = 0): + super().__init__(count, length) + self.cache_config = config + self.states = DeepseekV4TpPartModel.create_bounded_att_states( + model_config=dict( + compress_ratios=config.compress_ratios, + head_dim=config.head_dim, + qk_rope_head_dim=config.rope_head_dim, + index_head_dim=config.indexer_head_dim, + sliding_window=config.sliding_window, + ), + args=SimpleNamespace(mtp_step=mtp_step), + model_logical_layer_ids=tuple(range(config.layer_num)), + request_capacity=count + 1, + device="cuda", + ) + self.swa_state = next(state for state in self.states if isinstance(state, FP8E8M0SlidingWindowState)) + self.config = self.swa_state.config + self.swa_buffer = self.req_to_swa = self.swa_state.swa_buffer + shapes = config.get_runtime_state_shapes(mtp_step) + buffers = [] + for window_size, target, shape in ((8, "kv", shapes[1]), (8, "indexer", shapes[2]), (128, "kv", shapes[3])): + state = next( + ( + item + for item in self.states + if isinstance(item, WindowCompressorState) + and item.config.window_size == window_size + and item.config.target == target + ), + None, + ) + buffers.append( + state.buffer if state is not None else torch.zeros((shape[0], count + 1, *shape[1:]), device="cuda") + ) + self.req_to_c4_state, self.req_to_c4_indexer_state, self.req_to_c128_state = buffers + self.window = SimpleNamespace( + swa_buffer=self.swa_buffer, + c4_buffer=buffers[0], + indexer_buffer=buffers[1], + c128_buffer=buffers[2], + request_capacity=count + 1, + swa_capacity_tokens=self.swa_state.swa_capacity_tokens, + HOLD_SWA_SLOT=self.swa_state.HOLD_SWA_SLOT, + ) + self.c4_state_tokens = self.req_to_c4_state.shape[2] + self.c128_state_tokens = self.req_to_c128_state.shape[2] + self.sliding_window = config.sliding_window + self.swa_capacity_tokens, self.HOLD_SWA_SLOT = self.window.swa_capacity_tokens, self.window.HOLD_SWA_SLOT + if memory is not None: + self.bind_mem_manager(memory) + + def bind_mem_manager(self, memory: DeepseekV4MemoryManager | None): + self.memory = memory + self.bind_hold_page(memory.allocator.hold_token_indices) + if not memory.manager._state_registration_finished: + for state in self.states: + memory.manager.register_state("main", state) + memory.manager.finish_state_registration() + memory.manager.full_att_allocator = memory.allocator + memory.manager.bounded_att_big_page_allocator = memory.big_page_buffers.allocator + memory.manager.bounded_att_cpu_load_big_page_slot = memory.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID + self._bind_page_views(memory.big_page_buffers, "big") + + def _bind_page_views(self, pool, kind): + for state in self.states: + if isinstance(state, WindowCompressorState) and not state.config.overlap_tokens: + continue + view = ( + pool.swa_state + if isinstance(state, FP8E8M0SlidingWindowState) + else (pool.indexer_state if state.config.target == "indexer" else pool.c4_state) + ) + atomic_pool = object.__new__(WindowPagesPool) + atomic_pool.buffer = view + setattr(state, f"{kind}_pages_pool", atomic_pool) + + def get_compressor_state(self, logical_layer_id, target="kv"): + return next( + state + for state in self.states + if isinstance(state, WindowCompressorState) + and state.config.target == target + and state.contains_logical_layer(logical_layer_id) + ) + + def get_compressed_kv_state(self, logical_layer_id): + return self.memory.full_states[logical_layer_id] + + def get_compressed_kv_view(self, logical_layer_id: int): + return self.memory.full_states[logical_layer_id].get_layer_view(logical_layer_id) + + def get_att_input_params(self, layer: int): + compressed = self.get_compressed_kv_view(layer).kv if layer in self.memory.full_states else None + return ( + self.req_to_swa[layer].view(-1, self.config.get_packed_kv_bytes(self.config.page_token_size)), + compressed, + ) + + def create_small_page_cache_manager(self, size: int): + self.small_page_buffers = DeepseekV4StateCacheManager(size, self.cache_config) + self._bind_page_views(self.small_page_buffers, "small") + self.memory.manager.bounded_att_small_page_allocator = self.small_page_buffers.allocator + return self.small_page_buffers + + def save_state(self, req: Any, slot: int, pool: DeepseekV4StateCacheManager): + for state in self.states: + if state.big_pages_pool is not None: + external = ( + pool.swa_state + if isinstance(state, FP8E8M0SlidingWindowState) + else (pool.indexer_state if state.config.target == "indexer" else pool.c4_state) + ) + state._copy_window( + req_idx=req.req_idx, + external=external[slot], + state_token_len=req.get_chuncked_input_token_len(), + write=True, + ) + + def restore_state(self, req: Any, pool: DeepseekV4StateCacheManager, slot: int): + for state in self.states: + if state.big_pages_pool is not None: + external = ( + pool.swa_state + if isinstance(state, FP8E8M0SlidingWindowState) + else (pool.indexer_state if state.config.target == "indexer" else pool.c4_state) + ) + state._copy_window( + req_idx=req.req_idx, external=external[slot], state_token_len=req.cur_kv_len, write=False + ) + + def free(self, reqs: list[int], indices: torch.Tensor = None): + super().free(reqs) + if indices is not None: + self.memory.free(indices) diff --git a/unit_tests/common/state_cache_manager/test_linear_att.py b/unit_tests/common/state_cache_manager/test_linear_att.py deleted file mode 100644 index bceb67f4f6..0000000000 --- a/unit_tests/common/state_cache_manager/test_linear_att.py +++ /dev/null @@ -1,64 +0,0 @@ -import pytest -import torch - -from lightllm.common.state_cache_manager import LinearAttCacheConfig - - -@pytest.mark.parametrize("draft_layers", [0, 1, 3]) -@pytest.mark.parametrize( - "layer_types,expected_indices,linear_layers", - [ - ( - ("linear_attention", "linear_attention", "linear_attention", "full_attention") * 2, - (0, 1, 2, 0, 3, 4, 5, 1), - 6, - ), - ( - ( - "deepseek_sparse_attention", - "linear_attention", - "linear_attention", - "deepseek_sparse_attention", - "linear_attention", - ), - (0, 0, 1, 1, 2), - 3, - ), - ( - ("linear_attention", "full_attention", "linear_attention") * 2 + ("full_attention",), - (0, 0, 1, 2, 1, 3, 2), - 4, - ), - ], -) -def test_hybrid_layer_mapping_and_mtp_tail(layer_types, expected_indices, linear_layers, draft_layers): - 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=128, - global_linear_k_heads=1, - global_linear_v_heads=1, - num_linear_k_heads=1, - num_linear_v_heads=1, - head_linear_k_dim=128, - head_linear_v_dim=128, - conv_kernel_size=4, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=torch.float32, - layer_types=layer_types, - draft_full_att_kv_layer_num=draft_layers, - ) - assert config.all_layer_num == len(layer_types) - assert config.linear_layer_num == linear_layers - full_layers = len(layer_types) - linear_layers - assert config.get_main_model_full_att_layer_num() == full_layers - assert config.get_full_att_kv_layer_num_with_draft_model() == full_layers + draft_layers - for layer, (kind, expected) in enumerate(zip(layer_types, expected_indices)): - get_index = ( - config.get_linear_state_layer_index if kind == "linear_attention" else config.get_full_att_kv_layer_index - ) - assert get_index(layer) == expected - for draft in range(draft_layers): - assert config.get_full_att_kv_layer_index(len(layer_types) + draft) == full_layers + draft diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index 9b03f2e192..1278e5e195 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -4,12 +4,9 @@ import pytest import torch -from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager -from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4CacheConfig -from lightllm.common.req_manager import DeepseekV4ReqManager +from unit_tests.common.dsv4_test_utils import DeepseekV4MemoryManager, DeepseekV4CacheConfig, DeepseekV4ReqManager from lightllm.models.deepseek_v4.triton_kernel.build_compressed_kv_indices import build_compressed_kv_indices from lightllm.utils.envs_utils import get_env_start_args -from lightllm.models.deepseek_v4.triton_kernel.cache_io import get_compressed_kv_buffer, get_indexer_k_buffer from lightllm.models.deepseek_v4.triton_kernel.swa_cache import get_swa_token_indices as _get_swa_token_indices @@ -17,9 +14,9 @@ def _compressed_kv_pools(manager): c4 = next(iter(manager.cache_config.layer_to_c4_idx)) c128 = next(iter(manager.cache_config.layer_to_c128_idx)) return ( - get_compressed_kv_buffer(manager, c4).unsqueeze(0), - get_indexer_k_buffer(manager, c4).unsqueeze(0), - get_compressed_kv_buffer(manager, c128).unsqueeze(0), + manager.full_states[c4].get_layer_view(c4).kv.unsqueeze(0), + manager.full_states[c4].get_layer_view(c4).indexer_k.unsqueeze(0), + manager.full_states[c128].get_layer_view(c128).kv.unsqueeze(0), ) @@ -57,6 +54,7 @@ def cache(monkeypatch, tmp_path, request): "mtp_step": mtp_step, "mtp_dynamic_verify": False, "enable_ep_moe": False, + "enable_dp_prompt_cache_fetch": False, } ), ) @@ -116,68 +114,76 @@ def test_big_page_pool_reserves_cpu_cache_staging_slots(cache): assert pool.get_free_cache_num() == capacity -def test_copy_mem_to_mem_preserves_token_slot_indices(cache): +def test_copy_mem_to_mem_preserves_complete_block_indices(cache): manager, _ = cache - source = torch.tensor([3, 257, 509], device="cuda") - destination = torch.tensor([1025, 1028, 1407], device="cuda") + # 压缩块内的 KV/scale/indexer 不可拆分；允许按非连续、乱序的完整块复制。 + source = torch.cat([torch.arange(start, start + 256, device="cuda") for start in (0, 512, 256)]) + destination = torch.cat([torch.arange(start, start + 256, device="cuda") for start in (1024, 1536, 1280)]) for buffer in manager.kv_buffer: buffer.zero_() - buffer[source] = torch.arange(1, 4, device="cuda", dtype=torch.uint8)[:, None, None] + values = torch.arange(1, 4, device="cuda", dtype=torch.uint8).repeat_interleave(256) + buffer[source] = values[:, None, None] manager.operator.copy_mem_to_mem(source, destination) for buffer in manager.kv_buffer: torch.testing.assert_close(buffer[destination], buffer[source], rtol=0, atol=0) def test_cache_io_copies_individual_token_slots(cache): - from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_compressed_kv_pages - manager, _ = cache source = torch.tensor([3, 17, 257, 509], device="cuda") destination = torch.tensor([1025, 1028, 1407, 2047], device="cuda") for buffer in manager.kv_buffer: buffer.zero_() buffer[source] = torch.arange(1, 5, device="cuda", dtype=torch.uint8)[:, None, None] - pages = torch.empty((1, 4 * manager.get_cell_size()), dtype=torch.uint8, pin_memory=True) - page_ids = torch.zeros(1, dtype=torch.int32, device="cuda") - copy_compressed_kv_pages(manager, source, page_ids, pages, to_pages=True, page_size=4) - copy_compressed_kv_pages(manager, destination, page_ids, pages, to_pages=False, page_size=4) + layouts = manager.manager.get_cpu_cache_operator().get_page_layouts() + pages = torch.empty( + (1, manager.manager.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True + ) + for state in manager.full_states.values(): + state.write_to_cpu(source, pages, [0], layout=layouts[state.state_name]) + state.read_from_cpu(pages, [0], destination, layout=layouts[state.state_name]) torch.cuda.synchronize() for buffer in manager.kv_buffer: torch.testing.assert_close(buffer[destination], buffer[source], rtol=0, atol=0) assert buffer[destination - 1].eq(0).all() untouched = torch.tensor([1024, 1026, 1406, 1408], device="cuda") assert buffer[untouched].eq(0).all() - snapshot = manager.get_index_kv_buffer(source) - assert all(isinstance(value, torch.Tensor) for value in snapshot.values()) - # The shared preload path truncates dictionary values along their token axis. - snapshot = {name: value[:, :2].cpu() for name, value in snapshot.items()} for buffer in manager.kv_buffer: buffer[destination] = 0 - manager.load_index_kv_buffer(destination[:2], snapshot) + # 同一 CPU 页只读前两个 token，不能覆盖目标中剩余槽位。 + for state in manager.full_states.values(): + state.read_from_cpu(pages, [0], destination[:2], layout=layouts[state.state_name]) for buffer in manager.kv_buffer: torch.testing.assert_close(buffer[destination[:2]], buffer[source[:2]], rtol=0, atol=0) assert buffer[destination[2:]].eq(0).all() def test_request_runtime_is_independent_of_compressed_kv_capacity(cache, monkeypatch): - from lightllm.common.kv_cache_mem_manager import MemoryManager - original, requests = cache - observed = [] - - def profile(manager, mem_fraction): - observed.append(requests.req_to_swa.flatten(1, 2).nbytes) - assert _state_buffer(requests, "c4").is_cuda and _state_buffer(requests, "c128").is_cuda - assert not hasattr(manager, "swa_buffer") - manager.size = 8192 - - monkeypatch.setattr(MemoryManager, "profile_size", profile) - manager = DeepseekV4MemoryManager(None, torch.bfloat16, 1, 512, 3, original.cache_config) - assert observed == [requests.req_to_swa.flatten(1, 2).nbytes] - assert manager.size == 8192 - for buffer in manager.kv_buffer: - assert buffer.shape[0] == manager.size + manager.page_size - assert sum(buffer.nbytes for buffer in manager.kv_buffer) == original.get_cell_size() * (8192 + manager.page_size) + runtime = requests.window + before = tuple( + tensor.data_ptr() + for tensor in ( + runtime.swa_buffer, + runtime.c4_buffer, + runtime.indexer_buffer, + runtime.c128_buffer, + ) + ) + # 扩展压缩 KV 容量不能重新创建或移动按请求数分配的运行态。 + for state in original.full_states.values(): + state.allocate_storage(token_capacity=16384) + assert state.storage.shape[1] == 16384 + state.page_size + assert before == tuple( + tensor.data_ptr() + for tensor in ( + runtime.swa_buffer, + runtime.c4_buffer, + runtime.indexer_buffer, + runtime.c128_buffer, + ) + ) + assert runtime.request_capacity == requests.max_request_num + 1 @pytest.mark.parametrize("page_size", [256, 1024, 4096]) @@ -241,10 +247,19 @@ def _test_backend(requests, token_capacity=8192): model = type("TestModel", (), {})() model.req_manager = requests + model.att_state_manager = requests.memory.manager + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] + model.tp_rank_ = 0 + model.mtp_previous_draft_models = [] + model.__dict__["att_state_role"] = "main" model.config = { "num_attention_heads": 64, - "compress_ratios": requests.cache_config.compress_rates, + "compress_ratios": requests.cache_config.compress_ratios, "index_topk": 512, + "head_dim": 512, + "qk_rope_head_dim": 64, + "index_head_dim": 128, + "sliding_window": requests.sliding_window, } model.tp_world_size_ = 1 model.max_req_num = requests.max_request_num @@ -270,6 +285,10 @@ def _prepare_prefill(requests, req_ids, starts, ends, microbatch_index=0): state.prepare_swa_workspace( SimpleNamespace( req_manager=requests, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, b_req_idx=req_ids, b_ready_cache_len=starts, b_seq_len=ends, @@ -384,7 +403,7 @@ def test_indexer_writer_only_writes_closed_groups(cache): _compressed_kv_pools(manager)[1].zero_() from lightllm.models.deepseek_v4.triton_kernel.destindex_copy_indexer_k import destindex_copy_indexer_k - destindex_copy_indexer_k(keys, slots[:259], positions, get_indexer_k_buffer(manager, 0), 64) + destindex_copy_indexer_k(keys, slots[:259], positions, manager.full_states[0].get_layer_view(0).indexer_k, 64) closed = torch.arange(3, 259, 4, device="cuda") actual = _read_packed_rows(_compressed_kv_pools(manager)[1], 0, slots[closed] // 4) # Reference quantization is independent of the cache writer and stays in the test. @@ -428,6 +447,11 @@ def test_model_selects_held_slots_without_allocating_swa(cache, model_kind): ] model = model_cls.__new__(model_cls) model.req_manager = requests + model.att_state_manager = requests.memory.manager + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(requests.cache_config.layer_num)] + model.tp_rank_ = 0 + model.mtp_previous_draft_models = [] + model.__dict__["att_state_role"] = "main" model_input = ModelInput( batch_size=1, total_token_num=257, @@ -640,6 +664,7 @@ def test_pd_kv_and_att_state_roundtrip(cache, length): "req_to_c128_state", ): setattr(memory, name, getattr(runtime, name)) + memory.state_token_len = length memory.alloc_paged_kv_move_buffer(1, manager.page_size) for start in range(0, length, manager.page_size): end = min(length, start + manager.page_size) @@ -654,26 +679,33 @@ def test_pd_kv_and_att_state_roundtrip(cache, length): atol=0, ) - # Copy all ring slots, including the SWA-only layer and partially filled compressor pools. + # PD 传递与 MTP 容量无关的逻辑窗口，不再传递无效候选槽位或 padding。 source_memory.write_mem_to_page_kv_move_buffer([], 0, 0, [source_memory], 1, "att_state", source) manager.kv_move_buffer.copy_(source_memory.kv_move_buffer) manager.read_page_kv_move_buffer_to_mem([], 0, 0, [manager], 1, "att_state", destination) + positions = torch.arange(max(0, length - requests.sliding_window), length, device="cuda") + for layer in range(requests.cache_config.layer_num): + torch.testing.assert_close( + _read_swa(requests, layer, destination, positions, prefill=False), + _read_swa(prefill, layer, source, positions, prefill=False), + rtol=0, + atol=0, + ) ring_size = requests.swa_capacity_tokens - torch.testing.assert_close( - requests.req_to_swa[:, destination], - prefill.req_to_swa[:, source], - rtol=0, - atol=0, - ) - # Neither another request's ring nor the prefill workspace belongs to this transfer. assert requests.req_to_swa.flatten(1, 2)[:, : destination * ring_size].eq(0).all() assert requests.req_to_swa.flatten(1, 2)[:, (destination + 1) * ring_size :].eq(0).all() - for source_buffer, target_buffer in ( - (prefill.req_to_c4_state, requests.req_to_c4_state), - (prefill.req_to_c4_indexer_state, requests.req_to_c4_indexer_state), - (prefill.req_to_c128_state, requests.req_to_c128_state), + for source_buffer, target_buffer, width in ( + (prefill.req_to_c4_state, requests.req_to_c4_state, 8), + (prefill.req_to_c4_indexer_state, requests.req_to_c4_indexer_state, 8), + (prefill.req_to_c128_state, requests.req_to_c128_state, 128), ): - torch.testing.assert_close(source_buffer[:, source], target_buffer[:, destination], rtol=0, atol=0) + positions = torch.arange(max(0, length - width), length, device="cuda") + torch.testing.assert_close( + source_buffer[:, source, positions % source_buffer.shape[2]], + target_buffer[:, destination, positions % target_buffer.shape[2]], + rtol=0, + atol=0, + ) prefill.free([source], src) requests.free([destination], dst) assert manager.allocator.can_use_mem_size == manager.size @@ -715,6 +747,7 @@ def test_pd_transport_completion_covers_peer_gpu_state(cache, length): cudart = ctypes.CDLL(f"libcudart.so.{torch.version.cuda.split('.')[0]}") assert cudart.cudaDeviceEnablePeerAccess(1, 0) in (0, 704) cudart.cudaGetLastError() + transport.state_token_len = peer.state_token_len = length transport.alloc_paged_kv_move_buffer(1, transport.page_size) mem_managers = [transport, peer] for start in range(0, length, transport.page_size): @@ -734,16 +767,23 @@ def test_pd_transport_completion_covers_peer_gpu_state(cache, length): completed = torch.cuda.Event() completed.record(copy_stream) completed.synchronize() - # 完成事件只需保证数据可读，无需目标 stream 已空闲；用独立 stream 校验整块 state。 + # 在独立流验证逻辑窗口；无效 MTP 槽位不属于 PD 数据。 with torch.cuda.stream(consumer_stream): - torch.testing.assert_close( - requests.req_to_swa[:, destination], - requests.req_to_swa[:, source], - rtol=0, - atol=0, - ) - for buffer in (requests.req_to_c4_state, requests.req_to_c4_indexer_state, requests.req_to_c128_state): - torch.testing.assert_close(buffer[:, destination], buffer[:, source], rtol=0, atol=0) + positions = torch.arange(max(0, length - requests.sliding_window), length, device="cuda:1") + for layer in range(requests.cache_config.layer_num): + torch.testing.assert_close( + _read_swa(requests, layer, destination, positions, prefill=False), + _read_swa(requests, layer, source, positions, prefill=False), + rtol=0, + atol=0, + ) + for buffer, width in ( + (requests.req_to_c4_state, 8), + (requests.req_to_c4_indexer_state, 8), + (requests.req_to_c128_state, 128), + ): + positions = torch.arange(max(0, length - width), length, device="cuda:1") % buffer.shape[2] + torch.testing.assert_close(buffer[:, destination, positions], buffer[:, source, positions], rtol=0, atol=0) for buffer in peer.kv_buffer: torch.testing.assert_close(buffer[src.long()], buffer[dst.long()], rtol=0, atol=0) requests.free([source, destination], torch.cat([src, dst])) @@ -771,6 +811,10 @@ def compress(req_idx, slots, start, end, decode=False): state = SimpleNamespace( mem_manager=manager, req_manager=requests, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, mem_index=slots[start:end], position_ids=positions, is_prefill=not decode, @@ -921,16 +965,13 @@ def test_swa_ring_attention_matches_full_cache_after_speculative_retries(cache, def build(): if dspark: build_dspark_swa_indices( - req_idx, - query_positions, - requests, - indices, - lengths, - write_slots, - sliding_window=128, + token_req_idx=req_idx, + position_ids=query_positions, + swa_state=requests.swa_state, + swa_indices=indices, + swa_lengths=lengths, + swa_write_slots=write_slots, block_size=token_num, - hold_req_id=requests.HOLD_REQUEST_ID, - hold_swa_slot=requests.HOLD_SWA_SLOT, ) else: build_swa_indices(req_idx, query_positions, requests, indices, lengths, write_slots) @@ -940,7 +981,16 @@ def build(): "Model", (), { - "config": {"num_attention_heads": 64, "compress_ratios": [0]}, + "config": dict( + num_attention_heads=64, + compress_ratios=[0, 0, 0], + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + sliding_window=128, + ), + "layers_infer": [SimpleNamespace(layer_num_=2)], + "tp_rank_": 0, "tp_world_size_": 1, "req_manager": requests, }, @@ -979,6 +1029,10 @@ def forward(): reference_kv = backend.get_att_input_params( SimpleNamespace( req_manager=requests, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, is_prefill=True, prefill_att_state=requests._test_prefill_states[source], ), @@ -1040,7 +1094,7 @@ def test_hybrid_radix_hit_fork_pause_abort_and_eviction(cache, monkeypatch, hash small = requests.create_small_page_cache_manager(2) block_num = 2048 // hash_page_size big_page_tokens = block_num * hash_page_size - radix = HybridAttPagedRadixCache(manager.size, 0, hash_page_size, block_num, manager, small) + radix = HybridAttPagedRadixCache(manager.size, 0, hash_page_size, block_num, manager.manager) args = get_env_start_args() args.linear_att_hash_page_size = hash_page_size args.linear_att_page_block_num = block_num @@ -1048,6 +1102,7 @@ def test_hybrid_radix_hit_fork_pause_abort_and_eviction(cache, monkeypatch, hash hit_len = tail_len if hit_extra_page else big_page_tokens for name, value in { "req_manager": requests, + "att_state_manager": manager.manager, "radix_cache": radix, "args": args, "is_hybrid_att_model": True, @@ -1197,6 +1252,10 @@ def test_c4_indexer_metadata_scoring_and_graph_replay(cache, monkeypatch, mode): indexer = indexer_module.DeepseekV4IndexInfer(0, config, 1) state = infer_struct_module.DeepseekV4InferStateInfo() state.mem_manager, state.req_manager = manager, requests + state.att_state_manager = requests.memory.manager + state.att_state_role = "main" + state.first_logical_layer_id = 0 + state.first_c4_logical_layer_id = 0 state.is_prefill = mode != "decode" state.microbatch_index, state.max_kv_seq_len = 0, 4096 state.b_req_idx = torch.tensor([req_idx] if state.is_prefill else [req_idx] * tokens, device="cuda") @@ -1253,7 +1312,8 @@ def select(): @pytest.mark.parametrize("is_prefill", [False, True]) -def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill): +@pytest.mark.parametrize("owned_layer_ids", [(0, 1, 2), (0, 2), (2,)]) +def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill, owned_layer_ids): from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace @@ -1262,6 +1322,8 @@ def test_forward_indices_are_isolated_and_graph_inputs_refresh(cache, is_prefill second, second_slots = _reserve(manager, requests, 512) backend = _test_backend(requests) model = backend.model + model.layers_infer = [layer for layer in model.layers_infer if layer.layer_num_ in owned_layer_ids] + has_c128_layer = 1 in owned_layer_ids model.graph = SimpleNamespace(can_run=lambda *args: True, graph_max_len_in_batch=16384) model.dsv4_workspace = DeepseekV4Workspace(model) rope = torch.zeros((1024, 32), device="cuda") @@ -1272,6 +1334,10 @@ def make_state(lengths, microbatch_index): state = DeepseekV4InferStateInfo() state.is_prefill, state.microbatch_index = is_prefill, microbatch_index state.req_manager, state.mem_manager = requests, manager + state.att_state_manager = requests.memory.manager + state.att_state_role = "main" + state.first_logical_layer_id = 0 + state.first_c4_logical_layer_id = 0 state.batch_size = 2 state.max_kv_seq_len = max(lengths) state.input_ids = torch.zeros(2, dtype=torch.int64, device="cuda") @@ -1290,7 +1356,11 @@ def make_state(lengths, microbatch_index): swa_before = captured_state.dsv4_swa_indices.clone() other_state.dsv4_swa_indices.fill_(-7) torch.testing.assert_close(captured_state.dsv4_swa_indices, swa_before) - assert captured_state.dsv4_c128_indices.shape == (2, 64 if is_prefill else 128) + if has_c128_layer: + assert captured_state.dsv4_c128_indices.shape == (2, 64 if is_prefill else 128) + else: + assert captured_state.dsv4_c128_indices is None + assert captured_state.dsv4_c128_lengths is None captured_state.init_c4_indices(512) other_state.init_c4_indices(512) captured_state.dsv4_c4_indices.fill_(11) @@ -1302,7 +1372,8 @@ def make_state(lengths, microbatch_index): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=stream): swa_output = captured_state.dsv4_swa_indices.clone() - c128_output = captured_state.dsv4_c128_indices.clone() + if has_c128_layer: + c128_output = captured_state.dsv4_c128_indices.clone() updated = make_state([258, 386], 0) assert updated.dsv4_c4_indices is None if is_prefill: @@ -1311,7 +1382,8 @@ def make_state(lengths, microbatch_index): captured_state.copy_for_cuda_graph(updated) graph.replay() torch.testing.assert_close(swa_output, updated.dsv4_swa_indices) - torch.testing.assert_close(c128_output, updated.dsv4_c128_indices) + if has_c128_layer: + torch.testing.assert_close(c128_output, updated.dsv4_c128_indices) assert captured_state.dsv4_c4_indices.eq(11).all() requests.free([first, second], torch.cat((first_slots, second_slots))) @@ -1333,7 +1405,15 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, monkeypatch.setitem(args, "mtp_mode", "dspark") monkeypatch.setitem(args, "enable_decode_microbatch_overlap", True) monkeypatch.setitem(args, "enable_prefill_microbatch_overlap", False) - config = dict(sliding_window=128, index_topk=512, num_attention_heads=64, compress_ratios=[4, 128, 0]) + config = dict( + sliding_window=128, + index_topk=512, + num_attention_heads=64, + compress_ratios=[4, 128, 0], + head_dim=512, + qk_rope_head_dim=64, + index_head_dim=128, + ) model = type("TestModel", (), {})() vars(model).update( config=config, @@ -1344,7 +1424,14 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, has_vision=False, args=args, req_manager=requests, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, + tp_rank_=0, tp_world_size_=1, + att_state_role="main", + layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], ) backend = DeepseekV4FlashMlaFp8SparseAttBackend(model) model.block_size = 3 @@ -1356,6 +1443,10 @@ def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, state.is_prefill, state.microbatch_index = mode == "prefill", microbatch_index state.is_draft_model = True state.req_manager, state.mem_manager = requests, manager + state.att_state_manager = requests.memory.manager + state.att_state_role = "main" + state.first_logical_layer_id = 0 + state.first_c4_logical_layer_id = 0 if state.is_prefill: state.batch_size = 1 state.input_ids = torch.ones(3, dtype=torch.int64, device="cuda") @@ -1420,7 +1511,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( from lightllm.server.multi_level_kv_cache import CacheTier from lightllm.server.router.model_infer.infer_batch import InferReq from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common - from lightllm.common.state_cache_manager.deepseek4 import DeepseekV4StateCacheManager + from unit_tests.common.dsv4_test_utils import DeepseekV4StateCacheManager manager, requests = cache source, src = _reserve(manager, requests, length) @@ -1448,6 +1539,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( tail = slot else: source_checkpoints[end] = slot + offload_before = manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID].clone() # The live ring no longer contains the prompt's continuation checkpoints. requests.req_to_swa.flatten(1, 2).fill_(99) _state_buffer(requests, "c4").fill_(88) @@ -1491,6 +1583,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( producer = torch.cuda.current_stream() context = common.g_infer_context monkeypatch.setattr(context, "req_manager", requests) + monkeypatch.setattr(context, "att_state_manager", manager.manager) monkeypatch.setattr( context, "radix_cache", @@ -1512,7 +1605,10 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( disable_linear_att_small_page_cpu_cache=False, ) module.backend = SimpleNamespace( - is_master_in_dp=True, dp_world_size=1, model=SimpleNamespace(mem_manager=manager, req_manager=requests) + is_master_in_dp=True, + dp_world_size=1, + is_hybrid_att_model=True, + model=SimpleNamespace(att_state_manager=manager.manager, req_manager=requests), ) module.cpu_cache_client = client module.cpu_cache_handle_queue = deque() @@ -1541,7 +1637,7 @@ def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( if tail is not None: torch.testing.assert_close( manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID], - expected_checkpoints[-1], + offload_before, rtol=0, atol=0, ) @@ -1597,16 +1693,8 @@ def allocate(req, count): rtol=0, atol=0, ) - if exhaust_small_pool: - assert restored.tail_small_page_buffer_id is None - else: - assert restored.tail_small_page_buffer_id is not None - torch.testing.assert_close( - requests.small_page_buffers.buffer[restored.tail_small_page_buffer_id], - expected_checkpoints[-1], - rtol=0, - atol=0, - ) + # CPU 碎页直接经 LOAD 大页恢复运行态，不重复占用普通小页。 + assert restored.tail_small_page_buffer_id is None dst = requests.req_to_token_indexs[destination, :length] for pool in _compressed_kv_pools(manager): torch.testing.assert_close( @@ -1668,6 +1756,10 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): slots = _swa_slots(requests, req, positions) state = SimpleNamespace( req_manager=requests, + swa_state=requests.swa_state, + get_compressor_state=requests.get_compressor_state, + get_compressed_kv_state=requests.get_compressed_kv_state, + get_compressed_kv_view=requests.get_compressed_kv_view, is_prefill=True, prefill_att_state=requests._test_prefill_states[req], dsv4_swa_write_slots=slots, @@ -1700,27 +1792,39 @@ def test_prefill_fused_writer_preserves_old_prefix_and_runtime_tail(cache): def test_cpu_cache_preserves_already_published_pages(cache): - from lightllm.models.deepseek_v4.triton_kernel.cache_io import copy_kv_buffer_to_cpu_cache - manager, requests = cache req, slots = _reserve(manager, requests, 2048) _randomize_compressed_kv(manager) layout = manager.cache_config pages = torch.full((3, layout.get_cpu_cache_big_page_bytes()), 221, dtype=torch.uint8, pin_memory=True) - page_ids = torch.tensor([1], dtype=torch.int32, device="cuda") - readies = torch.ones(1, dtype=torch.bool, device="cuda") + page_ids = [1] + readies = [True] state_id = manager.big_page_buffers.alloc_one_state_cache() manager.big_page_buffers.buffer[state_id].fill_(7) - checkpoint_ids = torch.tensor([state_id], dtype=torch.int64, device="cuda") - copy_kv_buffer_to_cpu_cache(manager, slots, page_ids, readies, checkpoint_ids, pages) + checkpoint_ids = [state_id] + manager.manager.get_cpu_cache_operator().write_req( + token_indices=slots, + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), + cpu_page_ready_flags=readies, + big_page_slot_ids=checkpoint_ids, + ) torch.cuda.synchronize() assert pages.eq(221).all() - readies.zero_() - copy_kv_buffer_to_cpu_cache(manager, slots, page_ids, readies, checkpoint_ids, pages) + readies[0] = False + manager.manager.get_cpu_cache_operator().write_req( + token_indices=slots, + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=manager.manager.get_cpu_cache_operator().get_page_layouts(), + cpu_page_ready_flags=readies, + big_page_slot_ids=checkpoint_ids, + ) torch.cuda.synchronize() offset = 0 - for pool in manager.kv_buffer: - expected = pool[slots.long()].cpu().reshape(-1) + for state in manager.manager.get_full_attention_states(): + expected = state.storage[0, slots.long()].cpu().reshape(-1) torch.testing.assert_close(pages[1, offset : offset + expected.numel()], expected, rtol=0, atol=0) offset += expected.numel() assert pages[1, layout.get_cpu_cache_full_att_bytes() :].eq(7).all() @@ -1767,6 +1871,10 @@ def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_le state = DeepseekV4InferStateInfo() state.is_prefill = True state.req_manager, state.mem_manager = requests, manager + state.att_state_manager = requests.memory.manager + state.att_state_role = "main" + state.first_logical_layer_id = 0 + state.first_c4_logical_layer_id = 0 state.input_ids = torch.tensor(input_ids[ready:], dtype=torch.int64, device="cuda") state.b_req_idx = torch.tensor([req_idx], dtype=torch.int32, device="cuda") state.b_seq_len = torch.tensor([len(input_ids)], dtype=torch.int32, device="cuda") @@ -1825,6 +1933,10 @@ def test_prefill_metadata_does_not_read_cuda_scalars(cache, padding_tokens): infer_state = DeepseekV4InferStateInfo() infer_state.is_prefill = True infer_state.req_manager = requests + infer_state.att_state_manager = requests.memory.manager + infer_state.att_state_role = "main" + infer_state.first_logical_layer_id = 0 + infer_state.first_c4_logical_layer_id = 0 infer_state.mem_manager = manager infer_state.input_ids = torch.ones(5, dtype=torch.int64, device="cuda") infer_state.b_req_idx = torch.tensor(req_ids, dtype=torch.int32, device="cuda") diff --git a/unit_tests/common/test_paged_mem_manager.py b/unit_tests/common/test_paged_mem_manager.py deleted file mode 100644 index 60a7c86337..0000000000 --- a/unit_tests/common/test_paged_mem_manager.py +++ /dev/null @@ -1,59 +0,0 @@ -from types import SimpleNamespace -import pickle - -import pytest -import torch - -from lightllm.common.kv_cache_mem_manager.mem_manager import MemoryManager -from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager - - -def test_shared_memory_exports_preserve_variable_pickle_lengths(monkeypatch, tmp_path): - from lightllm.common.kv_cache_mem_manager import mem_manager 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 = MemoryManager.__new__(MemoryManager) - try: - manager.write_to_shm(SimpleNamespace(req_to_token_indexs=None)) - 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() - - -def test_kv_cache_allocation_requires_complete_pages(): - manager = MemoryManager.__new__(MemoryManager) - manager.page_size = 4 - manager.allocator = SimpleNamespace(alloc=lambda size: torch.arange(size, dtype=torch.int32)) - - with pytest.raises(AssertionError, match="must be a multiple of page_size 4"): - manager.alloc(5) - - assert manager.alloc(8).tolist() == list(range(8)) - - -@pytest.mark.parametrize( - "manager_class", - [MemoryManager, Deepseek2MemoryManager, PPLINT4KVMemoryManager, PPLINT8KVMemoryManager], -) -def test_kv_cache_buffer_size_requires_complete_pages(manager_class): - manager = manager_class.__new__(manager_class) - manager.page_size = 4 - - with pytest.raises(AssertionError, match="KV cache size 5 must be a multiple of page_size 4"): - manager._init_buffers(size=5, dtype=torch.float16, head_num=1, head_dim=8, layer_num=1) diff --git a/unit_tests/common/test_quantized_paged_mem_manager.py b/unit_tests/common/test_quantized_paged_mem_manager.py deleted file mode 100644 index 84fe266613..0000000000 --- a/unit_tests/common/test_quantized_paged_mem_manager.py +++ /dev/null @@ -1,33 +0,0 @@ -from types import SimpleNamespace - -import torch - -from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager -from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager - - -def test_quantized_kv_buffers_include_the_full_hold_page(monkeypatch): - allocations = [] - - def fake_empty(shape, **kwargs): - allocations.append(shape) - return SimpleNamespace(shape=shape) - - monkeypatch.setattr(torch, "empty", fake_empty) - - int8_manager = PPLINT8KVMemoryManager.__new__(PPLINT8KVMemoryManager) - int8_manager.page_size = 4 - int8_manager.group_quant_size = 8 - int8_manager._init_buffers(size=32, dtype=torch.float16, head_num=2, head_dim=16, layer_num=3) - - int4_manager = PPLINT4KVMemoryManager.__new__(PPLINT4KVMemoryManager) - int4_manager.page_size = 4 - int4_manager.group_quant_size = 8 - int4_manager._init_buffers(size=32, dtype=torch.float16, head_num=2, head_dim=16, layer_num=3) - - assert allocations == [ - (3, 36, 4, 16), - (3, 36, 4, 2), - (3, 36, 4, 8), - (3, 36, 4, 2), - ] diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index 224c308f28..09a5e10448 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -10,7 +10,7 @@ from lightllm.server.router.model_infer.mode_backend import base_backend -class _FakeMemManager: +class _FakeAllocator: page_size = 4 def __init__(self): @@ -24,16 +24,15 @@ def alloc(self, size): return result -def test_bind_mem_manager_initializes_hold_request_with_reserved_page(): +def test_bind_hold_page_initializes_only_padding_request(): req_manager = ReqManager.__new__(ReqManager) req_manager.HOLD_REQUEST_ID = 2 req_manager.req_list = SimpleNamespace(is_all_free=lambda: True) req_manager.req_to_token_indexs = torch.full((3, 8), -1, dtype=torch.int32) - mem_manager = SimpleNamespace(page_size=4, HOLD_TOKEN_MEMINDEXES=(100, 101, 102, 103)) + req_manager.page_size = 4 + req_manager.bind_hold_page((100, 101, 102, 103)) - req_manager.bind_mem_manager(mem_manager) - - assert req_manager.mem_manager is mem_manager + assert not hasattr(req_manager, "mem_manager") assert req_manager.req_to_token_indexs[2].tolist() == [100, 101, 102, 103] * 2 assert req_manager.req_to_token_indexs[:2].eq(-1).all() @@ -43,31 +42,7 @@ def test_hold_request_indexes_require_all_requests_released(): req_manager.req_list = SimpleNamespace(is_all_free=lambda: False) with pytest.raises(AssertionError, match="all requests are released"): - req_manager.init_hold_request_indexs() - - -def _make_context(monkeypatch): - mem_manager = _FakeMemManager() - req_manager = SimpleNamespace( - mem_manager=mem_manager, - req_to_token_indexs=torch.full((2, 16), -1, dtype=torch.int32), - ) - context = SimpleNamespace( - args=SimpleNamespace(page_size=4), - req_manager=req_manager, - radix_cache=None, - ) - monkeypatch.setattr(generic_pre_process, "g_infer_context", context) - monkeypatch.setattr(base_backend, "g_infer_context", context) - backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) - backend.args = context.args - return context, backend - - -def _make_req(req_idx): - req = SimpleNamespace(req_idx=req_idx, cur_kv_len=0, hold_kv_len=0) - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - return req + req_manager.bind_hold_page((100, 101, 102, 103)) def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypatch): @@ -77,17 +52,17 @@ def test_request_reuses_reserved_page_tail_before_allocating_next_page(monkeypat mem_indexes = backend._alloc_req_kv_mem(req, alloc_token_num=4) assert mem_indexes.tolist() == [0, 1, 2, 3] assert req.hold_kv_len == 4 - assert context.req_manager.mem_manager.alloc_sizes == [4] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [4] assert context.req_manager.req_to_token_indexs[0, :4].tolist() == [0, 1, 2, 3] req.cur_kv_len = 3 assert backend._alloc_req_kv_mem(req, alloc_token_num=0) is None - assert context.req_manager.mem_manager.alloc_sizes == [4] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [4] req.cur_kv_len = 4 backend._alloc_req_kv_mem(req, alloc_token_num=4) assert req.hold_kv_len == 8 - assert context.req_manager.mem_manager.alloc_sizes == [4, 4] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [4, 4] assert context.req_manager.req_to_token_indexs[0, :8].tolist() == list(range(8)) @@ -121,8 +96,8 @@ def __getitem__(self, key): target = _CopyTarget() context = SimpleNamespace( + att_state_manager=SimpleNamespace(full_att_allocator=_FakeAllocator()), req_manager=SimpleNamespace( - mem_manager=_FakeMemManager(), req_to_token_indexs=_ReqToTokenIndexes(target), ), radix_cache=None, @@ -175,7 +150,12 @@ def test_need_token_num_distinguishes_compute_and_page_allocation(): ], ) def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( - monkeypatch, hold_kv_len, can_alloc_token_num, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause + monkeypatch, + hold_kv_len, + can_alloc_token_num, + batch_max_tokens, + expected_alloc_sizes, + second_req_wait_pause, ): context, backend = _make_context(monkeypatch) backend.args.enable_cpu_cache = False @@ -209,7 +189,7 @@ def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( context.req_manager.req_to_token_indexs[req.req_idx, :hold_kv_len] = torch.arange( start_index, start_index + hold_kv_len, dtype=torch.int32 ) - context.req_manager.mem_manager.next_index = 2 * hold_kv_len + context.att_state_manager.full_att_allocator.next_index = 2 * hold_kv_len backend._filter_not_ready_reqs = lambda req_ids: reqs copy_modes = [] original_alloc_req_kv_mem = backend._alloc_req_kv_mem @@ -229,7 +209,7 @@ def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): assert prefill_reqs == [reqs[0]] assert decode_reqs == [] assert prefill_tokens == 3 - assert context.req_manager.mem_manager.alloc_sizes == expected_alloc_sizes + assert context.att_state_manager.full_att_allocator.alloc_sizes == expected_alloc_sizes assert reqs[0].hold_kv_len == 8 assert not reqs[0].wait_pause assert reqs[1].hold_kv_len == hold_kv_len @@ -265,7 +245,7 @@ def test_decode_scheduler_uses_non_blocking_request_table_copy(monkeypatch): req.get_cur_total_len = lambda: 5 req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) context.req_manager.req_to_token_indexs[0, :4] = torch.arange(4, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 4 + context.att_state_manager.full_att_allocator.next_index = 4 backend._filter_not_ready_reqs = lambda req_ids: [req] copy_modes = [] @@ -287,7 +267,7 @@ def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): assert decode_reqs == [req] assert prefill_tokens == 0 assert req.hold_kv_len == 16 - assert context.req_manager.mem_manager.alloc_sizes == [12] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [12] assert copy_modes == [True] @@ -327,7 +307,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step req.get_cur_total_len = lambda: req.cur_kv_len + 1 req.decode_need_token_num = MethodType(InferReq.decode_need_token_num, req) context.req_manager.req_to_token_indexs[0, :page_size] = torch.arange(page_size, dtype=torch.int32) - context.req_manager.mem_manager.next_index = page_size + context.att_state_manager.full_att_allocator.next_index = page_size backend._filter_not_ready_reqs = lambda req_ids: [req] for _ in range(expected_alloc_token_num): @@ -336,7 +316,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step req.cur_kv_len += 1 assert req.hold_kv_len == page_size + expected_alloc_token_num - assert context.req_manager.mem_manager.alloc_sizes == [expected_alloc_token_num] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [expected_alloc_token_num] def test_decode_reserves_mtp_headroom(monkeypatch): @@ -352,7 +332,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): req.args = context.args context.req_manager.req_to_token_indexs[0, :4] = torch.arange(4, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 4 + context.att_state_manager.full_att_allocator.next_index = 4 token_num, alloc_token_num = InferReq.decode_need_token_num(req) assert token_num == 6 assert alloc_token_num == 8 @@ -363,7 +343,7 @@ def test_decode_reserves_mtp_headroom(monkeypatch): assert run_reqs == [req, req, req] assert model_input.b_seq_len.tolist() == [4, 5, 6] assert req.hold_kv_len == 12 - assert context.req_manager.mem_manager.alloc_sizes == [8] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [8] assert context.req_manager.req_to_token_indexs[0, :12].tolist() == list(range(12)) @@ -380,7 +360,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): req.get_radix_cache_shared_len = lambda: 0 req.args = context.args context.req_manager.req_to_token_indexs[0, :3] = torch.arange(3, dtype=torch.int32) - context.req_manager.mem_manager.next_index = 3 + context.att_state_manager.full_att_allocator.next_index = 3 token_num, alloc_token_num = InferReq.decode_need_token_num(req) assert token_num == alloc_token_num == 6 backend._alloc_req_kv_mem(req, alloc_token_num) @@ -388,7 +368,7 @@ def test_page_size_one_uses_the_same_scheduler_preallocation(monkeypatch): model_input, _ = generic_pre_process.prepare_decode_inputs([req]) assert req.hold_kv_len == 9 - assert context.req_manager.mem_manager.alloc_sizes == [6] + assert context.att_state_manager.full_att_allocator.alloc_sizes == [6] assert context.req_manager.req_to_token_indexs[0, :9].tolist() == list(range(9)) @@ -457,35 +437,6 @@ def test_linear_attention_frees_preallocated_pages_before_first_forward(monkeypa assert free_token_indexes[0].tolist() == list(range(8)) -def _make_paused_req(req_idx, target_kv_len, page_size, match_kv_len=0, chunk_kv_len=None): - if chunk_kv_len is None: - chunk_kv_len = target_kv_len - req = SimpleNamespace( - args=SimpleNamespace(page_size=page_size), - req_id=req_idx, - req_idx=req_idx, - cur_kv_len=0, - hold_kv_len=0, - paused=True, - match_call_count=0, - shared_kv_node=None, - shm_req=SimpleNamespace(is_paused=True, shm_cur_kv_len=0), - get_cur_total_len=lambda: target_kv_len, - get_chuncked_input_token_len=lambda: chunk_kv_len, - ) - - def match_radix_cache(): - req.match_call_count += 1 - req.cur_kv_len = match_kv_len - req.hold_kv_len = match_kv_len - req.shm_req.shm_cur_kv_len = match_kv_len - - req._match_radix_cache = match_radix_cache - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - req.prefill_need_token_num = MethodType(InferReq.prefill_need_token_num, req) - return req - - def test_recover_paused_reqs_uses_page_allocation_need(monkeypatch): infer_context = InferenceContext.__new__(InferenceContext) infer_context.args = SimpleNamespace(page_size=4) @@ -494,7 +445,6 @@ def test_recover_paused_reqs_uses_page_allocation_need(monkeypatch): freed_indexes = [] infer_context.req_manager = SimpleNamespace( req_to_token_indexs=torch.arange(16, dtype=torch.int32).reshape(2, 8), - free_token=lambda indexes: freed_indexes.extend(indexes.tolist()), ) infer_context.get_can_alloc_token_num = lambda: 5 large_req = _make_paused_req(req_idx=0, target_kv_len=5, page_size=4) @@ -519,10 +469,14 @@ def test_recover_paused_reqs_requires_capacity_for_the_full_sequence(monkeypatch infer_context.radix_cache = None infer_context.req_manager = SimpleNamespace( req_to_token_indexs=torch.arange(8, dtype=torch.int32)[None, :], - free_token=lambda indexes: None, ) infer_context.get_can_alloc_token_num = lambda: 4 - req = _make_paused_req(req_idx=0, target_kv_len=8, page_size=4, chunk_kv_len=4) + req = _make_paused_req( + req_idx=0, + target_kv_len=8, + page_size=4, + chunk_kv_len=4, + ) monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", False) monkeypatch.setattr(infer_batch, "custom_cat", lambda tensors: torch.cat(tensors)) @@ -538,7 +492,12 @@ def test_recover_paused_reqs_accounts_for_rematched_prefix(monkeypatch): infer_context.args = SimpleNamespace(page_size=4) infer_context.backend = SimpleNamespace(disable_chunked_prefill=False) infer_context.get_can_alloc_token_num = lambda: 12 - req = _make_paused_req(req_idx=0, target_kv_len=9, page_size=4, match_kv_len=8) + req = _make_paused_req( + req_idx=0, + target_kv_len=9, + page_size=4, + match_kv_len=8, + ) monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", False) infer_context.recover_paused_reqs([req], is_master_in_dp=True) @@ -547,3 +506,62 @@ def test_recover_paused_reqs_accounts_for_rematched_prefix(monkeypatch): assert req.paused is False assert req.shm_req.is_paused is False assert req.match_call_count == 1 + + +def _make_context(monkeypatch): + mem_manager = _FakeAllocator() + req_manager = SimpleNamespace( + req_to_token_indexs=torch.full((2, 16), -1, dtype=torch.int32), + ) + context = SimpleNamespace( + args=SimpleNamespace(page_size=4), + req_manager=req_manager, + att_state_manager=SimpleNamespace(full_att_allocator=mem_manager), + radix_cache=None, + ) + monkeypatch.setattr(generic_pre_process, "g_infer_context", context) + monkeypatch.setattr(base_backend, "g_infer_context", context) + backend = base_backend.ModeBackend.__new__(base_backend.ModeBackend) + backend.args = context.args + return context, backend + + +def _make_req(req_idx): + req = SimpleNamespace(req_idx=req_idx, cur_kv_len=0, hold_kv_len=0) + req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) + return req + + +def _make_paused_req( + req_idx, + target_kv_len, + page_size, + match_kv_len=0, + chunk_kv_len=None, +): + if chunk_kv_len is None: + chunk_kv_len = target_kv_len + req = SimpleNamespace( + args=SimpleNamespace(page_size=page_size), + req_id=req_idx, + req_idx=req_idx, + cur_kv_len=0, + hold_kv_len=0, + paused=True, + match_call_count=0, + shared_kv_node=None, + shm_req=SimpleNamespace(is_paused=True, shm_cur_kv_len=0), + get_cur_total_len=lambda: target_kv_len, + get_chuncked_input_token_len=lambda: chunk_kv_len, + ) + + def match_radix_cache(): + req.match_call_count += 1 + req.cur_kv_len = match_kv_len + req.hold_kv_len = match_kv_len + req.shm_req.shm_cur_kv_len = match_kv_len + + req._match_radix_cache = match_radix_cache + req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) + req.prefill_need_token_num = MethodType(InferReq.prefill_need_token_num, req) + return req diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py new file mode 100644 index 0000000000..c07fd123f0 --- /dev/null +++ b/unit_tests/conftest.py @@ -0,0 +1,40 @@ +"""Shared fixtures for tests that construct runtime resources.""" + +from __future__ import annotations + + +from contextlib import suppress +from uuid import uuid4 + +import pytest + + +@pytest.fixture +def kv_allocator_environment(monkeypatch: pytest.MonkeyPatch): + """Provide a service identity and reclaim the allocator's real shared counters.""" + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.allocators import full_att_kv_allocator + from lightllm.utils.envs_utils import get_unique_server_name + + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "test_kv_allocator_" + uuid4().hex) + get_unique_server_name.cache_clear() + counters = [] + shared_int = full_att_kv_allocator.SharedInt + + def track_counter(name: str): + counter = shared_int(name) + counters.append(counter) + return counter + + monkeypatch.setattr(full_att_kv_allocator, "SharedInt", track_counter) + yield + AttentionStateManager._get_states.cache_clear() + names = set() + for counter in counters: + if counter.shm.name not in names: + with suppress(FileNotFoundError): + counter.shm.unlink() + names.add(counter.shm.name) + counter.shm.close() + get_unique_server_name.cache_clear() diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py index 8032134734..8394d18b48 100644 --- a/unit_tests/models/deepseek_v4/test_model.py +++ b/unit_tests/models/deepseek_v4/test_model.py @@ -7,6 +7,29 @@ from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer +@pytest.mark.parametrize("capture", [False, True], ids=["eager", "capture"]) +@pytest.mark.parametrize("draft", [False, True], ids=["main", "dspark"]) +def test_pre_layer_decode_uses_current_graph_flag(monkeypatch: pytest.MonkeyPatch, capture: bool, draft: bool) -> None: + from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo + from lightllm.models.deepseek_v4.layer_infer import pre_layer_infer + from lightllm.models.deepseek_v4_dspark.layer_infer.pre_layer_infer import DeepseekV4DSparkPreLayerInfer + + layer_class = DeepseekV4DSparkPreLayerInfer if draft else pre_layer_infer.DeepseekV4PreLayerInfer + layer = layer_class.__new__(layer_class) + layer.hc_mult = 4 + state = DeepseekV4InferStateInfo() + state.is_cuda_graph_capture = capture + hidden = torch.ones((1, 8)) + state.mtp_draft_input_hiddens = hidden + monkeypatch.setattr(pre_layer_infer.LlamaMultimodalPreLayerInfer, "token_forward", lambda *args: hidden) + monkeypatch.setattr(pre_layer_infer, "hc_expand", lambda embeddings, streams: embeddings.repeat(1, streams)) + + output = layer.token_forward(input_ids=None, infer_state=state, layer_weight=None) + assert output.shape == (1, 32) + # 捕获时保留输入引用供后续回放使用；普通 decode 可以释放不再使用的 draft 特征。 + assert state.mtp_draft_input_hiddens is (hidden if capture else None) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_rotary_cache_keeps_unscaled_swa_and_yarn_compression(): model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) @@ -61,7 +84,7 @@ def test_unsupported_parallel_layouts_fail_before_inference(monkeypatch, balance def test_memory_check_rejects_sequence_larger_than_full_token_pool(): model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - model.mem_manager = SimpleNamespace(size=60416) + model.att_state_manager = SimpleNamespace(full_att_allocator=SimpleNamespace(size=60416)) model.batch_max_tokens = 8192 model.max_seq_length = 262169 model.args = SimpleNamespace(performance_mode=None) @@ -72,7 +95,7 @@ def test_memory_check_rejects_sequence_larger_than_full_token_pool(): def test_memory_check_still_requires_one_prefill_batch(): model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - model.mem_manager = SimpleNamespace(size=8192) + model.att_state_manager = SimpleNamespace(full_att_allocator=SimpleNamespace(size=8192)) model.batch_max_tokens = 8192 model.max_seq_length = 262169 model.args = SimpleNamespace(performance_mode=None) diff --git a/unit_tests/models/deepseek_v4/test_vision_integration.py b/unit_tests/models/deepseek_v4/test_vision_integration.py index 76322216e6..8290ab7d89 100644 --- a/unit_tests/models/deepseek_v4/test_vision_integration.py +++ b/unit_tests/models/deepseek_v4/test_vision_integration.py @@ -204,7 +204,7 @@ def test_swa_index_adds_bidirectional_image_visibility(): HOLD_SWA_SLOT=256, ) requests = SimpleNamespace( - cache_config=SimpleNamespace(swa_page_token_num=128), + config=SimpleNamespace(page_token_size=128), swa_capacity_tokens=128, HOLD_SWA_SLOT=128, HOLD_REQUEST_ID=1, @@ -276,11 +276,15 @@ def allocate(req, count): module.args = SimpleNamespace(page_size=256) module.backend = SimpleNamespace( is_master_in_dp=True, + is_hybrid_att_model=False, _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table), - mem_manager=SimpleNamespace( - operator=SimpleNamespace(load_cpu_cache_to_gpu=lambda **kwargs: loaded.append(req.cur_kv_len)) + att_state_manager=SimpleNamespace( + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda **kwargs: loaded.append(req.cur_kv_len + kwargs["token_indices"].numel()), + ), ), ), ) @@ -288,6 +292,7 @@ def allocate(req, count): lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), deref_pages=lambda page_list: released.extend(page_list), ) + module.cpu_cache_client.cpu_kv_cache_tensor = torch.empty((1, 1)) module.init_sync_group = None module.need_sync_compute_stream = lambda: False monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 8192) diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 2d6819e98d..530b95436a 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -1,69 +1,64 @@ +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources import dataclasses from types import SimpleNamespace import pytest import torch -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager -from lightllm.common.req_manager import Glm5NextReqManager -from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.common.attention_state.allocators import FullAttKVAllocator + +from lightllm.common.attention_state.full_att import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, +) +from lightllm.common.req_manager import ReqManager +from lightllm.models.glm5_next.config import Glm5NextResourceConfig 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 -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture -def model_config(): - return { - "kv_lora_rank": 512, - "num_hidden_layers": 4, - "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], - "index_kpool": 4, - "index_head_dim": 128, - "linear_attn_config": { - "num_heads": 8, - "head_dim": 128, - "short_conv_kernel_size": 4, - "kda_layers": [0, 1, 2], +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("kv_allocator_environment"), +] + + +def test_nope_cache_config_uses_native_mla_width(): + config = Glm5NextResourceConfig.from_model_config( + { + "kv_lora_rank": 512, + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "index_kpool": 4, + "index_head_dim": 128, + "linear_attn_config": { + "num_heads": 8, + "head_dim": 128, + "short_conv_kernel_size": 4, + "kda_layers": [0, 1, 2], + }, }, - } - - -@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) -@pytest.mark.parametrize("page_size", [1, 16]) -def test_nope_cache_layout_and_hold_page(monkeypatch, model_config, kv_type, page_size): - args = StartArgs(tp=4, data_type="bfloat16", llm_kv_type=kv_type) - config = Glm5NextCacheConfig.from_model_config(model_config, args) - nested_config = Glm5NextCacheConfig.from_model_config({"text_config": model_config}, args) - assert nested_config == config - # The allocator reserves an entire HOLD page, including when KV is quantized. - monkeypatch.setattr(Glm5NextMemManager, "_init_linear_att_buffers", lambda self: None) - mem_manager_class = FP8Glm5NextMemManager if kv_type == "fp8kv_dsa" else Glm5NextMemManager - mem = object.__new__(mem_manager_class) - mem.page_size = page_size - mem._init_buffers(32, config.full_att_dtype, 1, config.full_att_head_dim, 1) - assert mem.kv_buffer.shape[1] == 32 + page_size - - assert config.full_att_head_dim == (720 if kv_type == "fp8kv_dsa" else 584) - assert config.full_att_dtype == (torch.uint8 if kv_type == "fp8kv_dsa" else torch.bfloat16) - assert config.conv_state_dtype == torch.bfloat16 + StartArgs(tp=4, data_type="bfloat16"), + ) + # 模型配置保留原始 latent 宽度；packed 行宽由 State 创建入口计算。 + assert config.full_att_head_dim == 512 @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) -@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step, kv_type): - from lightllm.models.glm5_next.model import Glm5NextTpPartModel - +def test_hybrid_checkpoint_restore_and_packed_kv_copy( + monkeypatch, + small_page, + tp_world_size, + mtp_step, +): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) args = StartArgs( tp=tp_world_size, - llm_kv_type=kv_type, data_type="bfloat16", linear_att_hash_page_size=4, linear_att_page_block_num=2, @@ -74,174 +69,182 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t set_unique_server_name(args) get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - use_fp8_kv = kv_type == "fp8kv_dsa" - dtype, head_dim = (torch.uint8, 720) if use_fp8_kv else (torch.bfloat16, 584) - config = Glm5NextCacheConfig( + config = Glm5NextResourceConfig( tp_world_size=tp_world_size, full_att_all_num_kv_heads=1, - full_att_dtype=dtype, - full_att_num_kv_heads=1, - full_att_head_dim=head_dim, + full_att_head_dim=512, global_linear_k_heads=2 * tp_world_size, global_linear_v_heads=2 * tp_world_size, - num_linear_k_heads=2, - num_linear_v_heads=2, head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - layer_types=("linear_attention",) * 3 + ("deepseek_sparse_attention",), - draft_full_att_kv_layer_num=int(mtp_step > 0), + full_attention_interval=4, + all_layer_num=4, + ) + full_layer_ids = config.get_full_att_layer_indices() + if mtp_step: + full_layer_ids += (config.all_layer_num,) + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + logical_config=PackedMLALogicalConfig( + logical_layer_groups=tuple((layer,) for layer in full_layer_ids), + format="packed_bf16", + tp_world_size=tp_world_size, + latent_dim=512, + ) + ), + page_size=args.page_size, ) - full_layers = config.get_full_att_kv_layer_num_with_draft_model() - model = object.__new__(Glm5NextTpPartModel) - model.args = args - model.max_total_token_num = 16 - model.linear_config = config - model.mem_fraction = 0.9 - model._init_mem_manager() - mem = model.mem_manager - req = Glm5NextReqManager(3, 16, mem, config) - att_kv = mem.get_att_input_params(3) - assert att_kv.shape == (17, 1, 584 if use_fp8_kv else 512) - assert att_kv.stride(0) == head_dim - index_bytes = mem.get_indexer_k_buffer(3) + full.allocate_storage(16) + req = ReqManager(3, 16) + from unit_tests.common.att_resources.model_helpers import create_full_attention_roles + + main_full, draft_full = create_full_attention_roles(full, config.all_layer_num) + resources = build_hybrid_resources( + config=config, + request_capacity=4, + mtp_step=mtp_step, + full_state=main_full, + full_allocator=FullAttKVAllocator(full.capacity, full.page_size), + draft_full_state=draft_full, + big_page_capacity=4, + small_page_capacity=2, + ) + full = main_full + linear = resources.main_kda_conv_bf16_ssm_fp32_state + tail = resources.main_compressor_indexer_w4_s4_bf16_state + assert not hasattr(req, "req_to_conv_state") + att_kv = full.get_layer_view(3).kv + assert att_kv.shape == (17, 1, 512) + assert att_kv.stride(0) == 584 + index_bytes = full.get_layer_view(3).indexer_k index_bytes.random_(0, 256) expected_index_bytes = index_bytes.clone() - new_kv = torch.randn(2, 1, 512, dtype=torch.bfloat16, device="cuda") + new_kv = torch.randn( + size=(2, 1, 512), + dtype=torch.bfloat16, + device="cuda", + ) destinations = torch.tensor([5, 9], dtype=torch.int32, device="cuda") - mem.operator.copy_kv_to_mem_manager(3, destinations, new_kv) - if use_fp8_kv: - values = att_kv[..., :448].view(torch.float8_e4m3fn).float().view(17, 1, 7, 64) - scales = att_kv[..., 576:583].view(torch.float8_e8m0fnu).float().unsqueeze(-1) - tail = att_kv[..., 448:576].view(torch.bfloat16).float() - actual = torch.cat(((values * scales).view(17, 1, 448), tail), dim=-1)[destinations] - torch.testing.assert_close(actual, new_kv.float(), atol=0.13, rtol=0.065) - assert torch.equal(tail[destinations], new_kv[..., 448:].float()) - else: - assert torch.equal(att_kv[destinations], new_kv) + full.write_full_att_kv(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 tail.get_layer_view(3).shape == (4, 4 + mtp_step, 256) if mtp_step: - assert req.get_indexer_tail_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 - 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() - if small_page: - req.save_state(source_req, slot, cache) - else: - req.save_big_page_states(torch.tensor([0], dtype=torch.int32, device="cuda"), [source_req], [slot]) + assert resources.draft0_compressor_indexer_w4_s4_bf16_state.get_layer_view(4).shape == (4, 4 + mtp_step, 256) + assert draft_full.get_layer_view(4).kv.shape == att_kv.shape + kind = "small" if small_page else "big" + allocator = resources.bounded_att_small_page_allocator if small_page else resources.bounded_att_big_page_allocator + slot = allocator.alloc(1)[0] + resources.reset_bounded_att_state([0]) + linear.conv_state[:, 0].normal_() + linear.ssm_state[:, 0].normal_() + conv = linear.conv_state[:, 0, ..., :3].clone() + ssm = linear.ssm_state[:, 0].clone() + resources.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=kind, + page_slot_ids=[slot], + ) 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() - 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) + linear.conv_state[:, 0].zero_() + linear.ssm_state[:, 0].zero_() + tail.buffer.normal_() + unchanged = tail.buffer.clone() + resources.restore_bounded_att_runtime_from_pages( + req_indices=[2], + state_token_lens=[8], + page_kind=kind, + page_slot_ids=[slot], + ) 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) - # 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 torch.equal(linear.conv_state[:, 2, ..., :3], conv) + assert torch.equal(linear.ssm_state[:, 2 * (mtp_step + 1)], ssm) + # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 + assert torch.equal(tail.buffer, unchanged) + resources.reset_bounded_att_state([2]) + assert not linear.conv_state[:, 2].any() + assert not linear.ssm_state[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].any() + assert not tail.buffer[:, 2, :, : tail.config.candidate_dim].any() + assert torch.isneginf(tail.buffer[:, 2, :, tail.config.candidate_dim :]).all() + assert torch.equal(tail.buffer[:, [0, 1, 3]], unchanged[:, [0, 1, 3]]) # 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) + packed_bytes = full.storage.view(torch.uint8) packed_bytes[:, 0].random_(0, 256) - mem.operator.copy_mem_to_mem(torch.tensor([0]), torch.tensor([7])) + full.copy_full_att_tokens(torch.tensor([0], device="cuda"), torch.tensor([7], device="cuda")) assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) - assert mem.get_cell_size() == head_dim * dtype.itemsize * full_layers - assert config.get_cpu_cache_full_att_bytes() == mem.get_cell_size() * 8 - - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_kv_buffer_to_cpu_cache, - copy_cpu_cache_to_kv_buffer, - ) - - cpu_pages = torch.zeros((1, config.get_cpu_cache_big_page_bytes()), dtype=torch.uint8, pin_memory=True) - 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 - packed_bytes[:, :8].random_(0, 256) - big.conv_state_cache.buffer[0].normal_() - big.ssm_state_cache.buffer[0].normal_() - expected_kv = packed_bytes[:, :8].clone() - expected_conv = big.conv_state_cache.buffer[0].clone() - expected_ssm = big.ssm_state_cache.buffer[0].clone() - common = dict( - mem_indexes=indexes, - page_indexes=zero, - big_page_buffer_ids=zero, - cpu_kv_conv_state=big.conv_state_cache.buffer, - cpu_kv_ssm_state=big.ssm_state_cache.buffer, - cpu_cache_tensor=cpu_pages, - tp_world_size=tp_world_size, - big_page_token_num=8, - linear_config=config, - ) - # Simulate TP writers to a shared CPU page: MLA/index KV is replicated, - # while each rank must retain its own KDA checkpoint region. - expected_states = [] - for rank in range(tp_world_size): - big.conv_state_cache.buffer[0].copy_(expected_conv + rank) - big.ssm_state_cache.buffer[0].copy_(expected_ssm + rank) - expected_states.append((big.conv_state_cache.buffer[0].clone(), big.ssm_state_cache.buffer[0].clone())) - copy_kv_buffer_to_cpu_cache(page_readies=ready, gpu_kv_full_att_state=mem.kv_buffer, tp_rank=rank, **common) - torch.cuda.synchronize() - for rank, (conv, ssm) in enumerate(expected_states): - packed_bytes[:, :8].zero_() - big.conv_state_cache.buffer[0].zero_() - big.ssm_state_cache.buffer[0].zero_() - copy_cpu_cache_to_kv_buffer(gpu_full_att_kv_state=mem.kv_buffer, tp_rank=rank, **common) - torch.cuda.synchronize() - assert torch.equal(packed_bytes[:, :8], expected_kv) - assert torch.equal(big.conv_state_cache.buffer[0], conv) - assert torch.equal(big.ssm_state_cache.buffer[0], ssm) @pytest.mark.parametrize("gate_type", ["silu", "sigmoid"]) def test_gated_norm_gate_type_and_strided_gate(gate_type): from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import gated_rmsnorm_forward - x = torch.randn(12, 128, device="cuda", dtype=torch.bfloat16) - gate = torch.randn(3, 8, 128, device="cuda", dtype=torch.bfloat16)[:, :4] + x = torch.randn( + size=(12, 128), + device="cuda", + dtype=torch.bfloat16, + ) + gate = torch.randn( + size=(3, 8, 128), + device="cuda", + dtype=torch.bfloat16, + )[:, :4] weight = torch.randn(128, device="cuda", dtype=torch.bfloat16) - actual = gated_rmsnorm_forward(x, weight, None, 1e-5, gate, gate_type=gate_type) + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-5, + z=gate, + gate_type=gate_type, + ) expected = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-5) * weight.float() z = gate.reshape(12, 128).float() expected *= z.sigmoid() if gate_type == "sigmoid" else torch.nn.functional.silu(z) - torch.testing.assert_close(actual, expected.bfloat16(), atol=0.008, rtol=0.008) + torch.testing.assert_close( + actual=actual, + expected=expected.bfloat16(), + atol=0.008, + rtol=0.008, + ) @pytest.mark.parametrize("add_one", [False, True]) def test_clamped_swiglu_preserves_gpt_oss_default(add_one): from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd - x = torch.linspace(-25, 25, 2048, device="cuda", dtype=torch.bfloat16).view(4, 512) - out = torch.empty(4, 256, device="cuda", dtype=torch.bfloat16) + x = torch.linspace( + start=-25, + end=25, + steps=2048, + device="cuda", + dtype=torch.bfloat16, + ).view(4, 512) + out = torch.empty( + size=(4, 256), + device="cuda", + dtype=torch.bfloat16, + ) kwargs = {} if add_one else {"clamp_up_add_one": False} alpha, limit = (1.702, 7.0) if add_one else (1.0, 10.0) - silu_and_mul_fwd(x, out, limit=limit, alpha=alpha, **kwargs) + silu_and_mul_fwd( + input=x, + output=out, + limit=limit, + alpha=alpha, + **kwargs, + ) gate, up = x.float().chunk(2, -1) gate = gate.clamp(max=limit) gate = (gate * torch.sigmoid(alpha * gate)).bfloat16().float() expected = gate * (up.clamp(-limit, limit) + int(add_one)) - torch.testing.assert_close(out, expected.bfloat16(), atol=0.008, rtol=0.008) + torch.testing.assert_close( + actual=out, + expected=expected.bfloat16(), + atol=0.008, + rtol=0.008, + ) diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py index d9a2086ebf..8b4e7aba69 100644 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -21,7 +21,7 @@ fp8_sparse_direct_decode, ) from lightllm.models.glm5_next.triton_kernel.pad_query_heads import _fused_strided_qpad_kernel -from lightllm.models.glm5_next.triton_kernel.destindex_copy_kv_flashmla_fp8 import destindex_copy_kv_flashmla_fp8 +from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( _build_prefill_row_table_kernel, @@ -115,7 +115,7 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid source[1].mul_(1e-5) mem_index = torch.randperm(31, device="cuda")[:11].to(index_dtype) packed_kv = torch.full((31, 1, 720), 71, device="cuda", dtype=torch.uint8) - destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, mem_index, packed_kv) groups = source[..., :448].float().reshape(11, 1, 7, 64) scales = torch.exp2(torch.ceil(torch.log2((groups.abs().amax(-1) / 448).clamp_min(1e-4)))) values = (groups / scales.unsqueeze(-1)).to(torch.float8_e4m3fn) @@ -130,7 +130,7 @@ def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_strid assert (packed_kv[unused] == 71).all() # The packer is graph-safe for both source layouts and mutable destination mapping. - graph, _ = _capture_cuda_graph(lambda: destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv)) + graph, _ = _capture_cuda_graph(lambda: write_fp8_e8m0_packed_kv(source, mem_index, packed_kv)) graph.replay() assert torch.equal(packed_kv[mem_index, :, :448], values.view(torch.uint8).reshape(11, 1, 448)) source.add_(0.125) @@ -173,7 +173,7 @@ def test_fp8_nope_decode_and_cuda_graph(num_heads): packed_kv = torch.zeros(4096, 1, 720, device="cuda", dtype=torch.uint8) source = torch.randn(4096, 1, 512, device="cuda", dtype=torch.bfloat16) mem_index = torch.arange(4096, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, mem_index, packed_kv) q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) # Includes a HOLD row, short tails, and three consecutive MTP positions. lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda", dtype=torch.int32) @@ -201,7 +201,7 @@ def check(actual): check(out) q.normal_() source.normal_() - destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, mem_index, packed_kv) lengths.copy_(lengths.roll(1)) indices.random_(0, 4096) indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) @@ -226,7 +226,7 @@ def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chun source = torch.randn(2 * seq_len, 1, 512, device="cuda", dtype=torch.bfloat16) prefill_kv = source[prefill_rows].clone().add_(0.25) packed_kv = torch.zeros(64, 1, 720, device="cuda", dtype=torch.uint8) - destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) kv_start_offsets = torch.tensor([0] * 3 + [seq_len] * 3, device="cuda", dtype=torch.int32) topk_indices = torch.full((6, 128), -1, device="cuda", dtype=torch.int32) @@ -274,7 +274,7 @@ def check(out): ragged_mem_index.copy_(ragged_mem_index.roll(1)) prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) topk_indices.copy_(topk_indices.roll(1, dims=1)) - destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) graph.replay() check(out) @@ -291,7 +291,7 @@ def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softm prefill_mem_index = ragged_mem_index[prefill_rows].long() prefill_kv = torch.randn(9, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] packed_kv = torch.zeros(2048, 1, 720, device="cuda", dtype=torch.uint8) - destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") topk_indices = torch.randint(num_tokens, (9, 2176), device="cuda", dtype=torch.int32) @@ -335,7 +335,7 @@ def check(actual): prefill_kv.normal_() ragged_mem_index.copy_(ragged_mem_index.roll(1)) prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) - destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + write_fp8_e8m0_packed_kv(source, ragged_mem_index, packed_kv) topk_indices.copy_(topk_indices.roll(1, dims=1)) graph.replay() check(output) @@ -369,7 +369,7 @@ def test_fp8_h64_largeq_native_prefill_permuted_ragged_graph(monkeypatch, num_q) ragged_n = 65536 packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32)[:ragged_n] q = torch.randn(num_q, 64, 512, device="cuda", dtype=torch.bfloat16) starts = torch.where( @@ -437,7 +437,7 @@ def test_fp8_h8_largeq_prefill_keeps_flashattention_path(monkeypatch): num_q = 16 packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) q = torch.randn(num_q, 8, 512, device="cuda", dtype=torch.bfloat16) ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) starts = torch.full((num_q,), 1024, device="cuda", dtype=torch.int32) @@ -477,7 +477,7 @@ def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) # Requests deliberately have different spans; starts are not fixed-stride. @@ -515,7 +515,7 @@ def check(actual): prefill_kv.normal_() ragged.copy_(ragged.roll(1)) prefill_mem.copy_(ragged[current_rows]) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) topk.copy_(topk.roll(1, dims=1)) graph.replay() check(output) @@ -531,7 +531,7 @@ def test_fp8_h64_direct_smallq_state_workspace_reuse_and_fresh_state(): packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() starts = torch.zeros(1, device="cuda", dtype=torch.int32) @@ -587,7 +587,7 @@ def run(active_state=state): q.normal_() source.normal_() prefill_kv.normal_() - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers @@ -600,7 +600,7 @@ def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 4096 @@ -657,7 +657,7 @@ def run(active_state=state): assert fresh_workspace["compact"].data_ptr() == workspace["compact"].data_ptr() source.normal_() - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) prefill_kv.normal_() q.normal_() eager = run() @@ -683,7 +683,7 @@ def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) indices = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) indices[:, :2049] = torch.randint(0, capacity, (1, 2049), device="cuda", dtype=torch.int32) @@ -751,7 +751,7 @@ def test_fp8_decode_infer_state_graph_qpad_and_mutation(num_heads): source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) indices[:, :2049] = torch.randint(0, capacity, (batch, 2049), device="cuda", dtype=torch.int32) @@ -808,7 +808,7 @@ def test_fp8_sparse_direct_decode_actual_entry_edges_and_graph(monkeypatch, batc source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) q = torch.randn(8, batch, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) counts = [2051] if batch == 1 else [0, 2048, 2049, 2050, 2051, 1985, 2051, 2051] @@ -896,7 +896,7 @@ def test_fp8_sparse_direct_decode_general_graph_factory_entry(monkeypatch, num_h source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) # BMM-style B/H stride: (512, batch * 512, 1). q = torch.randn(num_heads, batch, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) assert q.stride() == (512, batch * 512, 1) @@ -979,7 +979,7 @@ def test_fp8_sparse_direct_decode_unknown_bound_fallback_and_strided_q(monkeypat source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) q = torch.randn(1, 8, 512, device="cuda", dtype=torch.bfloat16) late = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) late[:, 2175] = 17 @@ -1024,7 +1024,7 @@ def test_fp8_sparse_direct_decode_strided_inputs_and_packed_row_padding(monkeypa capacity = 4096 source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) indices = torch.full((batch, 1, 4352), -1, device="cuda", dtype=index_dtype)[..., ::2] indices[..., :2051] = torch.randint(0, capacity, (batch, 1, 2051), device="cuda", dtype=index_dtype) @@ -1103,7 +1103,7 @@ def test_fp8_sparse_direct_decode_offsets_beyond_int32(strided_input): topk[..., :7] = torch.randint(0, capacity, (batch, 1, 7), device="cuda", dtype=topk.dtype) packed = torch.empty(capacity, 1, 720, device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) assert can_use_fp8_sparse_direct_decode(q, packed, topk, 512, 2051, is_cuda_graph_capture=True) actual = fp8_sparse_direct_decode(q, packed, topk, 0.0625) expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), topk[:, 0]) @@ -1119,7 +1119,7 @@ def test_fp8_sparse_direct_decode_large_workload_retains_native_graph(monkeypatc torch.manual_seed(53700 + num_heads + batch) packed = torch.empty(4096, 1, 720, device="cuda", dtype=torch.uint8) source = torch.randn(4096, 1, 512, device="cuda", dtype=torch.bfloat16) - destindex_copy_kv_flashmla_fp8(source, torch.arange(4096, device="cuda", dtype=torch.int32), packed) + write_fp8_e8m0_packed_kv(source, torch.arange(4096, device="cuda", dtype=torch.int32), packed) q = torch.randn(batch, num_heads, 512, device="cuda", dtype=torch.bfloat16) topk = torch.full((batch, 1, 2176), -1, device="cuda", dtype=torch.int32) topk[..., :2051] = torch.randint(0, 4096, (batch, 1, 2051), device="cuda", dtype=torch.int32) @@ -1155,7 +1155,7 @@ def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallba packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) mem = torch.arange(capacity, device="cuda", dtype=torch.int32) - destindex_copy_kv_flashmla_fp8(source, mem, packed) + write_fp8_e8m0_packed_kv(source, mem, packed) ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 2048 current_local = torch.full((num_q,), 2048, device="cuda", dtype=torch.int32) @@ -1275,3 +1275,47 @@ def fake_context(*args): ) assert state._nsa_prefill_att(torch.empty((1, heads, 512), dtype=torch.bfloat16), kv, control).shape[1] == heads assert not queries + + +def test_fp8_state_preserves_kv_and_indexer_across_cpu_and_pd(monkeypatch): + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state + + args = SimpleNamespace(cpu_cache_token_page_size=4, pd_kv_page_size=4) + monkeypatch.setattr("lightllm.common.attention_state.manager.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.attention_state.cpu_cache_operator.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.attention_state.pd_operator.get_env_start_args", lambda: args) + monkeypatch.setattr( + "lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_state.get_env_start_args", lambda: args + ) + config = PackedMLALogicalConfig( + logical_layer_groups=((3,),), + latent_dim=512, + format="packed_group64_e8m0", + ) + state = create_packed_mla_state(config=config, page_size=4) + state.allocate_storage(token_capacity=16) + state.storage.zero_() + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + indices = torch.arange(4, dtype=torch.int64, device="cuda") + kv = torch.randn((4, 1, 512), dtype=torch.bfloat16, device="cuda") + state.write_full_att_kv(logical_layer_id=3, token_indices=indices, kv=kv) + expected = torch.zeros((4, 1, 584), dtype=torch.uint8, device="cuda") + write_fp8_e8m0_packed_kv(kv, indices, expected) + torch.testing.assert_close(state.get_layer_view(logical_layer_id=3).kv[:4], expected, rtol=0, atol=0) + state.get_layer_view(logical_layer_id=3).indexer_k[:4].random_(0, 256) + cpu_layout = manager.get_cpu_cache_operator().get_page_layouts()[state.state_name] + cpu_pages = torch.empty((1, manager.get_cpu_cache_operator().get_page_bytes()), dtype=torch.uint8, pin_memory=True) + state.write_to_cpu(token_indices=indices, cpu_page_buffers=cpu_pages, cpu_page_indices=[0], layout=cpu_layout) + state.read_from_cpu(cpu_page_buffers=cpu_pages, cpu_page_indices=[0], token_indices=indices + 4, layout=cpu_layout) + pd_layout = manager.get_pd_operator().get_pd_full_att_page_layouts()[state.state_name] + pd_pages = torch.empty((1, manager.get_pd_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(token_indices=indices, pd_page_buffers=pd_pages, pd_page_index=0, layout=pd_layout) + state.read_from_pd_page(pd_page_buffers=pd_pages, token_indices=indices + 8, pd_page_index=0, layout=pd_layout) + torch.cuda.synchronize() + for begin in (4, 8): + actual = state.get_layer_view(logical_layer_id=3) + torch.testing.assert_close(actual.kv[begin : begin + 4], actual.kv[:4], rtol=0, atol=0) + torch.testing.assert_close(actual.indexer_k[begin : begin + 4], actual.indexer_k[:4], rtol=0, atol=0) diff --git a/unit_tests/models/glm5_next/test_indexer_topk.py b/unit_tests/models/glm5_next/test_indexer_topk.py index fce877448c..76646c5224 100644 --- a/unit_tests/models/glm5_next/test_indexer_topk.py +++ b/unit_tests/models/glm5_next/test_indexer_topk.py @@ -12,7 +12,7 @@ @pytest.fixture def indexer(): return Glm5NextNsaInfer( - 0, {"index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5}, 1 + 0, {"index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5, "index_kpool": 4}, 1 ) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 37485746ba..628547a591 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -35,16 +35,6 @@ def setup(): triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) -def _reference_kda(q, k, v, gate, beta, a, bias, state): - q = q.float() * torch.rsqrt(q.float().square().sum(-1, keepdim=True) + 1e-6) / q.shape[-1] ** 0.5 - k = k.float() * torch.rsqrt(k.float().square().sum(-1, keepdim=True) + 1e-6) - decay = (-5 * torch.sigmoid(a.exp()[:, None] * (gate.float() + bias))).exp() - state = state * decay[..., None] - delta = (v.float() - torch.einsum("hkv,hk->hv", state, k)) * beta.float().sigmoid()[:, None] - state = state + k[..., None] * delta[:, None, :] - return torch.einsum("hkv,hk->hv", state, q), state - - @pytest.mark.parametrize("tokens", [1, 3, 65, 129]) def test_kda_chunk_and_decode_match_recurrence(tokens): heads, dim = 2, 128 @@ -594,18 +584,26 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): storage = torch.zeros(tokens + 9, 1, 584, device="cuda", dtype=torch.bfloat16) tail = torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16) ragged = torch.randperm(tokens + 9, device="cuda", dtype=torch.int32)[:tokens] - manager = SimpleNamespace( - get_indexer_k_buffer=lambda _: storage.view(torch.uint8)[:, :, -132:], + full_att_state = SimpleNamespace( + get_layer_view=lambda logical_layer_id: SimpleNamespace(indexer_k=storage.view(torch.uint8)[:, :, -132:]), ) infer = SimpleNamespace( - mem_manager=manager, + att_state_role="main", + att_state_manager=SimpleNamespace( + get_state=lambda state_type_filter, **kwargs: ( + SimpleNamespace(get_layer_view=lambda logical_layer_id: tail) + if "compressor_indexer" in state_type_filter + else full_att_state + ), + ), + is_draft_model=False, is_prefill=True, b_mtp_index=torch.zeros(1, device="cuda", dtype=torch.int32), mem_index=ragged, # 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]), 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), @@ -617,7 +615,9 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): ragged_mem_index=ragged, ) indexer = Glm5NextNsaInfer( - 0, {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim, "rms_norm_eps": 1e-5}, 1 + 0, + {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim, "rms_norm_eps": 1e-5, "index_kpool": 4}, + 1, ) _, full = indexer._get_indices(hidden, hidden, infer, state, weights) assert full[-1, 2048].item() == tokens - 1 # Always-selected incomplete tail. @@ -625,11 +625,11 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): assert full[-1, :2049].unique().numel() == 2049 assert (full[-1, :2049] < tokens).all() # Restore an aligned prefix using only token KV, then decode across the - # next pool boundary. The runtime tail starts empty, as on a page restore. + # next pool boundary. Stale window data is irrelevant at a complete group boundary. prefix = 2060 storage = storage.clone() storage[ragged[prefix:].long()] = 0 - tail.zero_() + tail.fill_(19) infer.is_prefill = False infer.b1_cu_q_seq_len[1] = 1 infer.max_q_seq_len = 1 @@ -653,3 +653,13 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): scores = (logits * q_scale * (heads ** -0.5 * dim ** -0.5)).sum(0) expected_groups = set(scores.topk(512).indices.tolist()) assert set((decoded[0, :2048:4] // 4).tolist()) == expected_groups + + +def _reference_kda(q, k, v, gate, beta, a, bias, state): + q = q.float() * torch.rsqrt(q.float().square().sum(-1, keepdim=True) + 1e-6) / q.shape[-1] ** 0.5 + k = k.float() * torch.rsqrt(k.float().square().sum(-1, keepdim=True) + 1e-6) + decay = (-5 * torch.sigmoid(a.exp()[:, None] * (gate.float() + bias))).exp() + state = state * decay[..., None] + delta = (v.float() - torch.einsum("hkv,hk->hv", state, k)) * beta.float().sigmoid()[:, None] + state = state + k[..., None] * delta[:, None, :] + return torch.einsum("hkv,hk->hv", state, q), state diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 08f8579174..8ac9c2d473 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -61,7 +61,7 @@ def config(): "index_head_dim": 128, "index_topk": 2048, "index_kpool": 4, - "linear_attn_config": {"num_heads": 4, "head_dim": 128, "short_conv_kernel_size": 4}, + "linear_attn_config": {"num_heads": 4, "head_dim": 128, "short_conv_kernel_size": 4, "kda_layers": [0, 1, 2]}, "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], "mhc": False, "n_routed_experts": 4, @@ -170,7 +170,10 @@ def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (indices, indices)) state = InferStateInfo() state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} - state.mem_manager = SimpleNamespace(get_att_input_params=lambda layer_index: kv) + state.att_state_role = "main" + state.att_state_manager = SimpleNamespace( + get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda logical_layer_id: SimpleNamespace(kv=kv)) + ) state.prefill_att_state = Glm5NextSparsePrefillState() state.b1_cu_q_seq_len = torch.arange(4, device="cuda", dtype=torch.int32) state.max_q_seq_len = 1 @@ -375,21 +378,23 @@ def test_glm_rejects_tpsp(config, tp_world_size): model._verify_params() -def test_native_drafts_share_caches_but_keep_config_and_layer_indices(config): +def test_native_drafts_share_weights_and_prepare_their_own_state(config): main = object.__new__(Glm5NextTpPartModel) main.config = dict(config, mhc=True) main.tp_world_size_ = 1 - main.args, main.load_way = StartArgs(), "HF" + main.args, main.load_way = StartArgs(data_type="bfloat16", mtp_mode="vanilla_with_att", mtp_step=2), "HF" main._verify_params() main._init_some_value() - main.layers_infer = [object() for _ in range(config["num_hidden_layers"])] + main.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(config["num_hidden_layers"])] main.pre_post_weight = SimpleNamespace(wte_weight_=object(), lm_head_weight_=object(), final_norm_weight_=object()) - main.req_manager, main.mem_manager, main.linear_config = object(), object(), object() + main.req_manager = object() + main.max_req_num, main.tp_rank_ = 0, 0 drafts = [] for step in range(2): draft = object.__new__(Glm5NextMTPModel) draft.main_model, draft.mtp_previous_draft_models = main, list(drafts) - draft.tp_world_size_, draft.data_type = 1, torch.bfloat16 + draft.tp_world_size_, draft.tp_rank_, draft.data_type = 1, 0, torch.bfloat16 + draft.args = main.args draft.load_way = "HF" draft._init_config() draft._verify_params() @@ -397,9 +402,8 @@ def test_native_drafts_share_caches_but_keep_config_and_layer_indices(config): draft._init_weights() draft._init_infer_layer() draft._init_some_value() - draft._init_req_manager() - draft._init_mem_manager() - draft._init_att_backend1() + draft.bind_req_manager(req_manager=main.req_manager) + draft.max_req_num = main.max_req_num assert main.config["mhc"] and not draft.config["mhc"] assert draft.config["layer_types"] is not main.config["layer_types"] assert draft.layers_num == 1 and draft.head_dim_ == 512 @@ -408,10 +412,31 @@ def test_native_drafts_share_caches_but_keep_config_and_layer_indices(config): assert draft.pre_post_weight.wte_weight_ is main.pre_post_weight.wte_weight_ assert draft.pre_post_weight.lm_head_weight_ is main.pre_post_weight.lm_head_weight_ assert draft.pre_post_weight.main_norm_weight_ is main.pre_post_weight.final_norm_weight_ - assert draft.req_manager is main.req_manager and draft.mem_manager is main.mem_manager - assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None + assert draft.req_manager is main.req_manager drafts.append(draft) + assert all(not hasattr(model, "att_state_manager") for model in (main, *drafts)) + assert all(not hasattr(model, "linear_config") for model in (main, *drafts)) + from contextlib import nullcontext + + for model in (main, *drafts): + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + main.prepare_attention_states() + assert main.att_state_manager.main_kda_conv_bf16_ssm_fp32_state.conv_state.is_cuda + assert main.att_state_manager.main_compressor_indexer_w4_s4_bf16_state.buffer.is_cuda + for step, draft in enumerate(drafts): + draft.prepare_attention_states() + assert draft.att_state_manager is main.att_state_manager + draft._init_att_backend1() + state = getattr(draft.att_state_manager, f"draft{step}_mla_packed_bf16_state") + assert state is not main.att_state_manager.main_mla_packed_bf16_state + assert not hasattr(draft, "full_att_state") + assert state.storage is None + assert state.config.logical_layer_groups == ((4 + step,),) + assert state.state_name == f"draft{step}_mla_packed_bf16_state" + assert draft.prefill_att_backend1 is None and draft.decode_att_backend1 is None + assert getattr(draft.att_state_manager, f"draft{step}_compressor_indexer_w4_s4_bf16_state").buffer.is_cuda + # NoPE state initialization must not require a model-owned RoPE cache. for model in (main, *drafts): state = model.infer_state_class() diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 00d4e326f0..82e029c5dd 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -2,6 +2,8 @@ import pytest import torch +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, KDAState from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -124,7 +126,8 @@ def prefill(values, destination, tail): @pytest.mark.parametrize("dynamic", [False, True]) -def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): +@pytest.mark.parametrize("role", ["main", "draft0"]) +def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): from types import SimpleNamespace from lightllm.common.basemodel.attention.base_att import AttControl @@ -191,14 +194,32 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic): uses_dynamic_spec_verify_layout=lambda: 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) + resources = AttentionStateManager(request_capacity=4) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=heads, + global_v_heads=heads, + key_dim=dim, + value_dim=dim, + conv_width=3, + conv_dtype=conv.dtype, + ssm_dtype=ssm.dtype, ) + for model_role in ("main", "draft0"): + linear = KDAState(config, request_capacity=4, mtp_step=2, device="meta") + linear.mtp_state_index = accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets) + if model_role == role: + linear.conv_state, linear.ssm_state = conv.unsqueeze(0), ssm.unsqueeze(0) + resources.register_state(model_role, linear) + resources.finish_state_registration() + manager = SimpleNamespace(HOLD_REQUEST_ID=3) 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.att_state_manager = resources + infer.att_state_role = role 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 fab579aae1..ff93312516 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -1,30 +1,42 @@ +from unit_tests.common.att_resources.model_helpers import window_states, build_hybrid_resources +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + import dataclasses +from types import SimpleNamespace import pytest import torch -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager, MemoryManager -from lightllm.common.req_manager import Glm5NextReqManager -from lightllm.common.state_cache_manager import Glm5NextCacheConfig +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.model_helpers import create_full_attention_roles + +from lightllm.common.attention_state.full_att import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, +) +from lightllm.server.router.model_infer.mode_backend.pd.resource_io import allocate_transfer_buffer, copy_page +from lightllm.models.glm5_next.config import Glm5NextResourceConfig from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools 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 -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("kv_allocator_environment"), +] -@pytest.fixture(params=[(0, False), (2, False), (0, True), (2, True)]) +@pytest.fixture(params=[0, 2]) def make_mems(monkeypatch, request): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) - mtp_step, use_fp8_kv = request.param - dtype, head_dim = (torch.uint8, 720) if use_fp8_kv else (torch.bfloat16, 584) - mem_manager_class = FP8Glm5NextMemManager if use_fp8_kv else Glm5NextMemManager + mtp_step = request.param args = StartArgs( data_type="bfloat16", - llm_kv_type="fp8kv_dsa" if use_fp8_kv else "None", + pd_kv_page_size=2048, linear_att_hash_page_size=4, linear_att_page_block_num=2, mtp_step=mtp_step, @@ -34,41 +46,64 @@ 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( + def create(tp, target_mtp_step=None): + local_mtp_step = mtp_step if target_mtp_step is None else target_mtp_step + config = Glm5NextResourceConfig( tp_world_size=tp, full_att_all_num_kv_heads=1, - full_att_dtype=dtype, - full_att_num_kv_heads=1, - full_att_head_dim=head_dim, + full_att_head_dim=512, global_linear_k_heads=8, global_linear_v_heads=8, - num_linear_k_heads=8 // tp, - num_linear_v_heads=8 // tp, head_linear_k_dim=128, head_linear_v_dim=128, conv_kernel_size=4, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, - layer_types=(("linear_attention",) * 3 + ("deepseek_sparse_attention",)) * 2, - draft_full_att_kv_layer_num=int(mtp_step > 0), + full_attention_interval=4, + all_layer_num=8, ) + full_layer_ids = config.get_full_att_layer_indices() + if mtp_step: + full_layer_ids += (config.all_layer_num,) mems = [] - for _ in range(tp): - mem = mem_manager_class(32, dtype, 1, head_dim, 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 - mems.append(mem) - mems[0].alloc_paged_kv_move_buffer(1, 4096) + for rank in range(tp): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp)) + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + logical_config=PackedMLALogicalConfig( + logical_layer_groups=tuple((layer,) for layer in full_layer_ids), + format="packed_bf16", + tp_rank=rank, + tp_world_size=tp, + latent_dim=512, + ) + ), + page_size=args.page_size, + ) + full.allocate_storage(32) + main, draft = create_full_attention_roles(full, config.all_layer_num) + resources = build_hybrid_resources( + config=config, + request_capacity=4, + mtp_step=local_mtp_step, + tp_rank=rank, + full_state=main, + full_allocator=FullAttKVAllocator(full.capacity, full.page_size), + draft_full_state=draft, + big_page_capacity=6, + small_page_capacity=2, + ) + mems.append( + SimpleNamespace( + full=main, + full_states=tuple(state for state in (main, draft) if state is not None), + linear=resources.main_kda_conv_bf16_ssm_fp32_state, + tail=resources.main_compressor_indexer_w4_s4_bf16_state, + tails=window_states(resources), + resources=resources, + ) + ) return mems return create @@ -76,71 +111,157 @@ 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) +@pytest.mark.parametrize("decode_mtp_step", [0, 1, 5]) +def test_pd_kv_and_runtime_state_roundtrip( + make_mems, + prefill_tp, + decode_tp, + remainder, + decode_mtp_step, +): + source, dest = make_mems(prefill_tp), make_mems(decode_tp, target_mtp_step=decode_mtp_step) mtp_size = get_env_start_args().mtp_step + 1 - layers = source[0].layer_num + dst_mtp_size = decode_mtp_step + 1 + layers = sum(len(state.config.logical_layer_groups) for state in source[0].full_states) src_req, dst_req = 0, 2 length = 8 + remainder src_indexes = torch.randperm(32)[:length].tolist() dst_indexes = torch.randperm(32)[:length].tolist() - token_bytes = source[0].head_dim * source[0].dtype.itemsize - packed = torch.randint(0, 256, (layers, length, 1, token_bytes), dtype=torch.uint8, device="cuda") - conv = torch.randn(6, 3, 8, 128, 3, dtype=torch.bfloat16, device="cuda") - ssm = torch.randn(6, 8, 128, 128, dtype=torch.float32, device="cuda") + packed = torch.randint( + low=0, + high=256, + size=(layers, length, 1, 1168), + dtype=torch.uint8, + device="cuda", + ) + conv = torch.randn( + size=(6, 3, 8, 128, 3), + dtype=torch.bfloat16, + device="cuda", + ) + ssm = torch.randn( + size=(6, 8, 128, 128), + dtype=torch.float32, + device="cuda", + ) # 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") + tail = torch.randn( + size=(layers, 3 + mtp_size, 256), + dtype=torch.bfloat16, + device="cuda", + ) for rank, mem in enumerate(source): 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) + offset = 0 + for state in mem.full_states: + end = offset + len(state.config.logical_layer_groups) + state.storage.view(torch.uint8)[:, src_indexes] = packed[offset:end] + offset = end + mem.linear.conv_state[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) + mem.linear.ssm_state[:, src_req * mtp_size].copy_(ssm[:, heads]) + offset = 0 + for state in mem.tails: + end = offset + len(state.config.logical_layer_ids) + state.buffer[:, src_req].copy_(tail[offset:end]) + offset = end untouched = [] + original_tails = [] for mem in dest: - 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_() + for state in mem.full_states: + state.storage.zero_() + mem.linear.conv_state.normal_() + mem.linear.ssm_state.normal_() + for state in mem.tails: + state.buffer.normal_() + original_tails.append(torch.cat([state.buffer[:, dst_req] for state in mem.tails]).clone()) untouched.append( [ - x.buffer[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() + x[:, 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), + (mem.linear.conv_state, 1), + (mem.linear.ssm_state, dst_mtp_size), + *((state.buffer, 1) for state in mem.tails), ) ] ) - source[0].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) + source_resources = [mem.resources for mem in source] + dest_resources = [mem.resources for mem in dest] + source_buffer = allocate_transfer_buffer(source_resources[0], 1, device="cuda") + dest_buffer = allocate_transfer_buffer(dest_resources[0], 1, device="cuda") + assert source_buffer.shape == dest_buffer.shape + assert ( + source_resources[0].get_pd_operator().get_pd_bounded_att_page_layouts().keys() + == dest_resources[0].get_pd_operator().get_pd_bounded_att_page_layouts().keys() + ) + for kind in ("kv", "att_state"): + copy_page( + node_att_managers=source_resources, + pd_page_buffers=source_buffer, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=prefill_tp, + token_indices=torch.tensor(src_indexes if kind == "kv" else [], dtype=torch.long, device="cuda"), + task_kind=kind, + req_idx=src_req, + state_token_len=length, + write_to_pd_page=True, + ) + dest_buffer.copy_(source_buffer) + copy_page( + node_att_managers=dest_resources, + pd_page_buffers=dest_buffer, + pd_page_index=0, + dp_rank_in_node=0, + tp_world_size=decode_tp, + token_indices=torch.tensor(dst_indexes if kind == "kv" else [], dtype=torch.long, device="cuda"), + task_kind=kind, + req_idx=dst_req, + state_token_len=length, + write_to_pd_page=False, + ) torch.cuda.synchronize() for rank, mem in enumerate(dest): 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) + for state, source_state in zip(mem.full_states, source[0].full_states): + expected = full_state_test_fields(source_state, {"kv": source_state.storage}) + for name, field in full_state_test_fields(state, {"kv": state.storage}).items(): + source_field = expected[name] + assert torch.equal( + field.tensor[:, dst_indexes].contiguous().view(torch.uint8), + source_field.tensor[:, src_indexes].contiguous().view(torch.uint8), + ) + assert torch.equal(mem.linear.conv_state[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) + assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_mtp_size], ssm[:, heads]) + restored_tail = torch.cat([state.buffer[:, dst_req] for state in mem.tails]) + expected_tail = original_tails[rank].clone() + for position in range(length - 4, length): + expected_tail[:, position % (4 + decode_mtp_step)] = tail[:, position % tail.shape[1]] + assert torch.equal(restored_tail, expected_tail) + if mem.linear.mtp_state_index is not None: + assert mem.linear.mtp_state_index[dst_req] == 0 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)), + ( + (mem.linear.conv_state, 1), + (mem.linear.ssm_state, dst_mtp_size), + *((state.buffer, 1) for state in mem.tails), + ), untouched[rank], ): indexes = torch.tensor([0, 1, 3], device="cuda") * stride - assert torch.equal(state.buffer[:, indexes], expected) + assert torch.equal(state[:, indexes], expected) # Continue through the next pool boundary on P and on the restored D state. # Different physical KV/request slots must yield identical pooled FP8 bytes. count = 4 - remainder - raw = torch.randn(count, 256, dtype=torch.bfloat16, device="cuda") + raw = torch.randn( + size=(count, 256), + dtype=torch.bfloat16, + device="cuda", + ) ape = torch.randn(4, 128, device="cuda") tensor = lambda values: torch.tensor(values, dtype=torch.int32, device="cuda") outputs = [] @@ -149,8 +270,8 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem all_indexes = indexes + new_indexes compress_pools( raw=raw, - tail=mem.req_to_indexer_tail.buffer[0], - packed_buffer=mem.get_indexer_k_buffer(3), + tail=mem.tail.buffer[0], + packed_buffer=mem.full.get_layer_view(3).indexer_k, ape=ape, lengths=tensor(range(length + 1, length + count + 1)), starts=tensor([0] * count), @@ -160,16 +281,147 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem seq_lens=tensor([length + count]), max_q_len=count, ) - outputs.append(mem.get_indexer_k_buffer(3)[new_indexes[-1]].clone()) + outputs.append(mem.full.get_layer_view(3).indexer_k[new_indexes[-1]].clone()) assert torch.equal(*outputs) -def test_pd_page_capacity_includes_tail_without_resizing(make_mems): +def test_pd_page_capacity_includes_tail_without_resizing(make_mems, monkeypatch, caplog): + monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) mem = make_mems(1)[0] - helper = mem.att_state_page_helper - # 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") - 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) - assert shape == (2, 2048, mem.layer_num, 1, mem.head_dim) + full_bytes = max( + layout.byte_end for layout in mem.resources.get_pd_operator().get_pd_full_att_page_layouts().values() + ) + runtime_bytes = max( + layout.byte_end for layout in mem.resources.get_pd_operator().get_pd_bounded_att_page_layouts().values() + ) + assert runtime_bytes > full_bytes + assert mem.resources.get_pd_operator().get_page_bytes() == runtime_bytes + assert "smaller than the complete attention runtime" in caplog.text + + +@pytest.mark.parametrize("source_step,target_step", [(0, 1), (0, 5), (3, 2), (5, 1)]) +@pytest.mark.parametrize("prefix", [1, 3, 4, 5, 11, 12]) +def test_pd_prefill_tail_continues_mtp_after_rejection(source_step, target_step, prefix): + """真实 prefill 导出后，在不同大小的 ring 上验证、拒绝回退并再次验证。""" + from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state + from lightllm.common.attention_state.types import StateTransferLayout + + torch.manual_seed(53) + config = WindowCompressorStateConfig( + logical_layer_ids=(0,), + window_size=4, + compress_stride=4, + candidate_dim=128, + target="indexer", + dtype=torch.bfloat16, + ) + source = create_window_compressor_state( + config=config, + request_capacity=2, + mtp_step=source_step, + device="cuda", + ) + target = create_window_compressor_state( + config=config, + request_capacity=2, + mtp_step=target_step, + device="cuda", + ) + raw = torch.randn( + size=(64, 256), + device="cuda", + dtype=torch.bfloat16, + ) + ape = torch.randn( + size=(4, 128), + device="cuda", + dtype=torch.bfloat16, + ) + packed = torch.zeros( + size=(64, 132), + device="cuda", + dtype=torch.uint8, + ) + ragged = torch.arange(64, device="cuda", dtype=torch.int32) + + def prefill(values, destination, tail): + count = len(values) + compress_pools( + raw=values, + tail=tail, + packed_buffer=destination, + ape=ape, + lengths=torch.arange( + start=1, + end=count + 1, + device="cuda", + dtype=torch.int32, + ), + starts=torch.zeros(count, device="cuda", dtype=torch.int32), + ragged=ragged, + req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), + cu_q_lens=torch.tensor([0, count], device="cuda", dtype=torch.int32), + seq_lens=torch.tensor([count], device="cuda", dtype=torch.int32), + max_q_len=count, + ) + + prefill(raw[:prefix], packed, source.get_layer_view(0)) + regions = source.get_pd_transfer_regions() + assert regions == target.get_pd_transfer_regions() + layout = StateTransferLayout( + state_name=source.state_name, + tokens_per_page=None, + byte_start=0, + byte_end=regions[-1].byte_end, + regions=regions, + ) + pages = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + source.write_to_pd_page( + req_idx=0, + pd_page_buffers=pages, + pd_page_index=0, + state_token_len=prefix, + layout=layout, + ) + target.buffer.fill_(17) + target.read_from_pd_page( + pd_page_buffers=pages, + req_idx=0, + pd_page_index=0, + state_token_len=prefix, + layout=layout, + ) + assert (target.buffer[:, 1] == 17).all() + width = target_step + 1 + # 第一次验证后回访已接受位置；新增候选历史必须由 D 的本地 ring 正确保留。 + for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: + raw[start : start + count] = torch.randn( + size=(count, 256), + device="cuda", + dtype=torch.bfloat16, + ) + lengths = torch.arange( + start=start + 1, + end=start + count + 1, + device="cuda", + dtype=torch.int32, + ) + compress_pools( + raw=raw[start : start + count], + tail=target.get_layer_view(0), + packed_buffer=packed, + ape=ape, + lengths=lengths, + starts=torch.zeros(count, device="cuda", dtype=torch.int32), + ragged=ragged, + req_idx=torch.zeros(count, device="cuda", dtype=torch.int32), + cu_q_lens=torch.arange(count + 1, device="cuda", dtype=torch.int32), + seq_lens=lengths, + max_q_len=1, + mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), + ) + expected = torch.zeros_like(packed) + prefill(raw[: start + count], expected, torch.zeros_like(source.get_layer_view(0))) + closing = torch.arange(start, start + count, device="cuda") + closing = closing[(closing + 1) % 4 == 0] + assert torch.equal(packed[closing], expected[closing]) diff --git a/unit_tests/models/sensenova_loop/test_loop.py b/unit_tests/models/sensenova_loop/test_loop.py index 9444485ab6..ea3311e66a 100644 --- a/unit_tests/models/sensenova_loop/test_loop.py +++ b/unit_tests/models/sensenova_loop/test_loop.py @@ -8,7 +8,7 @@ from lightllm.models.registry import ModelRegistry from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types -from lightllm.common.state_cache_manager.linear_att import LinearAttCacheConfig +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -35,10 +35,14 @@ def _config(layers=32, start=4, end=28, count=3): def test_cache_slots_follow_execution_depth_including_unaligned_loop(monkeypatch, start, end, count): config = _config(start=start, end=end, count=count) args = StartArgs(model_dir="unused", tp=4, dp=1, data_type="bfloat16", mtp_mode="eagle_with_att", mtp_step=3) - monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_env_start_args", lambda: args) - monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_added_mtp_kv_layer_num", lambda: 1) + monkeypatch.setattr("lightllm.models.qwen3next.config.get_env_start_args", lambda: args) monkeypatch.setattr("transformers.configuration_utils.PretrainedConfig.get_config_dict", lambda _: (config, {})) - cache = LinearAttCacheConfig.load_from_args() + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + + normalized = SenseNovaLoopTpPartModel.normalize_model_config(config) + cache = Qwen3NextResourceConfig.from_model_config(normalized, args) + full_layers = cache.get_full_att_layer_indices() + linear_layers = tuple(i for i in range(cache.all_layer_num) if i not in full_layers) indices = get_loop_layer_indices(config["loop_settings"], 32) kinds = get_loop_layer_types(config["loop_settings"], config["text_config"]) assert len(kinds) == 32 + (end - start) * (count - 1) @@ -46,18 +50,18 @@ def test_cache_slots_follow_execution_depth_including_unaligned_loop(monkeypatch slots = {} for execution, (physical, kind) in enumerate(zip(indices, kinds)): if kind == "full_attention": - slot = cache.get_full_att_kv_layer_index(execution) + slot = full_layers.index(execution) assert slot == full full += 1 else: - slot = cache.get_linear_state_layer_index(execution) + slot = linear_layers.index(execution) assert slot == linear linear += 1 assert slot not in slots.setdefault(physical, set()) slots[physical].add(slot) - assert cache.get_main_model_full_att_layer_num() == full - assert cache.linear_layer_num == linear - assert cache.get_full_att_kv_layer_index(len(indices)) == full # Independent MTP draft. + assert len(full_layers) == full + assert len(linear_layers) == linear + assert len(set(full_layers) | set(linear_layers)) == len(indices) if (start, end, count) == (4, 28, 3): assert (cache.all_layer_num, full, linear) == (80, 20, 60) if (start, end, count) == (4, 28, 2): @@ -122,7 +126,7 @@ def test_mhc_registry_and_embedded_mtp_reuse_existing_qwen_draft(tmp_path, strea draft.weight_dir_, draft.finetune_config, draft.tp_world_size_ = str(tmp_path), None, 4 draft._init_config() assert draft.config["n_layer"] == draft.config["full_attention_interval"] == 1 - assert draft.config["layer_types"] == config["text_config"]["layer_types"] + assert draft.config["layer_types"] == ["full_attention"] def test_loop_tokenizer_loads_without_hf_loop_autoconfig(tmp_path): @@ -158,21 +162,27 @@ def test_model_and_worker_use_the_same_loop_cache_layout(monkeypatch, custom_lay if custom_layout: config["text_config"]["layer_types"] = ["full_attention", "linear_attention"] * 16 original = deepcopy(config) - args = StartArgs(tp=4, dp=1, data_type="bfloat16") - monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_added_mtp_kv_layer_num", lambda: 1) - monkeypatch.setattr("lightllm.common.state_cache_manager.linear_att.get_env_start_args", lambda: args) + args = StartArgs(tp=4, dp=1, data_type="bfloat16", linear_att_hash_page_size=256) + monkeypatch.setattr("lightllm.models.qwen3next.config.get_env_start_args", lambda: args) monkeypatch.setattr("transformers.configuration_utils.PretrainedConfig.get_config_dict", lambda _: (config, {})) - monkeypatch.setattr( - "lightllm.models.qwen3next.model.ReqManagerForMamba", - lambda *args, linear_config: SimpleNamespace(linear_config=linear_config), + config["text_config"]["num_attention_heads"] = 16 + original = deepcopy(config) + normalized = SenseNovaLoopTpPartModel.normalize_model_config(config) + raw_cache = Qwen3NextResourceConfig.from_model_config(normalized, args) + logical_ids = tuple(range(normalized["num_hidden_layers"])) + (full,) = SenseNovaLoopTpPartModel.create_full_att_states( + model_config=normalized, + args=args, + model_logical_layer_ids=logical_ids, ) - raw_cache = LinearAttCacheConfig.load_from_args() - model = object.__new__(SenseNovaLoopTpPartModel) - model.args = args - model.max_req_num, model.batch_max_tokens, model.max_seq_length = 2, 8, 16 - model._init_req_manager() - assert dataclasses.asdict(raw_cache) == dataclasses.asdict(model.linear_config) - assert model.req_manager.linear_config is model.linear_config + (bounded,) = SenseNovaLoopTpPartModel.create_bounded_att_states( + model_config=normalized, + args=args, + model_logical_layer_ids=logical_ids, + ) + full_layers = raw_cache.get_full_att_layer_indices() + assert full.config.logical_layer_groups == tuple((i,) for i in full_layers) + assert bounded.config.logical_layer_ids == tuple(i for i in logical_ids if i not in full_layers) assert config == original diff --git a/unit_tests/models/test_qwen3_dspark_model_output.py b/unit_tests/models/test_qwen3_dspark_model_output.py index 8e2085fa00..8f8c191b2c 100644 --- a/unit_tests/models/test_qwen3_dspark_model_output.py +++ b/unit_tests/models/test_qwen3_dspark_model_output.py @@ -65,7 +65,7 @@ def test_parallel_block_decode_commits_target_hiddens_directly(model_class): model = model_class.__new__(model_class) model._cos_cached = torch.arange(24).view(6, 4) model._sin_cached = model._cos_cached + 100 - model.mem_manager = object() + model.att_state_manager = object() model.pre_post_weight = object() target_hiddens = torch.arange(6, dtype=torch.float32).view(2, 3) @@ -104,7 +104,9 @@ def context_forward(self, hidden, infer_state, layer_weight): assert isinstance(output, ModelOutput) assert output.logits.shape == (2, 1) infer_state = observed_states[0] - assert infer_state.mem_manager is model.mem_manager + assert not hasattr(infer_state, "full_att_state") + assert infer_state.att_state_manager is model.att_state_manager + assert infer_state.is_draft_model assert infer_state.mem_index is mem_indexes assert torch.equal(infer_state.position_cos, model._cos_cached[[2, 4]]) assert torch.equal(infer_state.position_sin, model._sin_cached[[2, 4]]) @@ -297,49 +299,39 @@ def test_fixed_dspark_does_not_require_confidence_head(monkeypatch): model._verify_params() -def test_qwen35_dspark_adapter_uses_current_checkpoint_rope_layout(monkeypatch): - def init_dspark_config(self): - self.config = { - "dflash_config": {"mask_token_id": 1}, - "rope_parameters": { - "rope_theta": 1_000_000, - "factor": 32.0, - "original_max_position_embeddings": 8192, - "rope_type": "yarn", - }, - } - - monkeypatch.setattr(Qwen3DSparkModel, "_init_config", init_dspark_config) - model = Qwen3_5DSparkModel.__new__(Qwen3_5DSparkModel) - - model._init_config() - - assert model.config["rope_scaling"] == model.config["rope_parameters"] - assert model.config["rope_theta"] == 1_000_000 - assert model.config["partial_rotary_factor"] == 1.0 - assert model.config["mask_token_id"] == 1 - - -def test_qwen35_dspark_adapter_preserves_custom_partial_rotary_layout(monkeypatch): - def init_dspark_config(self): - self.config = { - "dflash_config": {"mask_token_id": 1}, - "rope_parameters": { - "mrope_interleaved": True, - "mrope_section": [11, 11, 10], - "partial_rotary_factor": 0.25, - "rope_theta": 10_000_000, - "rope_type": "default", - }, - } - - monkeypatch.setattr(Qwen3DSparkModel, "_init_config", init_dspark_config) - model = Qwen3_5DSparkModel.__new__(Qwen3_5DSparkModel) - - model._init_config() - - assert model.config["partial_rotary_factor"] == 0.25 - assert model.config["rope_scaling"] == model.config["rope_parameters"] +def test_qwen35_dspark_adapter_uses_current_checkpoint_rope_layout(): + raw_config = { + "dflash_config": {"mask_token_id": 1}, + "rope_parameters": { + "rope_theta": 1_000_000, + "factor": 32.0, + "original_max_position_embeddings": 8192, + "rope_type": "yarn", + }, + } + config = Qwen3_5DSparkModel.normalize_model_config(raw_config) + + assert config["rope_scaling"] == config["rope_parameters"] + assert config["rope_theta"] == 1_000_000 + assert config["partial_rotary_factor"] == 1.0 + assert config["mask_token_id"] == 1 + + +def test_qwen35_dspark_adapter_preserves_custom_partial_rotary_layout(): + raw_config = { + "dflash_config": {"mask_token_id": 1}, + "rope_parameters": { + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + "partial_rotary_factor": 0.25, + "rope_theta": 10_000_000, + "rope_type": "default", + }, + } + config = Qwen3_5DSparkModel.normalize_model_config(raw_config) + + assert config["partial_rotary_factor"] == 0.25 + assert config["rope_scaling"] == config["rope_parameters"] def test_qwen3_parallel_block_draft_applies_partial_rotary_factor(monkeypatch): diff --git a/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py b/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py index 8816e71d6b..1b2675737a 100644 --- a/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py +++ b/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py @@ -148,7 +148,7 @@ def test_factory_creates_legacy_controller_for_legacy_strategy(): def test_factory_calculates_adaptive_cache_capacities_across_local_dp_replicas_and_applies_gpu_ratio(monkeypatch): from lightllm.utils import envs_utils, kv_cache_utils - cpu_cache_meta = SimpleNamespace(page_num=10, calcu_one_page_size=lambda: 64) + cpu_cache_meta = SimpleNamespace(page_num=10, page_bytes=64) monkeypatch.setattr(kv_cache_utils, "calcu_cpu_cache_meta", lambda: cpu_cache_meta) monkeypatch.setattr(envs_utils, "get_cache_placement_gpu_capacity_ratio", lambda: 0.5) args = SimpleNamespace( 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..df68cc14bd 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 @@ -12,7 +12,7 @@ def test_disk_roundtrip_preserves_raw_pages_before_mm_boundary_trim(cpu_cache_fa light_mem = pytest.importorskip("light_mem") from lightllm.server.multi_level_kv_cache.disk_cache_worker import DiskCacheWorker - client = cpu_cache_factory(CpuKVCacheMeta(3, 1, 1, 1, 4096, torch.uint8, 0, torch.uint8)) + client = cpu_cache_factory(CpuKVCacheMeta(3, 4096)) pages, _ = client.allocate_pages([11, 22, 33], False) for index in pages: client.cpu_kv_cache_tensor[index].fill_(index + 1) @@ -69,19 +69,23 @@ def allocate(req, count): return torch.arange(count, dtype=torch.int32) module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.args = SimpleNamespace(page_size=1) + module.args = SimpleNamespace(page_size=1, cpu_cache_token_page_size=256) module.cpu_cache_client = client module.init_sync_group = None module.need_sync_compute_stream = lambda: False module.backend = SimpleNamespace( is_master_in_dp=True, + is_hybrid_att_model=False, _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=torch.arange(1024).view(1, -1)), - mem_manager=SimpleNamespace( - operator=SimpleNamespace( - load_cpu_cache_to_gpu=lambda **kwargs: copied_pages.extend(kwargs["page_indexes"].tolist()) - ) + att_state_manager=SimpleNamespace( + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: copied_pages.extend( + cpu_page_indices + ), + ), ), ), ) 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..04e03ee770 100644 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -9,6 +9,8 @@ from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache from lightllm.utils import shm_utils +pytestmark = pytest.mark.usefixtures("kv_allocator_environment") + @pytest.fixture(scope="module", autouse=True) def service_name(): @@ -95,7 +97,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")) @@ -245,33 +246,73 @@ def test_case9(): assert torch.equal(unmerged_node_d.token_id_key, torch.tensor([6], dtype=torch.int64)) -def test_case10(): - """ - 测试场景：测试 flush_cache 函数 - """ - print("\nTest Case 10: Testing flush_cache function\n") - tree = RadixCache(100, 0) - 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( - torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu"), update_refs=True - ) - assert tree_node is not None - assert size == 3 +def test_flush_cache_requires_released_references_and_returns_allocator_pages(): + from lightllm.common.attention_state.allocators import FullAttKVAllocator + + allocator = FullAttKVAllocator(100, pin_memory=False) + tree = RadixCache(100, 0, att_state_manager=SimpleNamespace(full_att_allocator=allocator)) + values = allocator.alloc(5) + key = torch.arange(5) + tree.insert(key[:3], values[:3]) + tree.insert(key, values) + node, matched_len, _ = tree.match_prefix(key[:3], update_refs=True) + assert matched_len == 3 + with pytest.raises(AssertionError, match="can not free"): + tree.flush_cache() + assert allocator.can_use_mem_size == 95 + tree.dec_node_ref_counter(node) tree.flush_cache() - tree_node, size, values = tree.match_prefix( - torch.tensor([1, 2, 3], dtype=torch.int64, device="cpu"), update_refs=True - ) - assert tree_node is None - assert size == 0 - assert tree.get_tree_total_tokens_num() == 0 - assert tree.get_refed_tokens_num() == 0 - assert len(tree.root_node.children) == 0 - assert tree.root_node.token_id_key.numel() == 0 - assert tree.root_node.token_mem_index_value.numel() == 0 + assert allocator.can_use_mem_size == 100 + assert tree.get_tree_total_tokens_num() == tree.get_refed_tokens_num() == 0 + assert not tree.root_node.children assert tree.root_node.ref_counter == 1 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") +def test_hybrid_flush_releases_shared_kv_and_state_slots(monkeypatch): + from unit_tests.common.att_resources.test_manager_helpers import composition + + monkeypatch.setattr( + "lightllm.common.attention_state.bounded_att.utils.get_env_start_args", + lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), + ) + resources = composition(device="cuda", with_tail=False) + allocator = resources.full_att_allocator + big = resources.bounded_att_big_page_allocator.alloc(1)[0] + small = resources.bounded_att_small_page_allocator.alloc(1)[0] + resources.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[big], + ) + resources.save_bounded_att_runtime_to_pages( + req_indices=[1], + state_token_lens=[12], + page_kind="small", + page_slot_ids=[small], + ) + values = allocator.alloc(12) + tree = HybridAttPagedRadixCache(12, 106, 4, 2, resources) + tree.insert( + torch.arange(12), + values, + block_hashs=[11, 22, 33], + block_state_idxs=[None, None, small], + len_to_big_page_id=SortedDict({8: big}), + ) + node, matched_len, _ = tree.match_prefix(torch.arange(12), block_hashs=[11, 22, 33], update_refs=True) + assert matched_len == 12 + with pytest.raises(AssertionError, match="can not free"): + tree.flush_cache() + tree.dec_node_ref_counter(node) + tree.flush_cache() + assert allocator.can_use_mem_size == 12 + assert resources.bounded_att_big_page_allocator.free_count == 3 + assert resources.bounded_att_small_page_allocator.free_count == 2 + assert not tree.root_node.children + + def test_page_aligned_insert_and_match(): tree = RadixCache(100, 99, page_size=4) values = torch.arange(100, 110, dtype=torch.int64) @@ -301,11 +342,11 @@ def test_page_aligned_insert_and_match(): assert tree.get_tree_total_tokens_num() == 12 -def test_page_size_must_match_mem_manager(): +def test_page_size_must_match_allocator(): mem_manager = SimpleNamespace(page_size=8) - with pytest.raises(ValueError, match="must match mem_manager page_size 8"): - RadixCache(100, 100, mem_manager=mem_manager, page_size=4) + with pytest.raises(ValueError, match="must match allocator page_size 8"): + RadixCache(100, 100, att_state_manager=SimpleNamespace(full_att_allocator=mem_manager), page_size=4) @pytest.mark.parametrize("update_refs", [False, True]) @@ -404,20 +445,25 @@ def test_page_key_bytes_does_not_share_tensor_memory(): assert key != TreeNode(page_size=4).get_child_key(token_ids) +def hybrid_resources(*, freed_big=None, freed_small=None): + return SimpleNamespace( + bounded_att_big_page_allocator=SimpleNamespace( + free=lambda slots: freed_big.extend(slots) if freed_big is not None else None + ), + bounded_att_small_page_allocator=SimpleNamespace( + free_count=0, free=lambda slots: freed_small.extend(slots) if freed_small is not None else None + ), + ) + + def test_hybrid_match_refs_only_valid_prefix_after_tail_state_eviction(): freed_state_ids = [] - small_page_buffers = SimpleNamespace( - get_free_cache_num=lambda: 0, - free_state_cache=lambda free_indexes: freed_state_ids.extend(free_indexes), - ) - 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, - small_page_buffers=small_page_buffers, + att_state_manager=hybrid_resources(freed_small=freed_state_ids), ) key = torch.arange(12, dtype=torch.int64) values = torch.arange(100, 112, dtype=torch.int64) @@ -457,10 +503,7 @@ 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) - ), - small_page_buffers=SimpleNamespace(free_state_cache=lambda **kwargs: None), + att_state_manager=hybrid_resources(freed_big=freed_big_states), ) return tree, freed_big_states @@ -526,10 +569,7 @@ 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) - ), - small_page_buffers=SimpleNamespace(get_free_cache_num=lambda: 0, free_state_cache=lambda **kwargs: None), + att_state_manager=hybrid_resources(freed_big=freed_big_states), ) # The older request leaves a small-page state at 384, before the image. old_values = torch.arange(1000, 1384) @@ -636,6 +676,7 @@ def test_hybrid_request_release_frees_or_transfers_image_page_slots( monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: context.args) context.radix_cache = tree + context.att_state_manager = tree.att_state_manager context.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(100, 132).view(1, -1)) req = InferReq.__new__(InferReq) req.req_idx = req.cur_output_len = 0 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py new file mode 100644 index 0000000000..0f0fd7ae88 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py @@ -0,0 +1,256 @@ +"""DP 前缀迁移：独立共享大页列表、目标页分配和源/目标索引配对。""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 +from lightllm.common.attention_state.allocators import StatePageAllocator +from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans + + +@pytest.mark.parametrize("hybrid,max_req_len", [(False, 32), (True, 32), (True, 4)]) +def test_dp_page_ids_use_separate_shared_array(monkeypatch: pytest.MonkeyPatch, hybrid: bool, max_req_len: int): + arrays = {} + + def shared_array(name: str, shape: tuple[int, ...], dtype: type) -> SimpleNamespace: + array = SimpleNamespace(arr=np.full(shape, -1, dtype=dtype), create_shm=lambda: None) + arrays[name] = array + return array + + monkeypatch.setattr(dp_shared_kv_trans, "ShmArray", shared_array) + monkeypatch.setattr(dp_shared_kv_trans, "get_dp_rank_in_node", lambda: 0) + monkeypatch.setattr(dp_shared_kv_trans, "get_env_start_args", lambda: SimpleNamespace(diverse_mode=False)) + backend = SimpleNamespace( + is_hybrid_att_model=hybrid, + args=SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2, max_req_total_len=max_req_len), + ) + module = dp_shared_kv_trans.DPKVSharedMoudle(max_req_num=3, dp_size_in_node=2, backend=backend) + assert arrays["dp_shared_req_infos"].arr.shape == (3, 2, 2) + if hybrid and max_req_len >= 8: + assert arrays["dp_shared_req_big_page_ids"].arr.shape == (3, 2, 4) + else: + assert module.shared_req_big_page_ids is None + assert list(arrays) == ["dp_shared_req_infos"] + + +def test_fill_orders_radix_and_owned_pages_and_clears_reused_rows(monkeypatch: pytest.MonkeyPatch): + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = SimpleNamespace(is_master_in_dp=True, node_nccl_group=None) + module.dp_rank_in_node = 0 + module.big_page_tokens = 8 + module.shared_req_infos = SimpleNamespace(arr=np.full((2, 2, 2), 99, dtype=np.int64)) + module.shared_req_big_page_ids = SimpleNamespace(arr=np.full((2, 2, 4), 99, dtype=np.int64)) + node = object() + context = SimpleNamespace(radix_cache=SimpleNamespace(get_big_page_ids_by_node=lambda _: [40])) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) + req = SimpleNamespace( + cur_kv_len=24, + req_idx=7, + shared_kv_node=node, + hybrid_len_to_big_page_id={24: 3, 16: 12, 32: 88}, + ) + second = SimpleNamespace( + cur_kv_len=8, + req_idx=5, + shared_kv_node=None, + hybrid_len_to_big_page_id={8: 6}, + ) + module.fill_reqs_info([req, second]) + assert module.shared_req_infos.arr[:, 0].tolist() == [[24, 7], [8, 5]] + assert module.shared_req_big_page_ids.arr[:, 0].tolist() == [[40, 12, 3, -1], [6, -1, -1, -1]] + assert (module.shared_req_big_page_ids.arr[:, 1] == 99).all() + + req.cur_kv_len = 8 + module.fill_reqs_info([req]) + assert module.shared_req_big_page_ids.arr[0, 0].tolist() == [40, -1, -1, -1] + module.backend.is_master_in_dp = False + req.cur_kv_len = 0 + module.fill_reqs_info([req]) + assert module.shared_req_infos.arr[0, 0, 0] == 8 + + +@pytest.mark.parametrize("capacity,src_kv_len", [(2, 26), (0, 14)]) +def test_build_reserves_only_new_big_pages( + monkeypatch: pytest.MonkeyPatch, + capacity: int, + src_kv_len: int, +): + module, req, allocator, allocations = _build_fixture(monkeypatch, capacity, src_kv_len) + tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) + (task,) = tasks + assert task.dst_req is req and task.src_req_idx == 1 + assert task.src_dp_rank == 1 + assert task.src_token_indices.tolist() == list(range(40, 32 + src_kv_len)) + assert len(task.dst_token_indices) == src_kv_len - 8 + assert allocations == [((src_kv_len + 3) // 4) * 4 - 8] + assert req.cur_kv_len == 8 and req.hybrid_len_to_big_page_id == {8: 77} + if src_kv_len == 26: + assert task.src_big_page_ids == [20, 21] + assert task.dst_big_page_ids == [0, 1] + assert allocator.free_count == 0 + else: + assert task.src_big_page_ids == task.dst_big_page_ids == [] + + +def test_missing_source_page_fails_before_allocating(monkeypatch: pytest.MonkeyPatch): + module, req, allocator, allocations = _build_fixture(monkeypatch, 2, 26) + module.shared_req_big_page_ids.arr[0, 1, 2] = -1 + with pytest.raises(AssertionError, match="missing prefix big pages"): + module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) + assert not allocations and allocator.free_count == 2 + + +@pytest.mark.parametrize("hybrid", [False, True]) +@pytest.mark.parametrize("local_dp_rank", [0, 1]) +@pytest.mark.parametrize("local_len,remote_len", [(24, 16), (24, 24), (31, 24), (0, 0)]) +def test_local_longest_prefix_skips_copy_but_keeps_barriers( + monkeypatch: pytest.MonkeyPatch, + hybrid: bool, + local_dp_rank: int, + local_len: int, + remote_len: int, +): + # 并列最长且本地为 DP 1 时，argmax 会优先选 DP 0；仍应在选择来源前跳过复制。 + allocate_pages = Mock(side_effect=AssertionError("unexpected big page allocation")) + allocate_tokens = Mock(side_effect=AssertionError("unexpected KV allocation")) + copy_resources = Mock(side_effect=AssertionError("unexpected DP copy")) + allocation_need = Mock(side_effect=AssertionError("unexpected allocation calculation")) + available_tokens = Mock(side_effect=AssertionError("unexpected capacity query")) + barriers = [] + manager = SimpleNamespace( + bounded_att_big_page_allocator=SimpleNamespace(alloc=allocate_pages), + get_dp_shared_operator=lambda: SimpleNamespace(copy_from_dp=copy_resources), + ) + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = SimpleNamespace( + is_hybrid_att_model=hybrid, + is_master_in_dp=True, + node_nccl_group="node", + model=SimpleNamespace(att_state_manager=manager), + dp_att_state_managers={}, + dp_req_to_token_indices={}, + _alloc_req_kv_mem=allocate_tokens, + ) + module.dp_rank_in_node = local_dp_rank + module.big_page_tokens = 8 if hybrid else 0 + module.shared_req_infos = SimpleNamespace(arr=np.zeros((1, 2, 2), dtype=np.int64)) + module.shared_req_infos.arr[0, 1 - local_dp_rank] = [remote_len, 9] + module.shared_req_big_page_ids = SimpleNamespace(arr=np.full((1, 2, 4), -1, dtype=np.int64)) if hybrid else None + owned_pages = {length: 70 + length // 8 for length in range(8, local_len + 1, 8)} if hybrid else {} + req = SimpleNamespace( + req_idx=7, + cur_kv_len=local_len, + hold_kv_len=local_len, + shared_kv_node=None, + hybrid_len_to_big_page_id=owned_pages.copy(), + get_cur_total_len=lambda: 32, + _kv_cache_alloc_need=allocation_need, + shm_req=SimpleNamespace(shm_cur_kv_len=local_len), + ) + monkeypatch.setattr( + dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=available_tokens) + ) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: barriers.append(group)) + + module.fill_reqs_info([req]) + assert module.shared_req_infos.arr[0, local_dp_rank].tolist() == [local_len, req.req_idx] + tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[local_dp_rank]) + assert tasks == [] + module.kv_trans(tasks) + + for operation in (allocate_pages, allocate_tokens, copy_resources, allocation_need, available_tokens): + operation.assert_not_called() + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == local_len + assert req.hybrid_len_to_big_page_id == owned_pages + # 即使不接收数据，也必须参加所有屏障，保证其他 DP 能完成读取后再释放来源资源。 + assert barriers == ["node", "node", "node"] + + +def _build_fixture(monkeypatch: pytest.MonkeyPatch, capacity: int, src_kv_len: int) -> tuple: + allocator = StatePageAllocator(capacity) + allocations = [] + req = SimpleNamespace( + req_idx=0, + cur_kv_len=8, + hold_kv_len=8, + hybrid_len_to_big_page_id={8: 77}, + get_cur_total_len=lambda: 32, + ) + req._kv_cache_alloc_need = lambda length: ((length + 3) // 4) * 4 - req.hold_kv_len + + def allocate(cur_req: SimpleNamespace, count: int) -> torch.Tensor: + allocations.append(count) + cur_req.hold_kv_len += count + return torch.arange(100, 100 + count, dtype=torch.int32) + + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = SimpleNamespace( + is_hybrid_att_model=True, + dp_world_size=2, + rank_in_dp=1, + node_nccl_group=None, + model=SimpleNamespace(att_state_manager=SimpleNamespace(bounded_att_big_page_allocator=allocator)), + dp_req_to_token_indices={1: torch.arange(64, dtype=torch.int32).reshape(2, 32)}, + _alloc_req_kv_mem=allocate, + ) + module.dp_rank_in_node = 0 + module.big_page_tokens = 8 + module.shared_req_infos = SimpleNamespace(arr=np.array([[[8, 0], [src_kv_len, 1]]], dtype=np.int64)) + module.shared_req_big_page_ids = SimpleNamespace(arr=np.array([[[77, -1, -1, -1], [10, 20, 21, -1]]])) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", SimpleNamespace(get_can_alloc_token_num=lambda: 100)) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) + return module, req, allocator, allocations + + +@pytest.mark.parametrize("prefix_len", [8, 16]) +def test_build_evicts_cached_state_pages(monkeypatch, kv_allocator_environment, prefix_len): + from sortedcontainers import SortedDict + from lightllm.common.attention_state.allocators import FullAttKVAllocator + from lightllm.server.router.dynamic_prompt import hybrid_att_radix_cache + from lightllm.server.router.model_infer.infer_batch import InferenceContext + from lightllm.server.router.model_infer.mode_backend import base_backend + + module, req, allocator, _ = _build_fixture(monkeypatch, 2, prefix_len) + req.cur_kv_len = req.hold_kv_len = 0 + req.hybrid_len_to_big_page_id = {} + module.shared_req_infos.arr[0, 0, 0] = 0 + full = FullAttKVAllocator(16, page_size=4, pin_memory=False) + manager = module.backend.model.att_state_manager + manager.full_att_allocator = full + manager.bounded_att_small_page_allocator = StatePageAllocator(0) + monkeypatch.setattr( + hybrid_att_radix_cache, "SharedArray", lambda name, shape, dtype: SimpleNamespace(arr=np.zeros(shape, dtype)) + ) + radix = hybrid_att_radix_cache.HybridAttPagedRadixCache(16, 0, 4, 2, manager) + slots = allocator.alloc(2) + for i, slot in enumerate(slots): + radix.insert( + key=torch.arange(8) + i * 8, + value=full.alloc(8).clone(), + block_hashs=[i * 2, i * 2 + 1], + block_state_idxs=[None, None], + len_to_big_page_id=SortedDict({8: slot}), + ) + requests = SimpleNamespace(req_to_token_indexs=torch.full((1, 32), -1, dtype=torch.int32)) + context = SimpleNamespace(att_state_manager=manager, radix_cache=radix, req_manager=requests) + context.get_can_alloc_token_num = lambda: InferenceContext.get_can_alloc_token_num(context) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) + monkeypatch.setattr(base_backend, "g_infer_context", context) + module.backend.args = SimpleNamespace(page_size=4) + module.backend._alloc_req_kv_mem = lambda req, count: base_backend.ModeBackend._alloc_req_kv_mem( + module.backend, req, count + ) + assert full.can_use_mem_size == allocator.free_count == 0 + assert context.get_can_alloc_token_num() == 16 + (task,) = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) + assert len(task.dst_big_page_ids) == prefix_len // 8 + assert len(task.dst_token_indices) == req.hold_kv_len == prefix_len + assert radix.get_tree_total_tokens_num() == 16 - prefix_len + assert torch.equal(requests.req_to_token_indexs[0, :prefix_len], task.dst_token_indices) + assert allocator.free_count == 0 + assert full.can_use_mem_size == 0 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 82422adf8b..aba9ff3906 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 @@ -97,7 +97,7 @@ def start_offload(req, cpu_kv_cache_stream): offload_calls.append((req.shm_req.request_id, CacheTier.DISK in req.cache_tiers)) return SimpleNamespace(req=req) - module._start_kv_cache_offload_task = start_offload + module.write_req_to_cpu = start_offload monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "is_hybrid_att_model", False) monkeypatch.setattr( multi_level_kv_cache_impl.g_infer_context, @@ -137,9 +137,9 @@ def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): context = InferenceContext() context.is_hybrid_att_model = True context.req_manager = SimpleNamespace(req_to_token_indexs=torch.tensor([[10, 11, 12]])) - context.radix_cache = SimpleNamespace( - small_page_buffers=SimpleNamespace(free_state_cache=freed_small_pages.extend), - big_page_buffers=SimpleNamespace(free_state_cache=freed_big_pages.extend), + context.att_state_manager = SimpleNamespace( + bounded_att_small_page_allocator=SimpleNamespace(free=freed_small_pages.extend), + bounded_att_big_page_allocator=SimpleNamespace(free=freed_big_pages.extend), ) req = SimpleNamespace( req_idx=0, @@ -181,14 +181,20 @@ def alloc_req_kv_mem(req, alloc_token_num): req.hold_kv_len = new_hold_kv_len return mem_indexes - operator = SimpleNamespace( - load_cpu_cache_to_gpu=lambda mem_indexes, **kwargs: loaded_indexes.extend(mem_indexes.tolist()) + resources = SimpleNamespace( + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded_indexes.extend( + token_indices.tolist() + ), + ), ) backend = SimpleNamespace( is_master_in_dp=False, + is_hybrid_att_model=False, radix_cache=None, model=SimpleNamespace( - mem_manager=SimpleNamespace(operator=operator), + att_state_manager=resources, req_manager=SimpleNamespace(req_to_token_indexs=table), ), _alloc_req_kv_mem=alloc_req_kv_mem, @@ -198,7 +204,7 @@ def alloc_req_kv_mem(req, alloc_token_num): module.backend = backend module.init_sync_group = object() module.need_sync_compute_stream = lambda: False - module.cpu_cache_client = SimpleNamespace() + module.cpu_cache_client = SimpleNamespace(cpu_kv_cache_tensor=None) context = SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table, mem_manager=SimpleNamespace(alloc=alloc)), get_can_alloc_token_num=lambda: 200, diff --git a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py new file mode 100644 index 0000000000..17d564987d --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import concurrent.futures +import socket +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import rpyc +import pytest + +from lightllm.server.router.model_infer.mode_backend.pd import nccl_kv_transporter as module + + +def test_page_events_are_recorded_after_send_and_receive(monkeypatch: pytest.MonkeyPatch): + page_buffers = module.torch.empty(size=(2, 8), dtype=module.torch.uint8, device="cpu") + transporter = SimpleNamespace(kv_move_buffer=page_buffers, args=SimpleNamespace(pd_detail_log=False)) + peer = module._NcclPeer(transporter, "remote") + stream = object() + operations = [] + comm = SimpleNamespace( + send=Mock(side_effect=lambda tensor, dst, stream: operations.append(("send", tensor.numel(), dst, stream))), + recv=Mock(side_effect=lambda tensor, src, stream: operations.append(("recv", tensor.numel(), src, stream))), + ) + monkeypatch.setattr(peer, "_ensure_comm", lambda is_server: comm) + monkeypatch.setattr(peer, "_get_stream", lambda: stream) + event = Mock() + event.record.side_effect = lambda stream: operations.append(("record", stream)) + monkeypatch.setattr(module.torch.cuda, "Event", lambda: event) + task = SimpleNamespace(src_page_index=1, dst_page_index=0) + + handle = peer.send_page(task) + assert operations == [("send", 8, 1, stream), ("record", stream)] + assert handle.mode == "send" + assert handle.finish_cuda_event is event + event.query.return_value = True + assert handle.check_status() == "DONE" + + operations.clear() + receive_queue = Mock() + monkeypatch.setattr(peer, "_get_recv_queue", lambda: receive_queue) + peer.start_recv(task) + queued_task = receive_queue.put.call_args.args[0] + assert queued_task.xfer_handle is task.xfer_handle + assert task.xfer_handle.mode == "recv" + assert task.xfer_handle.finish_cuda_event is None + peer._recv_page(queued_task) + assert operations == [("recv", 8, 0, stream), ("record", stream)] + assert task.xfer_handle.finish_cuda_event is event + assert task.xfer_handle.check_status() == "DONE" + assert event.query.call_count == 2 + + +def test_receive_waits_for_record_to_return(monkeypatch: pytest.MonkeyPatch): + record_started = threading.Event() + finish_record = threading.Event() + event = Mock() + + def record(stream: object) -> None: + record_started.set() + assert finish_record.wait(timeout=5) + + event.record.side_effect = record + monkeypatch.setattr(module.torch.cuda, "Event", lambda: event) + transporter = SimpleNamespace( + kv_move_buffer=module.torch.empty(size=(1, 8), dtype=module.torch.uint8), + args=SimpleNamespace(pd_detail_log=False), + ) + peer = module._NcclPeer(transporter, "remote") + monkeypatch.setattr(peer, "_ensure_comm", lambda is_server: SimpleNamespace(recv=Mock())) + monkeypatch.setattr(peer, "_get_stream", object) + handle = module._NcclXferHandle(peer_name="remote", mode="recv") + task = SimpleNamespace(dst_page_index=0, xfer_handle=handle) + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + receive = executor.submit(peer._recv_page, task) + try: + assert record_started.wait(timeout=5) + assert handle.check_status() == "PROC" + assert handle.finish_cuda_event is None + event.query.assert_not_called() + finally: + finish_record.set() + receive.result(timeout=5) + assert handle.check_status() == "DONE" + assert handle.finish_cuda_event is event + event.query.assert_called_once() + + +def test_receive_submission_failure_wakes_waiter(monkeypatch: pytest.MonkeyPatch): + transporter = SimpleNamespace( + kv_move_buffer=module.torch.empty(size=(1, 8), dtype=module.torch.uint8), + send_error_info_to_prefill_node=Mock(), + ) + peer = module._NcclPeer(transporter, "remote") + comm = SimpleNamespace(recv=Mock(side_effect=RuntimeError("recv failed"))) + monkeypatch.setattr(peer, "_ensure_comm", lambda is_server: comm) + monkeypatch.setattr(peer, "_get_stream", object) + monkeypatch.setattr(peer, "_drop_comm", Mock()) + handle = module._NcclXferHandle(peer_name="remote", mode="recv") + task = SimpleNamespace(dst_page_index=0, xfer_handle=handle) + + peer._recv_page(task) + + assert handle.check_status() == "ERR" + assert handle.error_info == "recv failed" + transporter.send_error_info_to_prefill_node.assert_called_once_with(task) + + +def test_concurrent_control_channels_retry_port_before_publishing_metadata(monkeypatch: pytest.MonkeyPatch): + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + first_port = probe.getsockname()[1] + + original = module.ThreadedServer + both_bound = threading.Barrier(2) + + def bind_before_listening(*args: object, **kwargs: object): + server = original(*args, **kwargs) + if kwargs["port"] == first_port: + # Force both SO_REUSEADDR sockets to bind before either listens. + both_bound.wait(timeout=10) + return server + + monkeypatch.setattr(module, "ThreadedServer", bind_before_listening) + channels = [] + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit( + module._NcclControlChannel, + "127.0.0.1", + first_port, + first_port + 4, + ) + for _ in range(2) + ] + for future in futures: + channels.append(future.result(timeout=15)) + assert len({channel.port for channel in channels}) == 2 + for index, channel in enumerate(channels): + payload = f"channel-{index}".encode() + connection = rpyc.connect("127.0.0.1", channel.port) + try: + connection.root.push_notif(payload) + finally: + connection.close() + assert channel.get_notifs() == [payload] + finally: + for channel in channels: + channel.close() 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 22ccc9e3cd..1da16313fd 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 @@ -13,24 +13,6 @@ from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_impl -def _bind_alloc_need(req, page_size): - req.args = SimpleNamespace(page_size=page_size) - req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) - return req - - -def _reserve_into_table(table, req, alloc_sizes, next_index, alloc_token_num): - if alloc_token_num == 0: - return None - alloc_sizes.append(alloc_token_num) - new_hold_kv_len = req.hold_kv_len + alloc_token_num - mem_indexes = torch.arange(next_index[0], next_index[0] + alloc_token_num, dtype=torch.int32) - table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes - next_index[0] += alloc_token_num - req.hold_kv_len = new_hold_kv_len - return mem_indexes - - @pytest.mark.parametrize("is_hybrid_att_model", [False, True]) def test_pd_decode_reserves_model_pages_but_transfers_only_logical_kv(monkeypatch, is_hybrid_att_model): table = torch.full((1, 16), -1, dtype=torch.int32) @@ -42,7 +24,13 @@ def test_pd_decode_reserves_model_pages_but_transfers_only_logical_kv(monkeypatc backend.args = SimpleNamespace(pd_kv_page_size=4, page_size=4) backend.model = SimpleNamespace(req_manager=SimpleNamespace(req_to_token_indexs=table)) backend.is_master_in_dp = False - backend._alloc_req_kv_mem = lambda req, size: _reserve_into_table(table, req, alloc_sizes, next_index, size) + backend._alloc_req_kv_mem = lambda req, size: _reserve_into_table( + table=table, + req=req, + alloc_sizes=alloc_sizes, + next_index=next_index, + alloc_token_num=size, + ) def create_task(**kwargs): tasks.append( @@ -98,10 +86,16 @@ def test_pd_prefill_rejects_unaligned_transfer_page_size(): req = SimpleNamespace(cur_kv_len=4, shm_req=SimpleNamespace(input_len=10)) with pytest.raises(AssertionError, match="pd_kv_page_size must be divisible by page_size"): - backend._prefill_chuncked_handle_func(req, next_token_id=1, next_token_prob=1.0, output_len=0) + backend._prefill_chuncked_handle_func( + req_obj=req, + next_token_id=1, + next_token_prob=1.0, + output_len=0, + ) -def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypatch): +@pytest.mark.parametrize("tp_size,tp_rank", [(1, 0), (2, 0), (2, 1)]) +def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypatch, tp_size, tp_rank): table = torch.full((1, 16), -1, dtype=torch.int32) table[0, :4] = torch.arange(4, dtype=torch.int32) alloc_sizes = [] @@ -117,14 +111,20 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat ) source_table = torch.arange(32, dtype=torch.int32).reshape(2, 16) backend = SimpleNamespace( - is_hybrid_att_model=False, model=SimpleNamespace(mem_manager=SimpleNamespace()), - dp_world_size=1, - rank_in_dp=0, + dp_world_size=tp_size, + rank_in_dp=tp_rank, node_nccl_group=None, - mem_managers=[SimpleNamespace(req_to_token_indexs=source_table)] * 2, + dp_req_to_token_indices={0: source_table, 1: source_table}, + is_hybrid_att_model=False, + ) + backend._alloc_req_kv_mem = lambda cur_req, size: _reserve_into_table( + table=table, + req=cur_req, + alloc_sizes=alloc_sizes, + next_index=next_index, + alloc_token_num=size, ) - backend._alloc_req_kv_mem = lambda cur_req, size: _reserve_into_table(table, cur_req, alloc_sizes, next_index, size) module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) module.backend = backend module.dp_rank_in_node = 0 @@ -142,5 +142,96 @@ def test_dp_cache_fetch_reserves_pages_and_keeps_logical_transfer_size(monkeypat assert req.cur_kv_len == 4 assert req.hold_kv_len == 12 assert len(tasks) == 1 - assert tasks[0].mem_indexes.tolist() == list(range(4, 10)) - assert len(tasks[0].max_kv_len_mem_indexes) == 6 + assert tasks[0].src_dp_rank == 1 + assert tasks[0].dst_token_indices.tolist() == list(range(4, 10)) + assert len(tasks[0].src_token_indices) == 6 + + +@pytest.mark.parametrize( + "node_size,tp_size,rank_in_node,expected_ranks", + [ + (4, 2, 0, [0, 2]), + (4, 2, 1, [1, 3]), + (4, 2, 2, [0, 2]), + (4, 2, 3, [1, 3]), + (6, 2, 5, [1, 3, 5]), + (2, 4, 1, [1]), + ], +) +def test_dp_imports_only_matching_tp_shards( + monkeypatch, + node_size, + tp_size, + rank_in_node, + expected_ranks, +): + from lightllm.server.router.model_infer.mode_backend import base_backend + from lightllm.server.router.model_infer import shared_att_resources_utils + + imported_ranks = [] + registered_ranks = [] + managers = {} + tables = {} + for rank in expected_ranks: + managers[rank] = SimpleNamespace( + get_dp_shared_operator=lambda rank=rank: SimpleNamespace( + register_page_pools_pin_memory=lambda: registered_ranks.append(rank), + ), + ) + tables[rank] = object() + + def load_storage(rank, kind): + assert kind == "dp" + imported_ranks.append(rank) + return managers[rank], tables[rank] + + local_manager = managers[rank_in_node] + backend = SimpleNamespace( + args=SimpleNamespace(running_max_req_size=8), + dp_size_in_node=max(1, node_size // tp_size), + dp_world_size=tp_size, + node_world_size=node_size, + rank_in_node=rank_in_node, + model=SimpleNamespace( + att_state_manager=local_manager, + req_manager=SimpleNamespace(req_to_token_indexs=tables[rank_in_node]), + ), + ) + monkeypatch.setattr(base_backend, "get_current_device_id", lambda: 0) + monkeypatch.setattr(base_backend.torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr(dp_shared_kv_trans, "DPKVSharedMoudle", lambda **kwargs: object()) + monkeypatch.setattr(shared_att_resources_utils, "load_att_resources", load_storage) + + base_backend.ModeBackend.init_dp_kv_shared(backend) + + remote_ranks = [rank for rank in expected_ranks if rank != rank_in_node] + assert imported_ranks == registered_ranks == remote_ranks + assert list(backend.dp_att_state_managers) == list(range(len(expected_ranks))) + assert list(backend.dp_req_to_token_indices) == list(range(len(expected_ranks))) + for dp_rank, rank in enumerate(expected_ranks): + assert backend.dp_att_state_managers[dp_rank] is managers[rank] + assert backend.dp_req_to_token_indices[dp_rank] is tables[rank] + + +def _bind_alloc_need(req, page_size): + req.args = SimpleNamespace(page_size=page_size) + req._kv_cache_alloc_need = lambda target_len: InferReq._kv_cache_alloc_need(req, target_len) + return req + + +def _reserve_into_table( + table, + req, + alloc_sizes, + next_index, + alloc_token_num, +): + if alloc_token_num == 0: + return None + alloc_sizes.append(alloc_token_num) + new_hold_kv_len = req.hold_kv_len + alloc_token_num + mem_indexes = torch.arange(next_index[0], next_index[0] + alloc_token_num, dtype=torch.int32) + table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes + next_index[0] += alloc_token_num + req.hold_kv_len = new_hold_kv_len + return mem_indexes diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py index 7bd0940c89..8f99c6add3 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py @@ -124,8 +124,7 @@ def test_notification_payload(backend, page_kind, method, stage, destination): if backend is NcclKVTransporter: transporter.host_ip = "127.0.0.1" transporter.control_channel = SimpleNamespace(port=1234, send_notif=Mock()) - transporter.kv_move_buffer = SimpleNamespace(dtype="float16") - transporter.page_size = transporter.num_layers = transporter.kv_head_num = transporter.head_dims = 1 + transporter.kv_move_buffer = SimpleNamespace(dtype="uint8", shape=(12, 128)) transporter._get_remote_metadata = Mock(return_value=SimpleNamespace(host_ip="127.0.0.2", control_port=4321)) transporter._get_peer = Mock(return_value=SimpleNamespace(start_recv=Mock())) else: @@ -190,3 +189,35 @@ def test_receiver_rejects_payload_with_decode_local_mark(): ) assert result.returncode != 0 assert b"decode modules are unavailable to the receiver" in result.stderr + + +def test_nixl_rejects_different_peer_page_capacity_before_registering_handles(): + transporter = NixlKVTransporter.__new__(NixlKVTransporter) + transporter.page_bytes = 128 + transporter.nixl_agent = SimpleNamespace(get_xfer_descs=Mock(), prep_xfer_dlist=Mock()) + with pytest.raises(AssertionError, match="same page byte capacity"): + transporter._create_paged_xfer_handles([(4096, 512, 0, None)], 2, "peer") + transporter.nixl_agent.get_xfer_descs.assert_not_called() + transporter._create_paged_xfer_handles([(4096, 256, 0, None)], 2, "peer") + transporter.nixl_agent.get_xfer_descs.assert_called_once_with([(4096, 128, 0), (4224, 128, 0)], "VRAM") + + +def test_nccl_checks_peer_page_bytes_before_accepting_agent(): + from lightllm.server.router.model_infer.mode_backend.pd.nccl_kv_transporter import NcclAgentMetadata + + transporter = NcclKVTransporter.__new__(NcclKVTransporter) + transporter.remote_agents = {} + transporter.kv_move_buffer = SimpleNamespace(shape=(2, 128)) + transporter.num_pages = 2 + assert pickle.loads(transporter.local_page_mem_desc) == {"num_pages": 2, "page_bytes": 128} + peer = SimpleNamespace( + agent_name="peer", + agent_metadata=pickle.dumps(NcclAgentMetadata("peer", "127.0.0.1", 1234, 0)), + page_reg_desc=pickle.dumps({"page_bytes": 256}), + ) + with pytest.raises(AssertionError, match="same page byte capacity"): + transporter.connect_add_remote_agent(peer) + assert not transporter.remote_agents + peer.page_reg_desc = transporter.local_page_mem_desc + transporter.connect_add_remote_agent(peer) + assert transporter.remote_agents["peer"] is peer diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py new file mode 100644 index 0000000000..726df8dbee --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py @@ -0,0 +1,126 @@ +"""PD 索引由普通 CPU Tensor 上传，使用实际拷贝流并跳过运行态任务。""" + +from types import SimpleNamespace +from unittest.mock import Mock +from queue import Queue + +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_trans_process +from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_trans_process + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("prefill", [True, False]) +def test_pd_uploads_cpu_indices_on_copy_stream_and_skips_runtime_pages( + monkeypatch: pytest.MonkeyPatch, + prefill: bool, +) -> None: + implementation = prefill_trans_process if prefill else decode_trans_process + module_class = implementation._PrefillTransModule if prefill else implementation._DecodeTransModule + device_id = torch.cuda.current_device() + args = SimpleNamespace(tp=2, dp=2, nnodes=1, pd_kv_page_num=2, pd_kv_page_size=4, pd_node_id=1) + resource = SimpleNamespace(get_pd_operator=lambda: SimpleNamespace(get_page_bytes=lambda: 64)) + uploads = [] + copies = [] + original_cuda = torch.Tensor.cuda + + def track_upload(tensor: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: + uploads.append( + (tensor.device.type, tensor.is_pinned(), kwargs.get("non_blocking"), torch.cuda.current_stream()) + ) + return original_cuda(tensor, *args, **kwargs) + + def capture_copy(**kwargs: object) -> None: + indices = kwargs["token_indices"] + copies.append((kwargs, None if indices is None else indices.clone())) + + monkeypatch.setattr(torch.Tensor, "cuda", track_upload) + monkeypatch.setattr(implementation, "copy_page", capture_copy) + monkeypatch.setattr(implementation, "create_kv_transporter", lambda **kwargs: SimpleNamespace()) + monkeypatch.setattr(implementation.threading.Thread, "start", lambda self: None) + if prefill: + module = module_class( + args=args, + device_id=device_id, + task_in_queue=None, + task_out_queue=None, + resource_managers=[resource, resource], + ) + else: + module = module_class( + args=args, + device_id=device_id, + task_in_queue=None, + task_out_queue=None, + resource_managers=[resource, resource], + up_status_in_queue=None, + ) + + # 构造时两个 DP 的预热均上传单个零索引。 + assert [indices.cpu().tolist() for _, indices in copies] == [[0], [0]] + assert [kwargs["dp_rank_in_node"] for kwargs, _ in copies] == [0, 1] + + tasks = [ + SimpleNamespace( + mem_indexes=indices, + page_kind=kind, + src_page_index=page, + dst_page_index=page, + prefill_dp_index=0, + decode_dp_index=0, + req_idx=1, + end_kv_index=7, + xfer_handle=None, + ) + for page, kind, indices in [(0, "kv", [5, 1, 7]), (1, "kv", [2]), (0, "att_state", None)] + ] + task_queue = module.local_copy_kv_queue if prefill else module.ready_page_task_queue + copy_loop = module.local_copy_kv_loop if prefill else module.read_page_to_mems_loop + monkeypatch.setattr(task_queue, "get", Mock(side_effect=[*tasks, _StopLoop()])) + with pytest.raises(_StopLoop): + copy_loop() + module.copy_cuda_stream.synchronize() + + copied_indices = [None if indices is None else indices.cpu().tolist() for _, indices in copies[2:]] + assert copied_indices == [[5, 1, 7], [2], None] + assert len(uploads) == 4 # 两次预热、两次 KV 上传；运行态任务不访问 mem_indexes。 + + # 下一批使用较短的索引，验证上传的数据与本次任务一致。 + tasks[0].mem_indexes = [9, 8] + monkeypatch.setattr(task_queue, "get", Mock(side_effect=[tasks[0], _StopLoop()])) + with pytest.raises(_StopLoop): + copy_loop() + module.copy_cuda_stream.synchronize() + assert copies[-1][1].cpu().tolist() == [9, 8] + assert all( + device == "cpu" and not pinned and non_blocking and stream == module.copy_cuda_stream + for device, pinned, non_blocking, stream in uploads + ) + assert all(kwargs["write_to_pd_page"] is prefill for kwargs, _ in copies) + assert all(indices is None or indices.is_cuda and indices.dtype == torch.long for _, indices in copies) + + +class _StopLoop(Exception): + pass + + +def test_failed_receive_is_not_read_or_copied(monkeypatch: pytest.MonkeyPatch) -> None: + module = decode_trans_process._DecodeTransModule.__new__(decode_trans_process._DecodeTransModule) + module.device_id = 0 + module.failed_queue = Queue() + handle = SimpleNamespace(check_status=Mock(return_value="ERR"), error_info="recv failed") + task = SimpleNamespace(xfer_handle=handle) + module.transporter = SimpleNamespace(check_task_status=Mock(return_value="ERR")) + module.ready_page_task_queue = SimpleNamespace(get=Mock(side_effect=[task, _StopLoop()])) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + copy = Mock() + monkeypatch.setattr(decode_trans_process, "copy_page", copy) + + with pytest.raises(_StopLoop): + module.read_page_to_mems_loop() + + assert module.failed_queue.get_nowait() is task + assert task.error_info == "recv failed" + copy.assert_not_called() diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py b/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py deleted file mode 100644 index 9e44ed774b..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_prompt_cache_preload.py +++ /dev/null @@ -1,52 +0,0 @@ -from types import SimpleNamespace - -import torch - -from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend - - -def test_preload_prompt_cache_keeps_only_complete_kv_pages(monkeypatch): - loaded = {} - - class FakeMemManager: - def alloc(self, size): - loaded["alloc_size"] = size - return torch.arange(size, dtype=torch.int32) - - def load_index_kv_buffer(self, indexes, buffers): - loaded["indexes"] = indexes.clone() - loaded["buffers"] = buffers - - class FakeRadixCache: - def __init__(self): - self.mem_manager = FakeMemManager() - - def insert(self, token_ids, indexes): - loaded["insert"] = (token_ids.clone(), indexes.clone()) - - def match_prefix(self, token_ids, update_refs): - loaded["match"] = (token_ids.clone(), update_refs) - - backend = ModeBackend.__new__(ModeBackend) - backend.logger = SimpleNamespace(info=lambda _: None) - backend.weight_dir = "/model" - backend.args = SimpleNamespace(page_size=4) - backend.radix_cache = FakeRadixCache() - kv_buffer = torch.arange(12).reshape(1, 6, 2) - monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) - monkeypatch.setattr(torch, "load", lambda *args, **kwargs: {"kv_buffer": kv_buffer}) - - backend.preload_prompt_cache_kv_buffer( - { - "prompt_cache_kv_buffer": {"rank_0": "prompt_cache.pt"}, - "prompt_cache_token_ids": [10, 11, 12, 13, 14, 15], - } - ) - - assert loaded["alloc_size"] == 4 - assert loaded["indexes"].tolist() == [0, 1, 2, 3] - assert loaded["buffers"]["kv_buffer"].shape == (1, 4, 2) - assert loaded["insert"][0].tolist() == [10, 11, 12, 13] - assert loaded["insert"][1].tolist() == [0, 1, 2, 3] - assert loaded["match"][0].tolist() == [10, 11, 12, 13] - assert loaded["match"][1] is True diff --git a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py b/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py index dff0e892a9..7a631e398d 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py @@ -1,7 +1,12 @@ from types import SimpleNamespace +from unittest.mock import Mock import pytest +import torch +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.req_manager import ReqManager +from lightllm.common.req_manager.base import _ReqLinkedList from lightllm.server.router.model_infer.mode_backend.rl_backend_ops import RlBackendOps from lightllm.utils.torch_memory_saver_utils import MemoryTag @@ -10,24 +15,41 @@ ("tags", "should_init_hold_indexes"), [ ([MemoryTag.WEIGHT], False), + ([MemoryTag.GRAPH], False), + ([], False), ([MemoryTag.KV_CACHE], True), (None, True), ], ) -def test_resume_memory_initializes_hold_indexes_only_for_kv_cache(tags, should_init_hold_indexes): - events = [] +def test_resume_memory_initializes_hold_indexes_only_for_kv_cache( + tags, should_init_hold_indexes, kv_allocator_environment +): + allocator = FullAttKVAllocator(16, page_size=4, pin_memory=False) + req_manager = ReqManager.__new__(ReqManager) + req_manager.req_list = _ReqLinkedList(2) + req_manager.HOLD_REQUEST_ID = 2 + req_manager.page_size = allocator.page_size + req_manager.req_to_token_indexs = torch.full((3, 8), -1, dtype=torch.int32) + resumed = [] + + def resume(tag): + assert req_manager.req_to_token_indexs.eq(-1).all() + resumed.append(tag) + model = SimpleNamespace( - torch_memory_saver=SimpleNamespace(resume=lambda tag: events.append(("resume", tag))), - req_manager=SimpleNamespace(init_hold_request_indexs=lambda: events.append(("init_hold_indexes", None))), + torch_memory_saver=SimpleNamespace(resume=resume), + req_manager=req_manager, + att_state_manager=SimpleNamespace(full_att_allocator=allocator), ) - backend = SimpleNamespace(model=model, logger=None) + backend = SimpleNamespace(model=model, logger=Mock()) ops = RlBackendOps(backend) ops._clear_cuda_cache = lambda: None - ops._resume_memory_tags(tags) + success, message = ops.resume_memory_occupation(tags) - init_events = [event for event in events if event[0] == "init_hold_indexes"] - assert bool(init_events) is should_init_hold_indexes - if should_init_hold_indexes: - last_resume_index = max(index for index, event in enumerate(events) if event[0] == "resume") - assert events.index(("init_hold_indexes", None)) > last_resume_index + assert success, message + assert resumed == list(RlBackendOps.MEMORY_TAG_ORDER if tags is None else tags) + expected = list(allocator.hold_token_indices) * 2 if should_init_hold_indexes else [-1] * 8 + assert req_manager.req_to_token_indexs[req_manager.HOLD_REQUEST_ID].tolist() == expected + assert req_manager.req_to_token_indexs[:2].eq(-1).all() + backend.logger.error.assert_not_called() diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py index ed42d27088..784848cb13 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py @@ -1,5 +1,6 @@ from types import SimpleNamespace +import pytest import torch from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import DFlashProposer @@ -114,3 +115,66 @@ def forward(model_input): proposal.schedule_scores, flat_draft_token_probs.reshape(2, block_size)[:, :2].float(), ) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_dflash_hidden_commit_uses_draft_state(device, monkeypatch): + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.attention_state.full_att import GQAFullAttentionState, GQAAttentionConfig + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.common.basemodel.layer_infer.template.transformer_layer_infer_template import TransformerLayerInferTpl + from lightllm.models.qwen3_dflash.model import Qwen3DFlashModel + + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + manager = AttentionStateManager(request_capacity=1) + states = [] + for role, layers in (("main", ((0,), (1,), (2,), (3,))), ("draft0", ((4,),))): + state = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=layers, global_head_num=1, head_dim=16, dtype=torch.float32), + device=device, + ) + state.allocate_storage(8) + state.k_storage.zero_() + state.v_storage.zero_() + manager.register_state(role, state) + states.append(state) + manager.finish_state_registration() + main, draft = states + ids = torch.tensor([1, 3], dtype=torch.int32, device=device) + if device == "cpu": + + def write(logical_layer_id, token_indices, kv): + assert logical_layer_id == 4 + draft.k_storage[0, token_indices.long()] = kv[:, :1] + draft.v_storage[0, token_indices.long()] = kv[:, 1:] + + monkeypatch.setattr(draft, "write_full_att_kv", write) + model = Qwen3DFlashModel.__new__(Qwen3DFlashModel) + model.infer_state_class = InferStateInfo + model.att_state_manager = manager + model.is_mtp_draft_model = True + model.mtp_previous_draft_models = [] + model._cos_cached = model._sin_cached = torch.zeros(8, 16, device=device) + model._select_mem_indexes = lambda _: ids + model.pre_post_weight = None + model.pre_infer = SimpleNamespace(context_forward=lambda _, state, weight: state.mtp_draft_input_hiddens) + + def commit(hidden, state, weight): + TransformerLayerInferTpl._post_cache_kv(SimpleNamespace(layer_num_=4), hidden, state, weight) + return hidden + + model.layers_infer = [SimpleNamespace(context_forward=commit)] + model.trans_layers_weight = [None] + hidden = torch.arange(64, dtype=torch.float32, device=device).view(2, 2, 16) + output = model._decode( + SimpleNamespace( + mtp_draft_input_hiddens=hidden, + batch_size=2, + b_seq_len=torch.tensor([2, 4], device=device), + ) + ) + assert output.logits.shape == (2, 1) + torch.testing.assert_close(draft.k_storage[0, ids.long()], hidden[:, :1]) + torch.testing.assert_close(draft.v_storage[0, ids.long()], hidden[:, 1:]) + assert not main.k_storage.any() and not main.v_storage.any() diff --git a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py new file mode 100644 index 0000000000..3c7ab1a6f5 --- /dev/null +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -0,0 +1,510 @@ +"""CPU service metadata must match the actual model storage and data path.""" + +from __future__ import annotations + +from unit_tests.common.att_resources.model_helpers import full_state_test_fields + +from unit_tests.common.att_resources.model_helpers import full_state_test_buffers + +from typing import TYPE_CHECKING +import dataclasses +import json +from types import SimpleNamespace + +import pytest +import torch + +from unit_tests.common.att_resources.model_helpers import main_linear_state, prepare_model_states +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + +from lightllm.models import Deepseek2TpPartModel, Deepseek3_2TpPartModel, LlamaTpPartModel +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args, get_unique_server_name +from lightllm.utils.kv_cache_utils import calcu_cpu_cache_meta + +if TYPE_CHECKING: + from collections.abc import Callable + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient + from lightllm.utils.kv_cache_utils import CpuKVCacheMeta + from pathlib import Path + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("kv_allocator_environment"), +] + + +@pytest.fixture +def runtime(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + for key, value in { + "LIGHTLLM_UNIQUE_SERVICE_NAME_ID": tmp_path.name, + "LIGHTLLM_CURRENT_RANK_IN_NODE": "0", + "LIGHTLLM_CURRENT_DEVICE_ID": "0", + }.items(): + monkeypatch.setenv(key, value) + yield + get_env_start_args.cache_clear() + get_added_mtp_kv_layer_num.cache_clear() + get_unique_server_name.cache_clear() + calcu_cpu_cache_meta.cache_clear() + + +CASES = [ + ("dense", "None", LlamaTpPartModel), + ("int4", "int4kv", LlamaTpPartModel), + ("int8", "int8kv", LlamaTpPartModel), + ("fp8_head", "fp8kv_sph", LlamaTpPartModel), + ("fp8_tensor", "fp8kv_spt", LlamaTpPartModel), + ("mla", "None", Deepseek2TpPartModel), + ("packed", "None", Deepseek3_2TpPartModel), + ("packed_fp8", "fp8kv_dsa", Deepseek3_2TpPartModel), + ("qwen", "None", Qwen3NextTpPartModel), + ("qwen_head", "fp8kv_sph", Qwen3NextTpPartModel), + ("qwen_tensor", "fp8kv_spt", Qwen3NextTpPartModel), + ("glm", "None", Glm5NextTpPartModel), +] + + +def make_runtime( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], + tp: int, + mtp_layers: int = 0, +): + hybrid = kind.startswith("qwen") or kind == "glm" + config = { + "model_type": "qwen3_5_text" if kind.startswith("qwen") else "glm5_next" if kind == "glm" else "llama", + "architectures": ["LlamaForCausalLM"], + "num_attention_heads": 8, + "num_key_value_heads": 2, + "num_hidden_layers": 4 if hybrid else 2, + "hidden_size": 128, + "head_dim": 16, + "kv_lora_rank": 512, + "qk_rope_head_dim": 64, + "full_attention_interval": 2, + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + "linear_conv_kernel_dim": 3, + "index_kpool": 4, + "index_head_dim": 128, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "linear_attn_config": {"num_heads": 4, "head_dim": 8, "short_conv_kernel_size": 3, "kda_layers": [0, 1, 2]}, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + args = StartArgs( + model_dir=str(tmp_path), + tp=tp, + dp=1, + data_type="bfloat16", + llm_kv_type=encoding, + page_size=4, + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + cpu_cache_token_page_size=8, + enable_cpu_cache=True, + cpu_cache_storage_size=0.001, + linear_att_cache_size=2, + mtp_mode="vanilla_with_att" if mtp_layers else None, + mtp_step=mtp_layers, + ) + if mtp_layers: + draft_config = dict( + config, + num_hidden_layers=1, + n_layer=1, + full_attention_interval=1, + # draft 只有一个自有层，不能沿用 main 的多层类型列表。 + layer_types=["deepseek_sparse_attention"] if kind == "glm" else ["full_attention"], + ) + draft_dir = tmp_path / "draft" + draft_dir.mkdir() + (draft_dir / "config.json").write_text(json.dumps(draft_config)) + args.mtp_draft_model_dir = [str(draft_dir)] * mtp_layers + draft_cls = Glm5NextMTPModel if kind == "glm" else model_cls + monkeypatch.setattr("lightllm.models.get_draft_model_class", lambda *_: draft_cls) + if encoding in ("fp8kv_sph", "fp8kv_spt"): + per_head = encoding == "fp8kv_sph" + calibration_layers = 2 + mtp_layers + shape = [calibration_layers, 4 if per_head else 2] + calibration = { + "num_head": 2, + "num_layers": calibration_layers, + "scales_shape": shape, + "scales": torch.full(shape, 0.125).tolist(), + "q_calibration": { + "num_head": 2, + "scales_shape": [calibration_layers, 2], + "scales": [[0.25, 0.5]] * calibration_layers, + }, + "qmin": -448, + "qmax": 448, + "quant_type": "per_head" if per_head else "per_tensor", + "architectures": config["architectures"][0], + } + path = tmp_path / "calibration.json" + path.write_text(json.dumps(calibration)) + args.kv_quant_calibration_config_path = str(path) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp)) + get_env_start_args.cache_clear() + get_added_mtp_kv_layer_num.cache_clear() + calcu_cpu_cache_meta.cache_clear() + monkeypatch.setattr("lightllm.utils.llm_utils.get_llm_model_class", lambda: model_cls) + meta = calcu_cpu_cache_meta() + managers = [] + for rank in range(tp): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) + model = model_cls.__new__(model_cls) + model.args, model.config = args, config + model.max_total_token_num, model.max_req_num = 16, 1 + model.tp_world_size_, model.tp_rank_ = tp, rank + prepare_model_states(model) + drafts = [] + for step in range(mtp_layers): + draft_cls = Glm5NextMTPModel if kind == "glm" else model_cls + draft = draft_cls.__new__(draft_cls) + draft.is_mtp_draft_model = True + draft.main_model, draft.mtp_previous_draft_models = model, list(drafts) + draft.args = args + draft.config = dict( + config, + num_hidden_layers=1, + n_layer=1, + full_attention_interval=1, + layer_types=["deepseek_sparse_attention"] if kind == "glm" else ["full_attention"], + ) + draft.max_total_token_num, draft.max_req_num = 16, 1 + draft.tp_world_size_, draft.tp_rank_ = tp, rank + prepare_model_states(draft, logical_layer_ids=(config["num_hidden_layers"] + step,)) + drafts.append(draft) + model.att_state_manager.finish_state_registration() + model.att_state_manager.load_full_att_calibration(args=args) + model.att_state_manager.initialize_storage( + args, token_capacity=model.max_total_token_num, mem_fraction=args.mem_fraction + ) + assert not hasattr(model, "mem_manager") + managers.append(model.att_state_manager) + return meta, managers + + +@pytest.mark.parametrize("kind,encoding,model_cls", CASES) +@pytest.mark.parametrize("tp", [1, 4]) +def test_cpu_service_describes_main_and_draft_as_distinct_resources( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], + tp: int, +): + meta, managers = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=model_cls, + tp=tp, + mtp_layers=2, + ) + for resources in managers: + (main,) = resources.get_full_attention_states("main") + (draft,) = resources.get_full_attention_states("draft0") + assert main is not draft + assert {layer for group in main.config.logical_layer_groups for layer in group}.isdisjoint( + layer for group in draft.config.logical_layer_groups for layer in group + ) + assert main.capacity == draft.capacity == resources.full_att_allocator.size + assert {main.state_name, draft.state_name} <= resources.get_cpu_cache_operator().get_page_layouts().keys() + layout_end = max(layout.byte_end for layout in resources.get_cpu_cache_operator().get_page_layouts().values()) + assert (layout_end + 15) // 16 * 16 == meta.page_bytes + assert resources.get_cpu_cache_operator().get_page_bytes() == meta.page_bytes + + +@pytest.mark.parametrize("kind,encoding,model_cls", CASES) +@pytest.mark.parametrize("tp", [1, 4]) +def test_cpu_service_layout_and_roundtrip_all_encodings( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + cpu_cache_factory: Callable[[CpuKVCacheMeta], CpuKvCacheClient], + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], + tp: int, +): + meta, managers = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=model_cls, + tp=tp, + ) + # No byte-size formulas duplicated in the test: check the service's actual + # meta-device descriptors against every real model's layout. + for resources in managers: + layout = next(iter(resources.get_cpu_cache_operator().get_page_layouts().values())) + assert 0 <= layout.byte_start < layout.byte_end <= meta.page_bytes + assert meta.page_bytes == resources.get_cpu_cache_operator().get_page_bytes() + meta.page_num = 3 + client = cpu_cache_factory(meta) + buffer = client.cpu_kv_cache_tensor + expected = [] + for resources in managers: + (full,) = resources.get_full_attention_states("main") + for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): + if "token" in field.axes: + field.tensor.fill_(7) + kwargs = {} + if resources.bounded_att_big_page_allocator is not None: + linear = main_linear_state(resources) + linear.conv_state.fill_(3) + linear.ssm_state.fill_(5) + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + resources.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[slot], + ) + kwargs = dict(big_page_slot_ids=[slot]) + torch.cuda.synchronize() + expected.append( + { + name: f.tensor[:, :8].contiguous().view(torch.uint8).clone() + for name, f in full_state_test_fields(full, full_state_test_buffers(full)).items() + if "token" in f.axes + } + ) + resources.get_cpu_cache_operator().write_req( + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + cpu_page_buffers=buffer, + cpu_page_indices=[2], + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + **kwargs, + ) + torch.cuda.synchronize() + for resources, fields in zip(managers, expected): + (full,) = resources.get_full_attention_states("main") + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer[:, 8:16].zero_() + kwargs = {} + if resources.bounded_att_big_page_allocator is not None: + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + kwargs = dict( + big_page_slot_ids=[slot], + req_idx=1, + ) + resources.get_cpu_cache_operator().read_req( + cpu_page_buffers=buffer, + cpu_page_indices=[2], + token_indices=torch.arange( + start=8, + end=16, + dtype=torch.long, + device="cuda", + ), + state_token_len=8, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + **kwargs, + ) + if resources.bounded_att_big_page_allocator is not None: + assert (main_linear_state(resources).conv_state[:, 1] == 3).all() + assert (main_linear_state(resources).ssm_state[:, 1] == 5).all() + for name, field in full_state_test_fields(full, full_state_test_buffers(full)).items(): + if "token" in field.axes: + assert torch.equal(field.tensor[:, 8:16].contiguous().view(torch.uint8), fields[name]) + assert not buffer[:2].any() + + +@pytest.mark.parametrize("hybrid", [False, True]) +def test_production_offload_preserves_ready_pages_without_host_sync( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + cpu_cache_factory: Callable[[CpuKVCacheMeta], CpuKvCacheClient], + hybrid: bool, +): + from lightllm.server.multi_level_kv_cache import CacheTier + from lightllm.server.router.model_infer.infer_batch import InferReq + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common + + kind, encoding, cls = ("qwen", "None", Qwen3NextTpPartModel) if hybrid else ("int4", "int4kv", LlamaTpPartModel) + meta, (resources,) = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind=kind, + encoding=encoding, + model_cls=cls, + tp=1, + ) + meta.page_num = 2 + client = cpu_cache_factory(meta) + pages, _ = client.allocate_pages([11, 22], disk_offload_enable=False) + client.cpu_kv_cache_tensor[pages[0]].fill_(123) + client.update_pages_status_to_ready([pages[0]], deref=False) + slots = resources.bounded_att_big_page_allocator.alloc(2) if hybrid else None + write_batches = [] + write_to_cpu = resources.get_cpu_cache_operator().write_req + + def capture_write(*args: object, **kwargs: object) -> None: + write_batches.append((kwargs["cpu_page_indices"], kwargs["cpu_page_ready_flags"])) + write_to_cpu(*args, **kwargs) + + monkeypatch.setattr(resources.get_cpu_cache_operator(), "write_req", capture_write) + if hybrid: + # 当前 bounded 页仍由主机复制，测试先准备好源快照。 + # TODO：CPU→CPU kernel 落地后，再覆盖尚未完成 D2H 的快照依赖。 + pool = main_linear_state(resources).big_pages_pool + for slot in slots: + pool.conv_state[slot].fill_(3) + pool.ssm_state[slot].fill_(5) + producer = torch.cuda.Stream() + producer.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(producer): + # 延迟当前生产流的 KV 写入；当前实现不在此处隐式插入 wait_stream， + # 生产者与搬运流的依赖由调用方负责组织。 + torch.cuda._sleep(10_000_000) + (full,) = resources.get_full_attention_states("main") + for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): + if "token" in field.axes: + field.tensor.fill_(7) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=16, + cache_tiers=(CacheTier.CPU,), + tail_small_page_buffer_id=None, + hybrid_len_to_big_page_id=dict(zip([8, 16], slots)) if hybrid else {}, + shared_kv_node=None, + cpu_cache_task_status=InferReq._CpuCacheTaskStatus.NOT_STARTED, + shm_req=SimpleNamespace( + token_hash_list=SimpleNamespace(get_all=lambda: [11, 22]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [8, 16]), + ), + ) + module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) + module.args = get_env_start_args() + module.cpu_cache_client = client + module.gloo_group = None + module.backend = SimpleNamespace( + is_master_in_dp=True, + is_hybrid_att_model=hybrid, + dp_world_size=1, + model=SimpleNamespace( + att_state_manager=resources, + req_manager=SimpleNamespace(req_to_token_indexs=torch.arange(16, device="cuda").view(1, -1)), + ), + ) + monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) + monkeypatch.setattr(common.g_infer_context, "is_hybrid_att_model", hybrid) + monkeypatch.setattr(common.g_infer_context, "radix_cache", None) + + def reject_host_sync(*args: object, **kwargs: object) -> None: + pytest.fail("CPU offload must use stream dependencies instead of host synchronization") + + copy_stream = torch.cuda.Stream() + waited_streams = [] + wait_stream = torch.cuda.Stream.wait_stream + + def record_wait(stream: torch.cuda.Stream, other: torch.cuda.Stream) -> None: + if stream == copy_stream: + waited_streams.append(other) + wait_stream(stream, other) + + with monkeypatch.context() as patch: + patch.setattr(torch.cuda.Stream, "synchronize", reject_host_sync) + patch.setattr(torch.cuda.Stream, "wait_stream", record_wait) + task = module.write_req_to_cpu(req, copy_stream) + assert waited_streams == [] + task.sync_event.synchronize() + assert write_batches == [(pages, [True, False])] + assert req.cpu_cache_task_status.is_running() + assert (client.cpu_kv_cache_tensor[pages[0]] == 123).all() + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer[:, :8].zero_() + kwargs = {} + if hybrid: + # 写出已完成，释放测试源快照，给读回目标腾出普通大页槽位。 + req.hybrid_len_to_big_page_id.clear() + resources.bounded_att_big_page_allocator.free(slots) + destination = resources.bounded_att_big_page_allocator.alloc(1)[0] + kwargs = dict(big_page_slot_ids=[destination], req_idx=1) + resources.get_cpu_cache_operator().read_req( + cpu_page_buffers=client.cpu_kv_cache_tensor, + cpu_page_indices=[pages[1]], + token_indices=torch.arange(8, dtype=torch.long, device="cuda"), + state_token_len=8, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + **kwargs, + ) + for field in full_state_test_fields(full, full_state_test_buffers(full)).values(): + if "token" in field.axes: + assert (field.tensor[:, :8] == 7).all() + if hybrid: + assert (main_linear_state(resources).conv_state[:, 1] == 3).all() + assert (main_linear_state(resources).ssm_state[:, 1] == 5).all() + + +def test_disk_preserves_resource_pages( + runtime: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + cpu_cache_factory: Callable[[CpuKVCacheMeta], CpuKvCacheClient], +): + import time + + pytest.importorskip("light_mem") + from lightllm.server.multi_level_kv_cache.disk_cache_worker import DiskCacheWorker + + meta, (resources,) = make_runtime( + monkeypatch=monkeypatch, + tmp_path=tmp_path, + kind="packed_fp8", + encoding="fp8kv_dsa", + model_cls=Deepseek3_2TpPartModel, + tp=1, + ) + meta.page_num = 2 + client = cpu_cache_factory(meta) + pages, _ = client.allocate_pages([11, 22], disk_offload_enable=True) + (full,) = resources.get_full_attention_states("main") + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.random_(0, 256) + resources.get_cpu_cache_operator().write_req( + token_indices=torch.arange(16, dtype=torch.long, device="cuda"), + cpu_page_buffers=client.cpu_kv_cache_tensor, + cpu_page_indices=pages, + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + ) + torch.cuda.synchronize() + expected = client.cpu_kv_cache_tensor.clone() + worker = DiskCacheWorker(0.125, client, str(tmp_path / "disk")) + task = worker.service.create(hash_128s=[11, 22], kv_page_indexer=torch.tensor(pages, dtype=torch.int32), mode="w") + deadline = time.monotonic() + 10 + while not task.ready(): + assert time.monotonic() < deadline, "resource page disk write did not finish" + time.sleep(0.01) + assert all(worker.service.query([11, 22])) + client.cpu_kv_cache_tensor.zero_() + assert worker.load_pages([11, 22], pages) + assert torch.equal(client.cpu_kv_cache_tensor, expected) + for storage_buffer in full_state_test_buffers(full).values(): + storage_buffer.zero_() + resources.get_cpu_cache_operator().read_req( + cpu_page_buffers=client.cpu_kv_cache_tensor, + cpu_page_indices=pages, + token_indices=torch.arange(16, dtype=torch.long, device="cuda"), + layouts=resources.get_cpu_cache_operator().get_page_layouts(), + ) + assert any((tmp_path / "disk").glob("cache_file*")) 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 3111abce81..07ab123d09 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 @@ -1,3 +1,4 @@ +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources import dataclasses import json from types import SimpleNamespace @@ -5,16 +6,27 @@ 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.state_cache_manager import LinearAttCacheConfig +from lightllm.common.attention_state.allocators import FullAttKVAllocator + +from lightllm.common.attention_state.full_att import ( + GQAFullAttentionState, + GQAAttentionConfig, +) +from lightllm.common.req_manager import ReqManager +from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule +from lightllm.server.multi_level_kv_cache import CacheTier +from lightllm.server.router.model_infer.infer_batch import InferReq +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.server.router.model_infer.infer_batch import g_infer_context from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name from lightllm.utils.kv_cache_utils import CpuKVCacheMeta -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.usefixtures("kv_allocator_environment"), +] @pytest.fixture(params=[False, True], ids=["default_chunks", "whole_mm_blocks"]) @@ -37,26 +49,42 @@ def linear_cache(monkeypatch, tmp_path, request): monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) get_env_start_args.cache_clear() get_unique_server_name.cache_clear() - config = LinearAttCacheConfig( + config = Qwen3NextResourceConfig( 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=3, conv_state_dtype=torch.bfloat16, ssm_state_dtype=torch.float32, + full_attention_interval=2, + all_layer_num=2, layer_types=("linear_attention", "full_attention"), ) - mem = Qwen3NextMemManager(384, torch.bfloat16, 1, 8, 1, config) - requests = ReqManagerForMamba(2, 512, mem, config) - requests.create_small_page_cache_manager(2) + full = GQAFullAttentionState( + GQAAttentionConfig( + logical_layer_groups=tuple((layer,) for layer in config.get_full_att_layer_indices()), + global_head_num=1, + head_dim=8, + dtype=torch.bfloat16, + ), + page_size=args.page_size, + ) + full.allocate_storage(384) + requests = ReqManager(2, 512) + resources = build_hybrid_resources( + config=config, + request_capacity=3, + mtp_step=0, + full_state=full, + full_allocator=FullAttKVAllocator(full.capacity, full.page_size), + big_page_capacity=5, + small_page_capacity=2, + ) + monkeypatch.setattr(g_infer_context, "att_state_manager", resources) monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) monkeypatch.setattr(g_infer_context, "req_manager", requests) monkeypatch.setattr( @@ -65,52 +93,60 @@ 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, ), ) - yield mem, requests, config + yield full, resources, requests, config get_env_start_args.cache_clear() get_unique_server_name.cache_clear() @pytest.mark.parametrize("token_num", [320, 384]) -def test_cpu_roundtrip_keeps_kv_continuous_and_restores_only_saved_states( - monkeypatch, linear_cache, cpu_cache_factory, token_num +def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( + monkeypatch, + linear_cache, + cpu_cache_factory, + token_num, ): from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache 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 + full, resources, requests, 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 + linear = resources.main_gdn_conv_bf16_ssm_fp32_state + big = linear.big_pages_pool + small = linear.small_pages_pool big_states = {} for end in range(128, token_num + 1, 128): - slot = big.alloc_one_state_cache() + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] big_states[end] = slot - for tensor in big.get_state_cache(slot): + for tensor in (big.conv_state[slot], big.ssm_state[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 = resources.bounded_att_small_page_allocator.alloc(1)[0] + for tensor in (small.conv_state[tail_slot], small.ssm_state[tail_slot]): tensor.normal_() - final_states = requests.small_page_buffers.get_state_cache(tail_slot) + final_states = (small.conv_state[tail_slot], small.ssm_state[tail_slot]) else: - final_states = big.get_state_cache(big_states[token_num]) + final_states = (big.conv_state[big_states[token_num]], big.ssm_state[big_states[token_num]]) expected_states = [tensor.clone() for tensor in final_states] - expected_big_states = {end: [t.clone() for t in big.get_state_cache(slot)] for end, slot in big_states.items()} - mem.kv_buffer.normal_() - expected_kv = mem.kv_buffer[:, :token_num].clone() + expected_big_states = { + end: [t.clone() for t in (big.conv_state[slot], big.ssm_state[slot])] for end, slot in big_states.items() + } + full.k_storage.normal_() + full.v_storage.normal_() + expected_k = full.k_storage[:, :token_num].clone() + expected_v = full.v_storage[:, :token_num].clone() client = cpu_cache_factory( - CpuKVCacheMeta(3, 1, 1, 1, config.get_cpu_cache_big_page_bytes(), torch.uint8, 0, torch.uint8), + CpuKVCacheMeta(3, resources.get_cpu_cache_operator().get_page_bytes()), ) - pages, _ = client.allocate_pages([11, 22, 33], disk_offload_enable=False) indexes = torch.arange(token_num, dtype=torch.int32, device="cuda") - page_ids = torch.arange(3, dtype=torch.int32, device="cuda") req = SimpleNamespace( req_idx=0, + cache_tiers=(CacheTier.CPU,), + cpu_cache_task_status=InferReq._CpuCacheTaskStatus.NOT_STARTED, cur_kv_len=token_num, hold_kv_len=token_num, hybrid_cache_len=token_num, @@ -123,27 +159,42 @@ def test_cpu_roundtrip_keeps_kv_continuous_and_restores_only_saved_states( input_len=token_num + 1, disk_prompt_cache_len=0, cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: pages), + token_hash_list=SimpleNamespace(get_all=lambda: [11, 22, 33]), token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256, token_num]), ), ) - mem.operator.offload_gpu_kv_to_cpu_cache(indexes, page_ids, torch.zeros_like(page_ids), client, req) - torch.cuda.synchronize() + # 通过接收 InferReq 的生产入口完成页来源查询、CPU 分配与快照写出。 + requests.req_to_token_indexs[0, :token_num] = indexes module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.args = get_env_start_args() module.cpu_cache_client = client + module.gloo_group = None + module.backend = SimpleNamespace( + is_master_in_dp=True, + is_hybrid_att_model=True, + dp_world_size=1, + model=SimpleNamespace(req_manager=requests, att_state_manager=resources), + ) + monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) + task = module.write_req_to_cpu(req, torch.cuda.Stream()) + assert task is not None + task.sync_event.synchronize() + pages = task.page_indexes.tolist() client.update_pages_status_to_ready(pages, deref=False) - big.free_state_cache(list(big_states.values())) + resources.bounded_att_big_page_allocator.free(list(big_states.values())) if tail_slot is not None: - requests.small_page_buffers.free_state_cache([tail_slot]) + resources.bounded_att_small_page_allocator.free([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 - mem.kv_buffer.zero_() - big.conv_state_cache.buffer.zero_() - big.ssm_state_cache.buffer.zero_() + full.k_storage.zero_() + full.v_storage.zero_() + big.conv_state.zero_() + big.ssm_state.zero_() def allocate(req, count): - allocated = mem.alloc(count) + allocated = resources.full_att_allocator.alloc(count) requests.req_to_token_indexs[req.req_idx, :count] = allocated.cuda() req.hold_kv_len = count return allocated @@ -152,7 +203,8 @@ 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(req_manager=requests, att_state_manager=resources), + is_hybrid_att_model=True, ) module.init_sync_group = None module.need_sync_compute_stream = lambda: False @@ -161,27 +213,36 @@ def allocate(req, count): module.load_cpu_cache_to_reqs([req]) assert req.cur_kv_len == req.hold_kv_len == token_num - assert torch.equal(mem.kv_buffer[:, :token_num], expected_kv) + restored_indices = requests.req_to_token_indexs[req.req_idx, :token_num] + assert torch.equal(full.k_storage[:, restored_indices], expected_k) + assert torch.equal(full.v_storage[:, restored_indices], expected_v) 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. + # Every published CPU page has a complete, length-checked linear state. for end, slot in req.hybrid_len_to_big_page_id.items(): - for actual, expected in zip(big.get_state_cache(slot), expected_big_states[end]): + for actual, expected in zip((big.conv_state[slot], big.ssm_state[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(linear.conv_state[:, 0].cpu(), expected_states[0]) + torch.testing.assert_close(linear.ssm_state[:, 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) + # Model data validity and multimodal resume boundaries are separate checks. + tree = HybridAttPagedRadixCache( + total_token_num=384, + rank_in_node=103, + hash_page_size=64, + big_page_num=2, + att_state_manager=resources, + ) cache_len = max(req.hybrid_len_to_big_page_id) tree.insert( - torch.arange(cache_len), + key=torch.arange(cache_len), + value=restored_indices[:cache_len].cpu(), block_hashs=list(range(1, cache_len // 64 + 1)), block_state_idxs=[None] * (cache_len // 64), len_to_big_page_id=req.hybrid_len_to_big_page_id, ) for length, expected in [(128, 128 if first_page_has_state else 0), (256, 256)]: node, matched, _ = tree.match_prefix( - torch.arange(length), + key=torch.arange(length), block_hashs=list(range(1, length // 64 + 1)), update_refs=True, mm_input_spans=req.mm_input_spans, @@ -189,9 +250,107 @@ def allocate(req, count): assert matched == expected tree.dec_node_ref_counter(node) req.shared_kv_node, _, _ = tree.match_prefix( - torch.arange(cache_len), + key=torch.arange(cache_len), block_hashs=list(range(1, cache_len // 64 + 1)), update_refs=True, mm_input_spans=req.mm_input_spans, ) tree.dec_node_ref_counter(req.shared_kv_node) + + +@pytest.mark.parametrize("shared_prefix", [False, True]) +@pytest.mark.parametrize("allocated_pages", [0, 1, 2]) +def test_offload_resolves_only_allocated_pages_and_requires_big_slots( + monkeypatch, + linear_cache, + shared_prefix, + allocated_pages, +): + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common + + full, resources, requests, config = linear_cache + slots = resources.bounded_att_big_page_allocator.alloc(allocated_pages) if allocated_pages else [] + shared_slots = [slots[0]] if shared_prefix and slots else [] + monkeypatch.setattr( + g_infer_context, "radix_cache", SimpleNamespace(get_big_page_ids_by_node=lambda node: shared_slots) + ) + big_slots = {128 * (i + 1): slot for i, slot in enumerate(slots) if not shared_prefix or i > 0} + req = SimpleNamespace( + req_idx=0, + shared_kv_node=None, + hybrid_len_to_big_page_id=big_slots, + tail_small_page_buffer_id=None, + cur_kv_len=256, + cache_tiers=(CacheTier.CPU,), + shm_req=SimpleNamespace( + token_hash_list=SimpleNamespace(get_all=lambda: [11, 22]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256]), + ), + ) + requested_hashes = [] + writes = [] + + def allocate_pages(hashes, disk_offload_enable): + requested_hashes.append(list(hashes)) + return list(range(allocated_pages)) + [-1] * (len(hashes) - allocated_pages), [False] * len(hashes) + + def write_to_cpu( + token_indices, + cpu_page_buffers, + cpu_page_indices, + **kwargs, + ): + writes.append( + (token_indices.numel(), cpu_page_indices, kwargs["big_page_slot_ids"], kwargs["small_page_slot_ids"]) + ) + + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.args = get_env_start_args() + module.gloo_group = None + module.backend = SimpleNamespace( + is_master_in_dp=True, + is_hybrid_att_model=True, + dp_world_size=1, + model=SimpleNamespace(att_state_manager=resources, req_manager=requests), + ) + module.cpu_cache_client = SimpleNamespace( + lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), + allocate_pages=allocate_pages, + cpu_kv_cache_tensor=torch.empty((2, 1), dtype=torch.uint8), + ) + monkeypatch.setattr(resources.get_cpu_cache_operator(), "write_req", write_to_cpu) + broadcasts = [] + monkeypatch.setattr(common.dist, "broadcast_object_list", lambda data, **kwargs: broadcasts.append(data)) + task = module.write_req_to_cpu(req, torch.cuda.Stream()) + if allocated_pages == 0: + assert task is None + assert req.cpu_cache_task_status == InferReq._CpuCacheTaskStatus.FINISHED + assert requested_hashes == [[11, 22]] + assert broadcasts == [[0]] + assert writes == [] + return + task.sync_event.synchronize() + assert broadcasts == [ + [ + { + "item_size": allocated_pages, + "page_list": list(range(allocated_pages)), + "ready_list": [False] * allocated_pages, + } + ] + ] + # 即使 CPU 只分配一页，也申请完整可缓存范围；未选中的 256 位置不要求提供槽位。 + assert requested_hashes == [[11, 22]] + assert writes == [(allocated_pages * 128, list(range(allocated_pages)), slots, [])] + + # 实际选中页缺失时直接暴露错误，不缩短写出范围，也不提交部分搬运。 + missing_length = allocated_pages * 128 + if shared_prefix and allocated_pages == 1: + shared_slots.clear() + else: + del req.hybrid_len_to_big_page_id[missing_length] + with pytest.raises(KeyError) as error: + module.write_req_to_cpu(req, torch.cuda.Stream()) + assert error.value.args == (missing_length,) + assert requested_hashes == [[11, 22], [11, 22]] + assert len(writes) == 1 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 f40dd1b98c..9779145cf4 100644 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -32,6 +32,57 @@ def test_chunk_ids_and_lengths_use_the_same_complete_boundaries(hybrid, spans): req.cur_kv_len = end +@pytest.mark.usefixtures("kv_allocator_environment") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") +def test_bounded_cache_miss_resets_runtime_and_preserves_pages(monkeypatch): + from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context + from unit_tests.common.att_resources.test_manager_helpers import composition, init_source + + monkeypatch.setattr( + "lightllm.common.attention_state.bounded_att.utils.get_env_start_args", + lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), + ) + resources = composition(device="cuda") + slots = init_source(resources) + resources.reset_bounded_att_state([1]) + resources.restore_bounded_att_runtime_from_pages( + req_indices=[1], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[slots[-1]], + ) + for state in resources.get_bounded_attention_states(): + state.mtp_state_index.fill_(2) + + monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) + monkeypatch.setattr(g_infer_context, "radix_cache", None) + monkeypatch.setattr(g_infer_context, "att_state_manager", resources) + req = InferReq.__new__(InferReq) + req.args = SimpleNamespace( + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + max_req_total_len=16, + page_size=4, + ) + req.req_idx = 1 + req.cur_kv_len = req.cur_output_len = 0 + req.sampling_param = SimpleNamespace(disable_prompt_cache=False) + req.shm_req = SimpleNamespace(input_len=1, hybrid_token_hash_list=SimpleNamespace(get_all=lambda: [], size=0)) + req._hybrid_match_radix_cache() + + for state in resources.get_bounded_attention_states(): + assert not state.conv_state[:, 1].any() + assert not state.ssm_state[:, 3:6].any() + assert state.mtp_state_index.tolist() == [2, 0, 2] + assert state.conv_state[:, 0].any() + pool = state.big_pages_pool + expected_conv, expected_ssm = (23, 47) if state is resources.main_kda_conv_fp32_ssm_fp32_state else (24, 48) + assert (pool.conv_state[slots[-1]] == expected_conv).all() + assert (pool.ssm_state[slots[-1]] == expected_ssm).all() + assert resources.bounded_att_big_page_allocator.free_count == 1 + assert req.cur_kv_len == req.hold_kv_len == req.shm_req.shm_cur_kv_len == 0 + + @pytest.mark.parametrize( "input_len,big_page_num,spans,expected_big,expected_small", [ @@ -49,7 +100,11 @@ def test_chunk_ids_and_lengths_use_the_same_complete_boundaries(hybrid, spans): ], ) def test_hybrid_chunks_save_state_only_at_original_boundaries( - input_len, big_page_num, spans, expected_big, expected_small + input_len, + big_page_num, + spans, + expected_big, + expected_small, ): from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext @@ -74,23 +129,15 @@ def test_hybrid_chunks_save_state_only_at_original_boundaries( context = InferenceContext() 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), - free_one_small_page_buffer=lambda: None, + context.radix_cache = SimpleNamespace(free_one_small_page_buffer=lambda: None) + context.att_state_manager = SimpleNamespace( + bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: [next(slots)]), + bounded_att_small_page_allocator=SimpleNamespace(alloc=lambda count: [7]), + save_bounded_att_runtime_to_pages=lambda req_indices, state_token_lens, page_kind, page_slot_ids: ( + big_snapshots if page_kind == "big" else small_snapshots + ).extend(state_token_lens), ) - def save_big_page_states(b_req_idx, reqs, buffer_indexes): - assert b_req_idx.tolist() == [req.req_idx for req in reqs] - for req, buffer_idx in zip(reqs, buffer_indexes): - if buffer_idx != -1: - big_snapshots.append(req.get_chuncked_input_token_len()) - - def save_state(req, buffer_idx, state_cache_manager): - small_snapshots.append(req.get_chuncked_input_token_len()) - - context.req_manager = SimpleNamespace(save_big_page_states=save_big_page_states, save_state=save_state) - while req.cur_kv_len < req.get_cur_total_len(): context.save_hybrid_state_to_cache(torch.tensor([0]), [req]) req.cur_kv_len = req.get_chuncked_input_token_len() @@ -107,7 +154,13 @@ def save_state(req, buffer_idx, state_cache_manager): [(1, True, 1201, 1024), (1, False, 700, 512), (128, True, 0, 1024), (128, False, 640, 512)], ) def test_full_attention_gpu_restore_stays_outside_image( - monkeypatch, tmp_path, cpu_cache_factory, page_size, same_image, gpu_end, cpu_end + monkeypatch, + tmp_path, + cpu_cache_factory, + page_size, + same_image, + gpu_end, + cpu_end, ): from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache @@ -144,7 +197,7 @@ def test_full_attention_gpu_restore_stays_outside_image( assert radix_cache.get_refed_tokens_num() == gpu_end hashes = compute_token_list_hash(source, 256) - client = cpu_cache_factory(CpuKVCacheMeta(len(hashes), 256, 1, 1, 8, torch.bfloat16, 0, torch.bfloat16)) + client = cpu_cache_factory(CpuKVCacheMeta(len(hashes), 4096)) pages, _ = client.allocate_pages(hashes, False) client.update_pages_status_to_ready(pages) manager = MultiLevelKVCacheManager.__new__(MultiLevelKVCacheManager) @@ -208,6 +261,7 @@ def test_mm_chunking_keeps_the_original_state_cache_position(monkeypatch, disabl ) req._init_all_state() assert req.hybrid_cache_len == 768 + assert req.mm_input_spans == ([(450, 600), (700, 900)] if disable_chunked_mm_input else []) @pytest.mark.parametrize("hybrid", [False, True]) @@ -243,7 +297,7 @@ def test_cpu_cache_match_and_load_use_a_complete_restore_point( from lightllm.utils.kv_cache_utils import CpuKVCacheMeta ends = list(range(2048, load_end + 1, 2048)) - client = cpu_cache_factory(CpuKVCacheMeta(max(1, len(ends)), 1, 1, 1, 16, torch.uint8, 0, torch.uint8)) + client = cpu_cache_factory(CpuKVCacheMeta(max(1, len(ends)), 16)) page_ids, _ = client.allocate_pages(list(range(10, 10 + len(ends))), disk_offload_enable=False) client.update_pages_status_to_ready(page_ids, deref=False) loaded = [] @@ -267,16 +321,21 @@ def allocate(req, count): return torch.arange(req.cur_kv_len, req.cur_kv_len + count, dtype=torch.int32) module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.args = SimpleNamespace(page_size=256) + module.args = SimpleNamespace(page_size=256, cpu_cache_token_page_size=256) module.backend = SimpleNamespace( is_master_in_dp=True, + is_hybrid_att_model=hybrid, _alloc_req_kv_mem=allocate, model=SimpleNamespace( req_manager=SimpleNamespace(req_to_token_indexs=table), - mem_manager=SimpleNamespace( - operator=SimpleNamespace( - load_cpu_cache_to_gpu=lambda **kwargs: loaded.append(req.cur_kv_len), - ) + att_state_manager=SimpleNamespace( + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda cpu_page_buffers, cpu_page_indices, token_indices, **kwargs: loaded.append( + req.hold_kv_len + ), + ), + bounded_att_big_page_allocator=SimpleNamespace(alloc=lambda count: list(range(count))), ), ), ) @@ -288,6 +347,7 @@ def allocate(req, count): monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None)) + req.hybrid_len_to_big_page_id = {} module.load_cpu_cache_to_reqs([req]) assert req.cur_kv_len == req.hold_kv_len == expected assert loaded == ([expected] if expected > loaded_start else []) diff --git a/unit_tests/server/router/test_kv_capacity.py b/unit_tests/server/router/test_kv_capacity.py new file mode 100644 index 0000000000..9626b03ea8 --- /dev/null +++ b/unit_tests/server/router/test_kv_capacity.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from lightllm.common.attention_state.allocators import full_att_kv_allocator + + +pytestmark = pytest.mark.usefixtures("kv_allocator_environment") + + +@pytest.mark.parametrize( + "tp,dp,nnodes,selected_ranks", + [(4, 1, 1, [0]), (8, 4, 1, [0, 2, 4, 6]), (8, 4, 2, [0, 2]), (8, 8, 2, [0, 1, 2, 3]), (8, 1, 2, [0])], +) +@pytest.mark.parametrize("reader_first", [False, True]) +def test_router_reads_each_local_dp_groups_allocator( + monkeypatch: pytest.MonkeyPatch, + tp: int, + dp: int, + nnodes: int, + selected_ranks: list[int], + reader_first: bool, +): + monkeypatch.setattr( + full_att_kv_allocator, "get_env_start_args", lambda: SimpleNamespace(tp=tp, dp=dp, nnodes=nnodes) + ) + reader = full_att_kv_allocator.KVCapacityReader() if reader_first else None + allocators = [] + for rank in range(tp // nnodes): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", str(rank)) + allocator = FullAttKVAllocator(16 + 4 * rank, 4, pin_memory=False) + assert allocator.shared_can_use_token_num.get_value() == allocator.size + allocator.alloc(4) + allocators.append(allocator) + reader = reader or full_att_kv_allocator.KVCapacityReader() + assert len(reader.counters) == len(selected_ranks) + for dp_rank, rank in enumerate(selected_ranks): + allocator = allocators[rank] + initial = allocator.size - 4 + assert reader.get_available_tokens(dp_rank) == initial + tokens = allocator.alloc(4) + assert reader.get_available_tokens(dp_rank) == initial - 4 + allocator.free(tokens) + assert reader.get_available_tokens(dp_rank) == initial + allocator.free_all() + assert reader.get_available_tokens(dp_rank) == allocator.size + + +def test_failed_page_operations_leave_shared_capacity_unchanged(): + allocator = FullAttKVAllocator(16, 4, pin_memory=False) + counter = allocator.shared_can_use_token_num + tokens = allocator.alloc(16) + with pytest.raises(AssertionError): + allocator.alloc(4) + with pytest.raises(AssertionError): + allocator.free(tokens[:3]) + assert counter.get_value() == allocator.can_use_mem_size == 0 + allocator.free(tokens) + assert counter.get_value() == allocator.can_use_mem_size == 16 diff --git a/unit_tests/server/test_dp_request_capacity.py b/unit_tests/server/test_dp_request_capacity.py index 1e63550f45..800f0cdda3 100644 --- a/unit_tests/server/test_dp_request_capacity.py +++ b/unit_tests/server/test_dp_request_capacity.py @@ -45,7 +45,6 @@ def test_reject_invalid_global_capacity(capacity): def _make_dp_queue(monkeypatch, balancer): monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: True) - monkeypatch.setattr("lightllm.server.router.req_queue.base_queue.get_fixed_kv_len", lambda: 0) args = StartArgs( running_max_req_size=32, dp=8, diff --git a/unit_tests/server/test_pd_start_args.py b/unit_tests/server/test_pd_start_args.py index 5f0c3e450f..3109fe11b9 100644 --- a/unit_tests/server/test_pd_start_args.py +++ b/unit_tests/server/test_pd_start_args.py @@ -91,28 +91,24 @@ def test_disabled_cpu_cache_keeps_page_size_unset(startup_args): assert startup_args.cpu_cache_token_page_size is None -def test_other_memory_manager_selection_uses_default(monkeypatch): - from types import SimpleNamespace - import lightllm.models as models - from lightllm.common.kv_cache_mem_manager import MemoryManager - from lightllm.common.kv_cache_mem_manager import mem_utils - - for name in ("DeepseekV4TpPartModel", "Deepseek3_2TpPartModel", "Deepseek2TpPartModel"): - monkeypatch.setitem(models.__dict__, name, type(name, (), {})) - monkeypatch.setattr(mem_utils, "get_llm_model_class", lambda: type("Llama", (), {})) - monkeypatch.setattr( - mem_utils, "get_env_start_args", lambda: SimpleNamespace(model_dir="unused", llm_kv_type="None") +def test_default_attention_factory_uses_gqa(): + import torch + from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState, create_gqa_state + + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=4, + head_dim=128, + dtype=torch.bfloat16, + ), + device="meta", ) - mem_utils.select_mem_manager_class.cache_clear() - try: - assert mem_utils.select_mem_manager_class() is MemoryManager - finally: - mem_utils.select_mem_manager_class.cache_clear() + assert isinstance(state, GQAFullAttentionState) def test_dsv4_cpu_cache_factory_resolves_packed_layout_without_overriding_kv_type(startup_args, monkeypatch, tmp_path): - from lightllm.common.kv_cache_mem_manager import DeepseekV4MemoryManager - from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class + from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts from lightllm.utils.envs_utils import get_env_start_args, get_added_mtp_kv_layer_num, get_llm_data_type from lightllm.utils.llm_utils import get_llm_model_class from lightllm.utils.kv_cache_utils import calcu_cpu_cache_meta @@ -143,18 +139,21 @@ def test_dsv4_cpu_cache_factory_resolves_packed_layout_without_overriding_kv_typ get_added_mtp_kv_layer_num, get_llm_data_type, get_llm_model_class, - select_mem_manager_class, calcu_cpu_cache_meta, ) for function in cached_functions: function.cache_clear() try: - assert select_mem_manager_class() is DeepseekV4MemoryManager + layouts = build_cpu_cache_layouts() + assert any("compressed_c4" in name for name in layouts) + assert any("compressed_c128" in name for name in layouts) + assert any("swa_w128_fp8" in name for name in layouts) + assert any("compressor_kv_w8_s4" in name for name in layouts) + assert not any("compressor_kv_w128" in name for name in layouts) meta = calcu_cpu_cache_meta() assert meta.page_num > 0 - assert meta.token_page_size == meta.layer_num == meta.num_heads == 1 assert args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - assert meta.calcu_one_page_size() == meta.head_dim + assert meta.page_bytes == max(layout.byte_end for layout in layouts.values()) assert get_env_start_args().llm_kv_type == "None" finally: for function in cached_functions: diff --git a/unit_tests/server/test_tokenizer.py b/unit_tests/server/test_tokenizer.py index 2accf0ce89..aa1c6d6273 100644 --- a/unit_tests/server/test_tokenizer.py +++ b/unit_tests/server/test_tokenizer.py @@ -1,5 +1,6 @@ """Tokenizer entry points must not eagerly load unrelated model implementations.""" +import json import os from pathlib import Path import subprocess @@ -13,6 +14,47 @@ ROOT = Path(__file__).resolve().parents[2] +@pytest.mark.parametrize("model_type", ["deepseek_v32", "deepseek_v3"]) +def test_deepseek_bytelevel_checkpoint_preserves_whitespace(tmp_path: Path, model_type: str) -> None: + from tokenizers import Tokenizer, decoders, pre_tokenizers, trainers + from tokenizers.models import BPE + from transformers import PreTrainedTokenizerFast + from lightllm.server.tokenizer import get_tokenizer + + text = "Question: What is 2 + 3?\nAnswer: #### 5" + backend = Tokenizer(BPE(unk_token="")) + backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False) + backend.decoder = decoders.ByteLevel() + backend.train_from_iterator( + [text], + trainers.BpeTrainer( + vocab_size=300, + initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), + special_tokens=[""], + ), + ) + fast = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="") + fast.save_pretrained(tmp_path) + # 真实 DeepSeek checkpoint 声明 LlamaTokenizerFast，但实际编码和解码均为 ByteLevel。 + config_path = tmp_path / "tokenizer_config.json" + config = json.loads(config_path.read_text()) + config.update(tokenizer_class="LlamaTokenizerFast", legacy=True) + config_path.write_text(json.dumps(config)) + (tmp_path / "config.json").write_text( + json.dumps( + { + "model_type": model_type, + "architectures": ["DeepseekV3ForCausalLM"], + } + ) + ) + + tokenizer = get_tokenizer(str(tmp_path)) + expected_ids = backend.encode(text).ids + assert tokenizer.encode(text, add_special_tokens=False) == expected_ids + assert tokenizer.decode(expected_ids, skip_special_tokens=True) == text + + @pytest.mark.parametrize("model_type", ["llama", "deepseek_v32", "deepseek_v4"]) def test_tokenizer_imports_are_isolated(model_type): result = subprocess.run( @@ -48,13 +90,6 @@ def find_spec(self, fullname, path=None, target=None): from lightllm.server import tokenizer as module assert not any(name.startswith("lightllm.models.") for name in sys.modules) -if model_type == "deepseek_v32": - original = module.AutoTokenizer.from_pretrained - def load(*args, **kwargs): - assert AutoConfig.for_model("deepseek_v32").model_type == "deepseek_v32" - return original(*args, **kwargs) - module.AutoTokenizer.from_pretrained = load - with tempfile.TemporaryDirectory() as directory: fast = PreTrainedTokenizerFast(tokenizer_object=Tokenizer(WordLevel({"hello": 0, "[UNK]": 1}, unk_token="[UNK]"))) fast.save_pretrained(directory) @@ -64,6 +99,7 @@ def load(*args, **kwargs): tokenizer = module.get_tokenizer(directory) assert tokenizer.encode("hello") == [0] if model_type == "deepseek_v32": + assert AutoConfig.for_model("deepseek_v32").model_type == "deepseek_v32" assert tokenizer.__class__.__module__ == "lightllm.models.deepseek3_2.tokenizer" prompt = tokenizer.apply_chat_template(messages=[{"role": "user", "content": "hello"}]) assert "hello" in prompt diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py index 6b57270125..2cde8ec635 100644 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ b/unit_tests/test_static_benchmark_paged_kv.py @@ -20,7 +20,7 @@ ) -class FakeMemManager: +class FakeKVAllocator: def __init__(self, page_size: int): self.page_size = page_size self.next_index = 0 @@ -35,15 +35,14 @@ def alloc(self, size: int) -> torch.Tensor: def make_executor(page_size: int = 4): - mem_manager = FakeMemManager(page_size=page_size) + allocator = FakeKVAllocator(page_size=page_size) req_manager = SimpleNamespace( - mem_manager=mem_manager, req_to_token_indexs=torch.full((4, 32), -1, dtype=torch.int32), ) - model = SimpleNamespace(mem_manager=mem_manager, req_manager=req_manager) + model = SimpleNamespace(att_state_manager=SimpleNamespace(full_att_allocator=allocator), req_manager=req_manager) args = SimpleNamespace(mtp_mode=None, mtp_step=0) executor = StaticBenchmarkExecutor(args=args, model=model, draft_models=[], token_source=None) - return executor, mem_manager, req_manager + return executor, allocator, req_manager def assert_complete_physical_pages(indexes: torch.Tensor, page_size: int): @@ -53,7 +52,7 @@ def assert_complete_physical_pages(indexes: torch.Tensor, page_size: int): def test_prefill_allocates_and_reuses_complete_pages_per_request(): - executor, mem_manager, req_manager = make_executor(page_size=4) + executor, allocator, req_manager = make_executor(page_size=4) req_idx = torch.tensor([0, 1], dtype=torch.int32) model_input = executor._make_prefill_input( @@ -62,7 +61,7 @@ def test_prefill_allocates_and_reuses_complete_pages_per_request(): ready_cache_len=0, ) - assert mem_manager.alloc_sizes == [4, 4] + assert allocator.alloc_sizes == [4, 4] assert model_input.b_is_decode_req.tolist() == [False, False] assert executor._hold_kv_len_by_req == {0: 4, 1: 4} assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :4], page_size=4) @@ -77,7 +76,7 @@ def test_prefill_allocates_and_reuses_complete_pages_per_request(): input_ids=torch.tensor([1, 2], dtype=torch.int64), max_kv_seq_len=4, ) - assert mem_manager.alloc_sizes == [4, 4] + assert allocator.alloc_sizes == [4, 4] # 跨过页边界后，每个请求分别追加一整页。 executor._make_decode_input( @@ -88,14 +87,14 @@ def test_prefill_allocates_and_reuses_complete_pages_per_request(): input_ids=torch.tensor([3, 4], dtype=torch.int64), max_kv_seq_len=5, ) - assert mem_manager.alloc_sizes == [4, 4, 4, 4] + assert allocator.alloc_sizes == [4, 4, 4, 4] assert executor._hold_kv_len_by_req == {0: 8, 1: 8} assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :8], page_size=4) assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :8], page_size=4) def test_mtp_expanded_rows_allocate_once_per_request(): - executor, mem_manager, req_manager = make_executor(page_size=4) + executor, allocator, req_manager = make_executor(page_size=4) executor.args.mtp_mode = "eagle_with_att" executor.args.mtp_step = 3 @@ -110,7 +109,7 @@ def test_mtp_expanded_rows_allocate_once_per_request(): ) assert model_input.batch_size == 4 - assert mem_manager.alloc_sizes == [8, 8] + assert allocator.alloc_sizes == [8, 8] assert executor._hold_kv_len_by_req == {0: 8, 1: 8} assert_complete_physical_pages(req_manager.req_to_token_indexs[0, :8], page_size=4) assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :8], page_size=4) diff --git a/unit_tests/utils/test_config_utils_fixed_kv.py b/unit_tests/utils/test_config_utils_fixed_kv.py deleted file mode 100644 index 116e299e24..0000000000 --- a/unit_tests/utils/test_config_utils_fixed_kv.py +++ /dev/null @@ -1,33 +0,0 @@ -from types import SimpleNamespace - -import pytest - -from lightllm.utils import config_utils - - -@pytest.fixture(autouse=True) -def clear_fixed_kv_len_cache(): - config_utils.get_fixed_kv_len.cache_clear() - yield - config_utils.get_fixed_kv_len.cache_clear() - - -def test_fixed_kv_len_supports_page_size_one(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=1, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {"prompt_cache_token_ids": [1, 2, 3]}) - - assert config_utils.get_fixed_kv_len() == 3 - - -def test_fixed_kv_len_uses_only_complete_pages(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=2, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {"prompt_cache_token_ids": [1, 2, 3]}) - - assert config_utils.get_fixed_kv_len() == 2 - - -def test_no_fixed_kv_supports_paged_kv(monkeypatch): - monkeypatch.setattr(config_utils, "get_env_start_args", lambda: SimpleNamespace(page_size=4, model_dir="model")) - monkeypatch.setattr(config_utils, "get_config_json", lambda _: {}) - - assert config_utils.get_fixed_kv_len() == 0