diff --git a/.cursor/rules/named-args-and-types.mdc b/.cursor/rules/named-args-and-types.mdc new file mode 100644 index 0000000000..a463ccac7d --- /dev/null +++ b/.cursor/rules/named-args-and-types.mdc @@ -0,0 +1,12 @@ +--- +description: 新代码的函数调用使用参数名，参数带类型标注 +alwaysApply: true +--- + +# 新代码的调用和类型 + +新写的函数调用使用 `param_name=value`。不要用 `*args`、`*activations`、`**kwargs` 把一组参数展开传进去。 + +新写的函数参数带类型标注。 + +修改旧代码时，原有调用和原有签名保持不动，避免把风格调整混进 review。这次改动里新加的函数、新加的调用按上面两条写。 diff --git a/.cursor/rules/no-test-only-production-api.mdc b/.cursor/rules/no-test-only-production-api.mdc new file mode 100644 index 0000000000..a9e877828e --- /dev/null +++ b/.cursor/rules/no-test-only-production-api.mdc @@ -0,0 +1,10 @@ +--- +description: 测试专用的实现和接口只放在测试里，不留在正式代码中 +alwaysApply: true +--- + +# 测试代码不进正式实现 + +正式路径不保留只为测试、benchmark 或对照方便而存在的函数、参数和分支。 + +需要这类行为时，在测试或 benchmark 自己实现，并调用正式接口。正式接口的参数对应真实推理会传入的数据，不为测试场景增加可选参数。 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 index 3fe6f2317e..82a07fc486 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py @@ -13,8 +13,24 @@ def create_gdn_state( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> GDNState: - """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" + """创建 GDN 请求运行态。replay 模式使用独立的 ReplayGDNState。""" + if ssm_state_mode == "replay": + from .replay_gdn_state import ReplayGDNState + + return ReplayGDNState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, + ) + if ssm_state_mode != "native": + raise ValueError(f"unknown GDN state mode: {ssm_state_mode}") return GDNState( config=config, request_capacity=request_capacity, diff --git a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py index b8a641a73d..82eb4026c5 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -46,6 +46,7 @@ def __init__( self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step + self.mtp_enabled = mtp_step > 0 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) @@ -67,7 +68,7 @@ def __init__( 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.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) self.big_pages_pool = None self.small_pages_pool = None @@ -126,7 +127,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: 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: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) def save_bounded_att_runtime_to_pages( @@ -170,7 +171,7 @@ def restore_bounded_att_runtime_from_pages( 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: + if self.mtp_enabled: for req_idx in req_indices: self.mtp_state_index[req_idx].fill_(0) @@ -178,10 +179,11 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + req_verify_width: 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: + if self.mtp_enabled: self.mtp_state_index[req_indices] = accepted_steps - 1 # ======================================================================== @@ -320,7 +322,7 @@ def copy_runtime_state_from_dp( ) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for dst_req in dst_req_indices: self.mtp_state_index[dst_req].fill_(0) @@ -375,7 +377,7 @@ def read_from_pd_page( write=False, ): copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) # ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py new file mode 100644 index 0000000000..137ce8ad32 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -0,0 +1,563 @@ +"""GDN replay 的 bounded attention 状态。 + +Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 +一层 decode 重放在 replay_gdn_decode。 +""" + +from __future__ import annotations + +from functools import lru_cache +from math import prod + +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.utils.log_utils import init_logger + +from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion +from ...utils import deepcopy_result, dtype_name +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .gdn_pages_pool import GDNPagesPool + +logger = init_logger(__name__) + + +class ReplayGDNState(BoundedAttentionState): + """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ssm_state_mode: str = "replay", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ) -> None: + """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" + assert ssm_state_mode == "replay", f"ReplayGDNState only implements replay mode, got {ssm_state_mode}" + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self.mtp_enabled = mtp_step > 0 + self.gate_lower_bound = gate_lower_bound + self._algorithm = "gdn" + self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} + layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + ( + layer_count, + request_capacity, + config.conv_channels, + config.conv_width + mtp_step, + ), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + layer_count, + request_capacity, + 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) + self.big_pages_pool = None + self.small_pages_pool = None + replay_config = dict(ssm_replay_config or {}) + capacity = int(replay_config.get("cache_len", 16)) + # 宽度为 1 逐 token 重放。开了 MTP 时默认先投影整个验证窗口。 + compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + verify_width = mtp_step + 1 + # 连续前缀挂在上面的 canonical SSM 上，不再另分配状态。 + # capacity 是 2 的幂，并且不小于一次验证窗口。 + assert self.ssm_state.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + self.kda = False + self.capacity = capacity + self.verify_width = verify_width + self.compute_mode = compute_mode + layers, slots, hv, k, v = self.ssm_state.shape + self.num_key_heads = config.num_k_heads + assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 + # 最后一个 slot 是 HOLD，replay 写零并且不记 history。 + self.hold_req_id = slots - 1 + # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 + self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) + # 按 launch 行和 key head 同步。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 + self.tile_sync_mark = torch.zeros( + (slots * (mtp_step + 6), self.num_key_heads), dtype=torch.int32, device=device + ) + # 归一化 k 有范围保证，bf16 激活用 fp16 保留更多尾数；d 无此保证，使用 fp32。 + # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 + # GDN 的 g 是每个头一个标量。 + reconstruction_dtype = torch.float16 if config.conv_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, self.num_key_heads, capacity, k) + self.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=device) + self.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=device) + self.gates = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device) + self.raw_keys = torch.empty(key_shape, dtype=config.conv_dtype, device=device) + self.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=config.conv_dtype, device=device) + self.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device) + + def get_state_type(self) -> str: + """返回算法名和 Conv/SSM 存储精度。""" + return ( + f"{self._algorithm}_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 + ) + + # ======================================================================== + # 二、计算视图、运行态重置、已接受快照与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射。""" + return logical_layer_id in self._logical_to_storage_layer + + def storage_layer(self, logical_layer_id: int) -> int: + """返回逻辑层在 history 和 canonical state 里的存储下标。""" + + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int): + index = self.storage_layer(logical_layer_id) + return self.conv_state[index], self.ssm_state[index] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + self.ssm_state[:, req_idx].zero_() + self.cursors[req_idx].zero_() + if self.mtp_enabled: + 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, + page_slot_ids: list[int], + ) -> None: + """把 prefill 基准 Conv 窗口和 canonical SSM 写入页。history 与 MTP 候选区不参与。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "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) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + conv_copies.append((self.conv_state[:, req_idx, :, :width], pool.conv_state[slot_id])) + ssm_copies.append((self.ssm_state[:, req_idx], 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, + page_slot_ids: list[int], + ) -> None: + """从页恢复基准 Conv/SSM。使目标槽的旧 history 失效。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "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) + width = self.config.conv_width + 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, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + for req_idx in req_indices: + self.cursors[req_idx].zero_() + if self.mtp_enabled: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, + ) -> None: + """MTP 验证后写入索引，并从整窗全接受后的 cursor 上剪掉未接受尾部。""" + if not self.mtp_enabled: + return + commit_ssm_prefix(self.mtp_state_index, req_indices, accepted_steps, req_verify_width, self.cursors) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间。不含 history 或 MTP 候选区。""" + return self._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._tensor_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" + # 本地目标保留本 TP rank 的分片形状。layout.regions 描述 CPU 全局字段。 + pool_page_fields = [ + self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids + ] + 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) -> 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" + # 大页是基准快照，与请求槽位无关。页池须在本进程注册，当前流访问完才能归还页面。 + 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: "ReplayGDNState", + ) -> None: + """复制 prefill 基准 Conv/SSM 到目标单槽。使目标槽的旧 history 失效。""" + 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" + width = self.config.conv_width + 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, :, :width], + self.conv_state[:, dst_req, :, :width], + ) + ) + ssm_copies.append((remote_dp_bounded_state.ssm_state[:, src_req], self.ssm_state[:, dst_req])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + for dst_req in dst_req_indices: + self.cursors[dst_req].zero_() + if self.mtp_enabled: + 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._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: + """写出 prefill 基准 Conv 窗口和 canonical SSM。源 history 不参与。""" + 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._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, 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: + """读入基准 Conv/SSM。使目标槽的旧 history 失效。""" + 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._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, 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) + self.cursors[req_idx].zero_() + if self.mtp_enabled: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、Replay history + # + # cursor 是已接受、尚未折回的长度，各层只读。 + # 最后一层返回后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 + # MTP 验证后再减去没接受的尾部。宽度为 1 时窗口就是这一个 token，没有尾部可剪。 + # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 把已接受前缀写回 canonical SSM。 + # ======================================================================== + + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: + """最后一层按整窗全接受推进 cursor。非最后一层直接返回。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: + return + advance_cursor(self, reqs, cu_seqlens) + + def fold_accepted_prefix(self, reqs: torch.Tensor) -> None: + """用原始 k/v 把 cursor 覆盖的已接受前缀折进 canonical SSM，并清掉这些请求的 cursor。""" + + if reqs.numel() == 0: + return + fold_prefix(self, reqs) + self.cursors.index_fill_(0, reqs.long(), 0) + + # ======================================================================== + # 七、私有辅助实现 + # ======================================================================== + + def _transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 只传有效历史窗口，不包含 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 + regions.append( + 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, + ) + ) + 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 + + @staticmethod + def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): + """描述一页里的 Conv 和 SSM 字段。""" + 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 index fe1fceb3bd..11d7425cd9 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -102,8 +102,12 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, ) -> None: - """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" + """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。 + + req_verify_width 与请求逐行对齐，是每个请求这次送去验证的 token 数。Replay 用它从全接受后的 cursor 上剪掉未接受尾部。 + """ ... 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 index 49c9ad0ce5..1742441389 100644 --- a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py @@ -13,8 +13,24 @@ def create_kda_state( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> KDAState: - """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" + """创建 KDA 请求运行态。replay 模式使用独立的 ReplayKDAState。""" + if ssm_state_mode == "replay": + from .replay_kda_state import ReplayKDAState + + return ReplayKDAState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, + ) + if ssm_state_mode != "native": + raise ValueError(f"unknown KDA state mode: {ssm_state_mode}") return KDAState( config=config, request_capacity=request_capacity, diff --git a/lightllm/common/attention_state/bounded_att/kda/kda_state.py b/lightllm/common/attention_state/bounded_att/kda/kda_state.py index e5480a3ab0..68097905d5 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -46,6 +46,7 @@ def __init__( self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 self.mtp_step = mtp_step + self.mtp_enabled = mtp_step > 0 self._logical_to_storage_layer = { logical_layer_id: storage_layer_index for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) @@ -67,7 +68,7 @@ def __init__( 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.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) self.big_pages_pool = None self.small_pages_pool = None @@ -126,7 +127,7 @@ def reset_bounded_att_state(self, req_indices: List[int]) -> None: 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: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) def save_bounded_att_runtime_to_pages( @@ -170,7 +171,7 @@ def restore_bounded_att_runtime_from_pages( 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: + if self.mtp_enabled: for req_idx in req_indices: self.mtp_state_index[req_idx].fill_(0) @@ -178,10 +179,11 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + req_verify_width: 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: + if self.mtp_enabled: self.mtp_state_index[req_indices] = accepted_steps - 1 # ======================================================================== @@ -320,7 +322,7 @@ def copy_runtime_state_from_dp( ) copy_bounded_att_tensors(copies=conv_copies) copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: for dst_req in dst_req_indices: self.mtp_state_index[dst_req].fill_(0) @@ -375,7 +377,7 @@ def read_from_pd_page( write=False, ): copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: + if self.mtp_enabled: self.mtp_state_index[req_idx].fill_(0) # ======================================================================== diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py new file mode 100644 index 0000000000..2c73e97d6f --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -0,0 +1,563 @@ +"""KDA replay 的 bounded attention 状态。 + +Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 +一层 decode 重放在 replay_kda_decode。 +""" + +from __future__ import annotations + +from functools import lru_cache +from math import prod + +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.utils.log_utils import init_logger + +from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion +from ...utils import deepcopy_result, dtype_name +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .kda_pages_pool import KDAPagesPool + +logger = init_logger(__name__) + + +class ReplayKDAState(BoundedAttentionState): + """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ssm_state_mode: str = "replay", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ) -> None: + """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" + assert ssm_state_mode == "replay", f"ReplayKDAState only implements replay mode, got {ssm_state_mode}" + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self.mtp_enabled = mtp_step > 0 + self.gate_lower_bound = gate_lower_bound + self._algorithm = "kda" + self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} + layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + ( + layer_count, + request_capacity, + config.conv_channels, + config.conv_width + mtp_step, + ), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + layer_count, + request_capacity, + 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) + self.big_pages_pool = None + self.small_pages_pool = None + replay_config = dict(ssm_replay_config or {}) + capacity = int(replay_config.get("cache_len", 16)) + # 宽度为 1 逐 token 重放。开了 MTP 时默认先投影整个验证窗口。 + compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + verify_width = mtp_step + 1 + # 连续前缀挂在上面的 canonical SSM 上，不再另分配状态。 + # capacity 是 2 的幂，并且不小于一次验证窗口。 + assert self.ssm_state.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + self.kda = True + self.capacity = capacity + self.verify_width = verify_width + self.compute_mode = compute_mode + layers, slots, hv, k, v = self.ssm_state.shape + self.num_key_heads = config.num_k_heads + assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 + # 最后一个 slot 是 HOLD，replay 写零并且不记 history。 + self.hold_req_id = slots - 1 + # cursor 是已接受、尚未折回的长度。各层只读，最后一层返回后才按整窗推进。 + self.cursors = torch.zeros(slots, dtype=torch.int32, device=device) + # 按 launch 行和 key head 同步。每个请求槽留 mtp_step + 6 行，盖住验证窗口和重复的 HOLD。 + self.tile_sync_mark = torch.zeros( + (slots * (mtp_step + 6), self.num_key_heads), dtype=torch.int32, device=device + ) + # 归一化 k 有范围保证，bf16 激活用 fp16 保留更多尾数；d 无此保证，使用 fp32。 + # 折回不用这份 d，而用下面的原始 k/v，避免把重建误差写进 SSM。g、beta 保持 fp32。 + # KDA 的 g 按 key 通道存。 + reconstruction_dtype = torch.float16 if config.conv_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, self.num_key_heads, capacity, k) + self.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=device) + self.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=device) + self.gates = torch.empty((layers, slots, hv, capacity, k), dtype=torch.float32, device=device) + self.raw_keys = torch.empty(key_shape, dtype=config.conv_dtype, device=device) + self.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=config.conv_dtype, device=device) + self.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=device) + + def get_state_type(self) -> str: + """返回算法名和 Conv/SSM 存储精度。""" + return ( + f"{self._algorithm}_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 + ) + + # ======================================================================== + # 二、计算视图、运行态重置、已接受快照与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射。""" + return logical_layer_id in self._logical_to_storage_layer + + def storage_layer(self, logical_layer_id: int) -> int: + """返回逻辑层在 history 和 canonical state 里的存储下标。""" + + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int): + index = self.storage_layer(logical_layer_id) + return self.conv_state[index], self.ssm_state[index] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + self.ssm_state[:, req_idx].zero_() + self.cursors[req_idx].zero_() + if self.mtp_enabled: + 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, + page_slot_ids: list[int], + ) -> None: + """把 prefill 基准 Conv 窗口和 canonical SSM 写入页。history 与 MTP 候选区不参与。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "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) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + conv_copies.append((self.conv_state[:, req_idx, :, :width], pool.conv_state[slot_id])) + ssm_copies.append((self.ssm_state[:, req_idx], 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, + page_slot_ids: list[int], + ) -> None: + """从页恢复基准 Conv/SSM。使目标槽的旧 history 失效。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "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) + width = self.config.conv_width + 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, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + for req_idx in req_indices: + self.cursors[req_idx].zero_() + if self.mtp_enabled: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, + ) -> None: + """MTP 验证后写入索引，并从整窗全接受后的 cursor 上剪掉未接受尾部。""" + if not self.mtp_enabled: + return + commit_ssm_prefix(self.mtp_state_index, req_indices, accepted_steps, req_verify_width, self.cursors) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间。不含 history 或 MTP 候选区。""" + return self._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._tensor_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" + # 本地目标保留本 TP rank 的分片形状。layout.regions 描述 CPU 全局字段。 + pool_page_fields = [ + self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids + ] + 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) -> 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" + # 大页是基准快照，与请求槽位无关。页池须在本进程注册，当前流访问完才能归还页面。 + 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: "ReplayKDAState", + ) -> None: + """复制 prefill 基准 Conv/SSM 到目标单槽。使目标槽的旧 history 失效。""" + 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" + width = self.config.conv_width + 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, :, :width], + self.conv_state[:, dst_req, :, :width], + ) + ) + ssm_copies.append((remote_dp_bounded_state.ssm_state[:, src_req], self.ssm_state[:, dst_req])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + for dst_req in dst_req_indices: + self.cursors[dst_req].zero_() + if self.mtp_enabled: + 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._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: + """写出 prefill 基准 Conv 窗口和 canonical SSM。源 history 不参与。""" + 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._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, 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: + """读入基准 Conv/SSM。使目标槽的旧 history 失效。""" + 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._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, 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) + self.cursors[req_idx].zero_() + if self.mtp_enabled: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、Replay history + # + # cursor 是已接受、尚未折回的长度，各层只读。 + # 最后一层返回后按整窗全接受推进：放不下时 cursor 变成窗口长度，否则加上窗口长度。 + # MTP 验证后再减去没接受的尾部。宽度为 1 时窗口就是这一个 token，没有尾部可剪。 + # 前缀放不下时，折回发生在 replay kernel 内。fold_accepted_prefix 把已接受前缀写回 canonical SSM。 + # ======================================================================== + + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor, cu_seqlens: torch.Tensor) -> None: + """最后一层按整窗全接受推进 cursor。非最后一层直接返回。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if not layers or logical_layer_id != layers[-1] or reqs.numel() == 0: + return + advance_cursor(self, reqs, cu_seqlens) + + def fold_accepted_prefix(self, reqs: torch.Tensor) -> None: + """用原始 k/v 把 cursor 覆盖的已接受前缀折进 canonical SSM，并清掉这些请求的 cursor。""" + + if reqs.numel() == 0: + return + fold_prefix(self, reqs) + self.cursors.index_fill_(0, reqs.long(), 0) + + # ======================================================================== + # 七、私有辅助实现 + # ======================================================================== + + def _transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 只传有效历史窗口，不包含 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 + regions.append( + 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, + ) + ) + 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 + + @staticmethod + def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): + """描述一页里的 Conv 和 SSM 字段。""" + 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/fp8_e8m0_swa_state.py b/lightllm/common/attention_state/bounded_att/sliding_window/fp8_e8m0_swa_state.py index a1d72102ac..5afdfab160 100644 --- 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 @@ -179,7 +179,12 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, + ) -> None: """ring 按绝对 token 位置寻址，无需提交额外游标。""" pass 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 index cea0a49af2..5d9a0ffe12 100644 --- 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 @@ -210,7 +210,12 @@ def restore_bounded_att_runtime_from_pages( write=False, ) - def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + def commit_bounded_att_state( + self, + req_indices: torch.Tensor, + accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, + ) -> None: """按绝对位置寻址；接受后继续覆盖 ring，不维护额外 MTP 游标。""" pass diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 6fb18c7ea1..c0a85c367e 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -396,13 +396,19 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + req_verify_width: torch.Tensor, role: str, ) -> None: - """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" - assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" + """传入 CUDA 请求索引、int32 接受步数和每个请求的验证宽度；无参与 State 时为空操作。 + + req_verify_width 与请求逐行对齐，是每个请求这次送去验证的 token 数。Replay 用它剪掉未接受尾部。 + """ + assert ( + req_indices.is_cuda and accepted_steps.is_cuda and req_verify_width.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) + state.commit_bounded_att_state(req_indices, accepted_steps, req_verify_width) # ======================================================================== # 四、CPU/DP/PD 操作均由对应 Operator 提供。 diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index 6a40a05183..08528f98f6 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -1,56 +1,77 @@ -from .import_utils import load_backend_class +from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend +from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend +from lightllm.common.basemodel.attention.linear.replay_gdn import ( + ReplayGDNFlashQlaAttBackend, + ReplayGDNTritonAttBackend, +) +from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend from lightllm.utils.backend_validator import validate from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) -linear_prefill_att_backend_classes = { - "flashqla": "linear.flashqla:FlashQlaLinearAttBackend", - "triton": "linear.triton:TritonLinearAttBackend", +gdn_prefill_backends = { + "flashqla": FlashQlaLinearAttBackend, + "triton": TritonLinearAttBackend, } -linear_decode_att_backend_classes = { - "triton": "linear.triton:TritonLinearAttBackend", +gdn_replay_prefill_backends = { + "flashqla": ReplayGDNFlashQlaAttBackend, + "triton": ReplayGDNTritonAttBackend, } +gdn_decode_backends = { + "triton": TritonLinearAttBackend, +} + +gdn_replay_decode_backends = { + "triton": ReplayGDNTritonAttBackend, +} + +kda_backends = { + "triton": KDALinearAttBackend, +} + +kda_replay_backends = { + "triton": ReplayKDALinearAttBackend, +} -def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): + +def get_gdn_linear_prefill_att_backend_class(index=1, priority_list=("flashqla", "triton")): args = get_env_start_args() - backend_str = _get_backend_str(args.llm_prefill_att_backend, index) - if backend_str != "auto": - return load_backend_class(linear_prefill_att_backend_classes[backend_str]) + classes = gdn_replay_prefill_backends if args.ssm_state_mode == "replay" else gdn_prefill_backends + return _select_backend(classes, args.llm_prefill_att_backend, index, priority_list, "prefill") - for backend_name in priority_list: - backend_class = linear_prefill_att_backend_classes[backend_name] - if backend_name == "triton": - logger.info("Linear prefill attention backend: triton.") - return load_backend_class(backend_class) - if validate(backend_name): - logger.info(f"Linear prefill attention backend: {backend_name} (validated).") - return load_backend_class(backend_class) - logger.warning("No linear prefill attention backend validation succeeded, falling back to Triton.") - return load_backend_class("linear.triton:TritonLinearAttBackend") +def get_gdn_linear_decode_att_backend_class(index=1, priority_list=("triton",)): + args = get_env_start_args() + classes = gdn_replay_decode_backends if args.ssm_state_mode == "replay" else gdn_decode_backends + return _select_backend(classes, args.llm_decode_att_backend, index, priority_list, "decode") -def get_qwen35_linear_decode_att_backend_class(index=1, priority_list=("triton",)): +def get_kda_linear_att_backend_class(): args = get_env_start_args() - backend_str = _get_backend_str(args.llm_decode_att_backend, index) + classes = kda_replay_backends if args.ssm_state_mode == "replay" else kda_backends + return classes["triton"] + + +def _select_backend(classes, backend_args, index, priority_list, stage): + backend_str = _get_backend_str(backend_args, index) if backend_str != "auto": - return load_backend_class(linear_decode_att_backend_classes[backend_str]) + return classes[backend_str] for backend_name in priority_list: - backend_class = linear_decode_att_backend_classes[backend_name] - if backend_name == "triton": - logger.info("Linear decode attention backend: triton.") - return load_backend_class(backend_class) - if validate(backend_name): - logger.info(f"Linear decode attention backend: {backend_name} (validated).") - return load_backend_class(backend_class) - - logger.warning("No linear decode attention backend validation succeeded, falling back to Triton.") - return load_backend_class("linear.triton:TritonLinearAttBackend") + if backend_name == "triton" or validate(backend_name): + if backend_name == "triton": + logger.info(f"Linear {stage} attention backend: triton.") + else: + logger.info(f"Linear {stage} attention backend: {backend_name} (validated).") + return classes[backend_name] + + logger.warning(f"No linear {stage} attention backend validation succeeded, falling back to Triton.") + return classes["triton"] def _get_backend_str(backend_args, index): diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py new file mode 100644 index 0000000000..8e4a081f73 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -0,0 +1,356 @@ +"""GDN replay 的推理入口。 + +Triton 和 FlashQla 各有一个 backend。它们只差 prefill chunk kernel， +decode 都走 ReplayGDNState。两个类都直接实现 backend 接口，不继承 native GDN。 +""" + +from __future__ import annotations + +import dataclasses +import weakref +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ( + ReplayGDNState, +) +from lightllm.common.basemodel.attention.base_att import ( + AttControl, + BaseAttBackend, + BaseDecodeAttState, + BasePrefillAttState, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import ( + fused_gdn_gating, +) +from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import ( + conv_pack_gdn_decode_inputs, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +from lightllm.utils.envs_utils import get_env_start_args + +from .replay_gdn_decode import replay + +if TYPE_CHECKING: + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class ReplayGDNTritonAttBackend(BaseAttBackend): + """GDN replay 的 Triton 入口。Prefill 使用仓内 chunk kernel。""" + + def __init__(self, model): + super().__init__(model=model) + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule + + _configure_gdn(self, model, chunk_gated_delta_rule) + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +class ReplayGDNFlashQlaAttBackend(BaseAttBackend): + """GDN replay 的 FlashQla 入口。Prefill 使用 flash_qla chunk kernel。""" + + def __init__(self, model): + super().__init__(model=model) + from flash_qla import chunk_gated_delta_rule + + _configure_gdn(self, model, chunk_gated_delta_rule) + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +@dataclasses.dataclass +class ReplayGDNPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + reqs = self.infer_state.b_req_idx + self.b_conv_buffer_idx = reqs + self.b_ssm_buffer_idx = reqs + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill and att_control.linear_att_prefill_dict is not None + params = att_control.linear_att_prefill_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend = self.backend + state = _replay_state(self, layer_num) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) + core_attn_out = self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight) + return core_attn_out, z + + def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): + backend = self.backend + g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer_weight.linear_conv1d.mm_param.weight, + bias=layer_weight.linear_conv1d.bias, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation=backend.activation, + ).transpose(0, 1) + query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv) + core_attn_out, last_state = backend.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=ssm_states[self.b_ssm_buffer_idx], + output_final_state=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[self.b_ssm_buffer_idx] = last_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + +@dataclasses.dataclass +class ReplayGDNDecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_cu_seqlens: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + role = self.infer_state.att_state_role + matches = [ + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayGDNState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayGDNState for {role}") + draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) + self._init_indices(draft_step, matches[0].mtp_state_index) + self.b_ssm_buffer_idx = self.b_conv_buffer_idx + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode and att_control.linear_att_decode_dict is not None + params = att_control.linear_att_decode_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend = self.backend + mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) + state = _replay_state(self, layer_num) + conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) + draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) + if draft_step > 0: + core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) + else: + core, z = self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) + return core, z + + def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: + infer = self.infer_state + if draft_step == 0: + self.b_conv_buffer_idx = infer.b_req_idx + self.b1_mtp_cu_q_seq_len = None + self.b_num_accepted_tokens = None + # 每行一个 token。replay 和 cursor 共用这份前缀和。 + self.b_cu_seqlens = torch.arange( + infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device + ) + return + if self.backend.uses_dynamic_spec_verify_layout(): + cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=mtp_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + self.b1_mtp_cu_q_seq_len = cu + self.b_cu_seqlens = cu + self.b_conv_buffer_idx = conv_idx + self.b_num_accepted_tokens = accepted + return + width = draft_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "GDN replay decode requires batch_size divisible by draft_step + 1, " + f"got batch_size={batch_size}, draft_step={draft_step}." + ) + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + self.b_cu_seqlens = self.b1_mtp_cu_q_seq_len + self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 + + def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state): + backend = self.backend + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer_weight.linear_conv1d.mm_param.weight, + layer_weight.linear_conv1d.bias, + self.b_conv_buffer_idx, + backend.activation, + backend.conv_kernel_dim, + backend.tp_num_k_heads, + backend.head_k_dim, + backend.tp_num_v_heads, + backend.head_v_dim, + ) + query = query.view(query.shape[0], query.shape[-2], query.shape[-1]) + key = key.view(key.shape[0], key.shape[-2], key.shape[-1]) + value = value.view(value.shape[0], value.shape[-2], value.shape[-1]) + output = replay( + state, + state.storage_layer(layer_num), + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b_cu_seqlens, + ) + return output, z + + def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state): + backend = self.backend + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer_weight.linear_conv1d.mm_param.weight, + mtp_step=draft_step, + bias=layer_weight.linear_conv1d.bias, + activation=backend.activation, + conv_state_indices=self.b_conv_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + query, key, value = torch.split( + mixed_qkv, + [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], + dim=-1, + ) + query = query.view(-1, backend.tp_num_k_heads, backend.head_k_dim) + key = key.view(-1, backend.tp_num_k_heads, backend.head_k_dim) + value = value.view(-1, backend.tp_num_v_heads, backend.head_v_dim) + output = replay( + state, + state.storage_layer(layer_num), + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b_cu_seqlens, + ) + return output + + +def _replay_state(att_state, layer_num: int) -> ReplayGDNState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=ReplayGDNState, + ) + + +def _configure_gdn(backend, model, prefill_kernel) -> None: + """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" + + config = model.config + tp_world_size = model.tp_world_size_ + backend.mtp_step = get_env_start_args().mtp_step + backend.num_v_heads = config["linear_num_value_heads"] + backend.num_k_heads = config["linear_num_key_heads"] + backend.head_k_dim = config["linear_key_head_dim"] + backend.head_v_dim = config["linear_value_head_dim"] + backend.key_dim = backend.head_k_dim * backend.num_k_heads + backend.value_dim = backend.head_v_dim * backend.num_v_heads + backend.conv_kernel_dim = config["linear_conv_kernel_dim"] + backend.activation = config["hidden_act"] + backend.tp_qkvz_dim = (backend.key_dim * 2 + backend.value_dim * 2) // tp_world_size + backend.tp_ba_dim = (backend.num_v_heads * 2) // tp_world_size + backend.tp_num_k_heads = backend.num_k_heads // tp_world_size + backend.tp_num_v_heads = backend.num_v_heads // tp_world_size + backend.tp_key_dim = backend.key_dim // tp_world_size + backend.tp_value_dim = backend.value_dim // tp_world_size + assert backend.num_v_heads % backend.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" + backend.num_v_heads_per_k_head = backend.num_v_heads // backend.num_k_heads + ssm_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} + backend.ssm_state_dtype = ssm_dtypes.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) + backend.prefill_kernel = prefill_kernel + + +def _split_qkvzba(backend, mixed_qkvzba): + qkv_dim = backend.tp_key_dim * 2 + backend.tp_value_dim + z_end = qkv_dim + backend.tp_value_dim + b_end = z_end + backend.tp_num_v_heads + mixed_qkv = mixed_qkvzba[:, :qkv_dim] + z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, backend.tp_num_v_heads, backend.head_v_dim) + b = mixed_qkvzba[:, z_end:b_end] + a = mixed_qkvzba[:, b_end:] + return mixed_qkv, z, b, a + + +def _rearrange_mixed_qkv(backend, mixed_qkv, decode=False): + query, key, value = torch.split( + mixed_qkv, + [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], + dim=-1, + ) + if decode: + batch_size = mixed_qkv.shape[0] + leading = (batch_size, 1) + else: + leading = (1, query.shape[0]) + query = query.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) + key = key.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) + value = value.view(*leading, backend.tp_num_v_heads, backend.head_v_dim) + return query, key, value diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py new file mode 100644 index 0000000000..ac7af9dbdc --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -0,0 +1,200 @@ +"""GDN replay decode。 + +History 张量和 cursor 在 ReplayGDNState。这里调优并执行一层重放。 +""" + +from __future__ import annotations + +import ast + +import torch +import triton + +from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( + _decode_replay, + check_replay_sync_mark, +) +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune + + +def get_configs() -> list[dict[str, int]]: + """返回 GDN replay 的候选 Triton 配置。""" + return [ + {"BV": bv, "num_warps": warps, "num_stages": stages} + for bv, warps in [(8, 1), (16, 1), (16, 2), (32, 1), (32, 2), (32, 4), (64, 2), (64, 4), (128, 4)] + for stages in (1, 2, 3) + ] + + +def get_static_key(state, q) -> dict[str, object]: + """区分头维度、状态精度、history 容量和投影模式。""" + _, _, hv, kd, vd = state.ssm_state.shape + return { + "v": 1, + "shape": str((q.shape[-2], hv, state.num_key_heads, kd, vd, state.verify_width)), + "dtype": str((q.dtype, state.ssm_state.dtype)), + "capacity": state.capacity, + "recon_dtype": str(state.keys.dtype), + "projection": "gdn_mma" if state.compute_mode == "precompute" else "inline", + } + + +def get_run_key(q, cu_seqlens) -> tuple[int, int]: + """调优键是 token 数和请求数。q 是 [tokens, heads, dim]。""" + return q.shape[0], cu_seqlens.numel() - 1 + + +def _check_replay_inputs(q, k, v, reqs, cu_seqlens) -> None: + assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" + assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" + assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" + + +def rebuild_inputs(state, linear_layer_idx, q, k, v, a, b, a_log, bias, reqs, cu_seqlens, run_config=None): + """为一层 scratch 和同形状激活分配调优输入。不复制服务状态池。""" + del linear_layer_idx, reqs, run_config + tokens, run_count = get_run_key(q, cu_seqlens) + # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 + batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state.verify_width)) + _, _, hv, kd, vd = state.ssm_state.shape + scratch = type(state)( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=state.num_key_heads, + global_v_heads=hv, + key_dim=kd, + value_dim=vd, + conv_width=1, + conv_dtype=q.dtype, + ssm_dtype=state.ssm_state.dtype, + ), + batch + 1, + mtp_step=state.verify_width - 1, + device=q.device, + ssm_replay_config={"cache_len": state.capacity, "compute_mode": state.compute_mode}, + gate_lower_bound=state.gate_lower_bound, + ) + scratch.ssm_state.normal_(0, 0.02) + copied = [ + torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) + ] + bench_reqs = torch.cat( + ( + torch.arange(batch, dtype=torch.int32, device=q.device), + torch.full((run_count - batch,), batch, dtype=torch.int32, device=q.device), + ) + ) + # 计时用均匀窗口。宽度为 1 时 verify_width 就是 1。 + cu = torch.arange(run_count + 1, dtype=torch.int32, device=q.device) + cu = (cu * state.verify_width).clamp_max(tokens) + return (), dict( + state=scratch, + linear_layer_idx=0, + q=copied[0], + k=copied[1], + v=copied[2], + a=copied[3], + b=copied[4], + a_log=copied[5], + bias=copied[6], + reqs=bench_reqs, + cu_seqlens=cu, + ) + + +def _run_key_distance(left: str, right: str) -> int: + return sum(abs(x - y) for x, y in zip(ast.literal_eval(left), ast.literal_eval(right))) + + +@autotune( + kernel_name="replay_gdn_decode:v1", + configs_gen_func=get_configs, + static_key_func=get_static_key, + run_key_func=get_run_key, + run_key_distance_func=_run_key_distance, + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + rebuild_input_func=rebuild_inputs, + warmup_all_exist_config=False, +) +def replay( + state, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor, + run_config: dict[str, int] | None = None, +) -> torch.Tensor: + """执行一层 GDN replay decode。Kernel 接收这一层的 SSM 和 history；state 提供共用的 cursor 与配置。 + + q、k、v 是 [tokens, heads, dim]。cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。宽度为 1 时它是 0, 1, 2, ...。 + """ + + assert not state.kda, "GDN replay step does not use the KDA gate" + _check_replay_inputs(q, k, v, reqs, cu_seqlens) + if run_config is None: + run_config = {"BV": 32, "num_warps": 1, "num_stages": 3} + hv, kd, vd = state.ssm_state.shape[-3:] + assert state.num_key_heads in (q.shape[-2], hv) + out = torch.empty_like(v) + bv = run_config["BV"] + # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 + precompute_projections = state.compute_mode == "precompute" and bv >= 16 + padded_verify_width = triton.next_power_of_2(state.verify_width) + check_replay_sync_mark( + sync_mark=state.tile_sync_mark, + reqs=reqs, + value_heads=hv, + key_heads=state.num_key_heads, + ) + _decode_replay[(triton.cdiv(vd, bv), hv, reqs.numel())]( + Q=q, + StepK=k, + StepV=v, + A=a, + B=b, + Alog=a_log, + Bias=bias, + State=state.ssm_state[linear_layer_idx], + Keys=state.keys[linear_layer_idx], + Deltas=state.deltas[linear_layer_idx], + Gates=state.gates[linear_layer_idx], + RawKeys=state.raw_keys[linear_layer_idx], + RawValues=state.raw_values[linear_layer_idx], + Betas=state.betas[linear_layer_idx], + Reqs=reqs, + Cursors=state.cursors, + SyncMark=state.tile_sync_mark, + Cu=cu_seqlens, + Out=out, + SQ=q.stride(0), + SK=k.stride(0), + SV=v.stride(0), + SA=a.stride(0), + SB=b.stride(0), + H=q.shape[-2], + HV=hv, + KH=state.num_key_heads, + K=kd, + V=vd, + L=state.capacity, + HOLD=state.hold_req_id, + BK=triton.next_power_of_2(kd), + BV=bv, + PRECOMPUTE_PROJECTIONS=precompute_projections, + BW=padded_verify_width, + KDA=False, + LOWER=state.gate_lower_bound, + BQK=max(16, 2 * padded_verify_width), + BR=triton.next_power_of_2(max(16, state.capacity + padded_verify_width)), + NV=triton.cdiv(vd, bv), + num_warps=run_config["num_warps"], + num_stages=run_config["num_stages"], + ) + return out diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py new file mode 100644 index 0000000000..af931aa477 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -0,0 +1,257 @@ +"""KDA replay 的推理入口。 + +这个 backend 只服务 ReplayKDAState。维度在这里自己读取，不继承 native KDA backend。 +""" + +from __future__ import annotations + +import dataclasses +import weakref +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( + ReplayKDAState, +) +from lightllm.common.basemodel.attention.base_att import ( + AttControl, + BaseAttBackend, + BaseDecodeAttState, + BasePrefillAttState, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( + chunk_kda_with_fused_gate, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + +from .replay_kda_decode import replay + +if TYPE_CHECKING: + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class ReplayKDALinearAttBackend(BaseAttBackend): + """KDA replay 推理 backend。GLM-5-Next 在 replay 模式下只构造这一个类。""" + + def __init__(self, model): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = model.args.mtp_step + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayKDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayKDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +@dataclasses.dataclass +class ReplayKDAPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + reqs = self.infer_state.b_req_idx + self.b_conv_buffer_idx = reqs + self.b_ssm_buffer_idx = reqs + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill + params = att_control.linear_att_prefill_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend: ReplayKDALinearAttBackend = self.backend + state = _replay_state(self, layer_num) + 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( + params["mixed_qkv"].transpose(0, 1), + layer_weight.get_merged_kda_conv_weight(), + bias=None, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + ).transpose(0, 1) + q, k, v = backend.split_qkv(mixed_qkv) + q = q.view(1, -1, backend.tp_num_heads, backend.head_dim) + k = k.view(1, -1, backend.tp_num_heads, backend.head_dim) + v = v.view(1, -1, backend.tp_num_heads, backend.head_dim) + raw_gate = params["raw_gate"].view(1, -1, backend.tp_hidden_size) + raw_beta = params["raw_beta"].view(1, -1, backend.tp_num_heads) + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states[self.b_ssm_buffer_idx].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + return output + + +@dataclasses.dataclass +class ReplayKDADecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_cu_seqlens: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + role = self.infer_state.att_state_role + matches = [ + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayKDAState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayKDAState for {role}") + self._init_indices(self.backend.mtp_step, matches[0].mtp_state_index) + self.b_ssm_buffer_idx = self.b_conv_buffer_idx + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode + params = att_control.linear_att_decode_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend: ReplayKDALinearAttBackend = self.backend + state = _replay_state(self, layer_num) + conv_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: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=backend.mtp_step, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update( + params["mixed_qkv"], + conv_states, + layer_weight.get_merged_kda_conv_weight(), + **conv_kwargs, + ) + q, k, v = backend.split_qkv(mixed_qkv) + q = q.view(-1, backend.tp_num_heads, backend.head_dim) + k = k.view(-1, backend.tp_num_heads, backend.head_dim) + v = v.view(-1, backend.tp_num_heads, backend.head_dim) + raw_gate = params["raw_gate"].reshape(-1, backend.tp_hidden_size) + raw_beta = params["raw_beta"].reshape(-1, backend.tp_num_heads) + layer = state.storage_layer(layer_num) + output = replay( + state_ssm=state.ssm_state[layer], + state_keys=state.keys[layer], + state_deltas=state.deltas[layer], + state_gates=state.gates[layer], + state_raw_keys=state.raw_keys[layer], + state_raw_values=state.raw_values[layer], + state_betas=state.betas[layer], + state_cursors=state.cursors, + state_tile_sync_mark=state.tile_sync_mark, + q=q, + k=k, + v=v, + a=raw_gate, + b=raw_beta, + a_log=layer_weight.linear_A_log.weight, + bias=layer_weight.linear_dt_bias.weight, + reqs=self.b_conv_buffer_idx, + cu_seqlens=self.b_cu_seqlens, + state_verify_width=state.verify_width, + state_compute_mode=state.compute_mode, + state_gate_lower_bound=state.gate_lower_bound, + ) + # 各层都读到了本步开始时的 cursor。最后一层返回后再按整窗全接受推进。 + state.finish_decode_token(layer_num, self.b_conv_buffer_idx, self.b_cu_seqlens) + return output + + def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: + infer = self.infer_state + if mtp_step == 0: + self.b_conv_buffer_idx = infer.b_req_idx + self.b1_mtp_cu_q_seq_len = None + self.b_num_accepted_tokens = None + # 每行一个 token。replay 和 cursor 共用这份前缀和。 + self.b_cu_seqlens = torch.arange( + infer.b_req_idx.numel() + 1, dtype=torch.int32, device=infer.b_req_idx.device + ) + return + if self.backend.uses_dynamic_spec_verify_layout(): + cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=mtp_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + self.b1_mtp_cu_q_seq_len = cu + self.b_cu_seqlens = cu + self.b_conv_buffer_idx = conv_idx + self.b_num_accepted_tokens = accepted + return + width = mtp_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "KDA replay decode requires batch_size divisible by mtp_step + 1, " + f"got batch_size={batch_size}, mtp_step={mtp_step}." + ) + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + self.b_cu_seqlens = self.b1_mtp_cu_q_seq_len + self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 + + +def _replay_state(att_state, layer_num: int) -> ReplayKDAState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="kda", + logical_layer_id=layer_num, + state_class=ReplayKDAState, + ) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py new file mode 100644 index 0000000000..3704ba4fd6 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -0,0 +1,246 @@ +"""KDA replay decode。 + +调用方传入这一层的 SSM、history 和 cursor。这里调优并执行一层重放。 +""" + +from __future__ import annotations + +import ast + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import ( + _decode_replay, + check_replay_sync_mark, +) +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune + + +def get_configs() -> list[dict[str, int]]: + """返回 KDA replay 的候选 Triton 配置。""" + return [ + {"BV": bv, "num_warps": warps, "num_stages": stages} + for bv, warps in [(8, 1), (16, 1), (16, 2), (32, 1), (32, 2), (32, 4), (64, 2), (64, 4), (128, 4)] + for stages in (1, 2, 3) + ] + + +def get_static_key( + state_ssm: torch.Tensor, + state_keys: torch.Tensor, + q: torch.Tensor, + state_verify_width: int, + state_compute_mode: str, + state_gate_lower_bound: float, +) -> dict[str, object]: + """区分头维度、状态精度、history 容量、gate 下界和投影模式。""" + _, hv, kd, vd = state_ssm.shape + return { + "v": 1, + "shape": str((q.shape[-2], hv, state_keys.shape[-3], kd, vd, state_verify_width)), + "dtype": str((q.dtype, state_ssm.dtype)), + "capacity": state_keys.shape[-2], + "lower": state_gate_lower_bound, + "recon_dtype": str(state_keys.dtype), + "projection": "kda_mma" if state_compute_mode == "precompute" else "inline", + } + + +def get_run_key(q: torch.Tensor, cu_seqlens: torch.Tensor) -> tuple[int, int]: + """调优键是 token 数和请求数。q 是 [tokens, heads, dim]。""" + return q.shape[0], cu_seqlens.numel() - 1 + + +def _check_replay_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> None: + assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must be [tokens, heads, dim]" + assert q.shape[0] == k.shape[0] == v.shape[0], "q, k, v must cover the same tokens" + assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" + + +def rebuild_inputs( + state_ssm: torch.Tensor, + state_keys: torch.Tensor, + state_deltas: torch.Tensor, + state_gates: torch.Tensor, + state_raw_keys: torch.Tensor, + state_raw_values: torch.Tensor, + state_betas: torch.Tensor, + state_cursors: torch.Tensor, + state_tile_sync_mark: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor, + state_verify_width: int, + state_compute_mode: str, + state_gate_lower_bound: float, + run_config: dict[str, int] | None = None, +) -> tuple[tuple, dict[str, object]]: + """为一层 scratch 和同形状激活分配调优输入。不复制服务状态池。""" + del state_cursors, reqs, run_config + tokens, run_count = get_run_key(q=q, cu_seqlens=cu_seqlens) + # 一行一个 token 时请求数就是 token 数。打包窗口再用 verify_width 折成请求数。 + batch = run_count if tokens == run_count else min(run_count, triton.cdiv(tokens, state_verify_width)) + slots = batch + 1 + _, hv, kd, vd = state_ssm.shape + key_heads, capacity = state_keys.shape[-3], state_keys.shape[-2] + device = q.device + rows_per_slot = state_tile_sync_mark.shape[0] // state_ssm.shape[0] + copied = [ + torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) + ] + bench_reqs = torch.cat( + ( + torch.arange(batch, dtype=torch.int32, device=device), + torch.full((run_count - batch,), batch, dtype=torch.int32, device=device), + ) + ) + # 计时用均匀窗口。宽度为 1 时 verify_width 就是 1。 + cu = torch.arange(run_count + 1, dtype=torch.int32, device=device) + cu = (cu * state_verify_width).clamp_max(tokens) + return (), dict( + state_ssm=torch.empty(slots, hv, kd, vd, device=device, dtype=state_ssm.dtype).normal_(0, 0.02), + state_keys=torch.empty(slots, key_heads, capacity, kd, device=device, dtype=state_keys.dtype), + state_deltas=torch.empty(slots, hv, capacity, vd, device=device, dtype=state_deltas.dtype), + state_gates=torch.empty(slots, hv, capacity, kd, device=device, dtype=state_gates.dtype), + state_raw_keys=torch.empty(slots, key_heads, capacity, kd, device=device, dtype=state_raw_keys.dtype), + state_raw_values=torch.empty(slots, hv, capacity, vd, device=device, dtype=state_raw_values.dtype), + state_betas=torch.empty(slots, hv, capacity, device=device, dtype=state_betas.dtype), + state_cursors=torch.zeros(slots, dtype=torch.int32, device=device), + state_tile_sync_mark=torch.zeros((slots * rows_per_slot, key_heads), dtype=torch.int32, device=device), + q=copied[0], + k=copied[1], + v=copied[2], + a=copied[3], + b=copied[4], + a_log=copied[5], + bias=copied[6], + reqs=bench_reqs, + cu_seqlens=cu, + state_verify_width=state_verify_width, + state_compute_mode=state_compute_mode, + state_gate_lower_bound=state_gate_lower_bound, + ) + + +def _run_key_distance(left: str, right: str) -> int: + return sum(abs(x - y) for x, y in zip(ast.literal_eval(left), ast.literal_eval(right))) + + +@autotune( + kernel_name="replay_kda_decode:v1", + configs_gen_func=get_configs, + static_key_func=get_static_key, + run_key_func=get_run_key, + run_key_distance_func=_run_key_distance, + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + rebuild_input_func=rebuild_inputs, + warmup_all_exist_config=False, +) +def replay( + state_ssm: torch.Tensor, + state_keys: torch.Tensor, + state_deltas: torch.Tensor, + state_gates: torch.Tensor, + state_raw_keys: torch.Tensor, + state_raw_values: torch.Tensor, + state_betas: torch.Tensor, + state_cursors: torch.Tensor, + state_tile_sync_mark: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor, + state_verify_width: int, + state_compute_mode: str, + state_gate_lower_bound: float, + run_config: dict[str, int] | None = None, +) -> torch.Tensor: + """执行一层 KDA replay decode。 + + state_* 张量已经是这一层的视图。ssm 是 [slots, value_heads, key_dim, value_dim]， + keys 是 [slots, key_heads, capacity, key_dim]。最后一个 slot 是 HOLD。 + q、k、v 是 [tokens, heads, dim]。cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。 + 宽度为 1 时它是 0, 1, 2, ...。 + """ + + _check_replay_inputs(q=q, k=k, v=v, reqs=reqs, cu_seqlens=cu_seqlens) + if run_config is None: + run_config = {"BV": 32, "num_warps": 1, "num_stages": 3} + _, hv, kd, vd = state_ssm.shape + key_heads = state_keys.shape[-3] + capacity = state_keys.shape[-2] + assert key_heads in (q.shape[-2], hv) + out = torch.empty_like(v) + bv = run_config["BV"] + # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 + precompute_projections = state_compute_mode == "precompute" and bv >= 16 + padded_verify_width = triton.next_power_of_2(state_verify_width) + check_replay_sync_mark( + sync_mark=state_tile_sync_mark, + reqs=reqs, + value_heads=hv, + key_heads=key_heads, + ) + _decode_replay[(triton.cdiv(vd, bv), hv, reqs.numel())]( + Q=q, + StepK=k, + StepV=v, + A=a, + B=b, + Alog=a_log, + Bias=bias, + State=state_ssm, + Keys=state_keys, + Deltas=state_deltas, + Gates=state_gates, + RawKeys=state_raw_keys, + RawValues=state_raw_values, + Betas=state_betas, + Reqs=reqs, + Cursors=state_cursors, + SyncMark=state_tile_sync_mark, + Cu=cu_seqlens, + Out=out, + SQ=q.stride(0), + SK=k.stride(0), + SV=v.stride(0), + SA=a.stride(0), + SB=b.stride(0), + H=q.shape[-2], + HV=hv, + KH=key_heads, + K=kd, + V=vd, + L=capacity, + HOLD=state_ssm.shape[0] - 1, + BK=triton.next_power_of_2(kd), + BV=bv, + PRECOMPUTE_PROJECTIONS=precompute_projections, + BW=padded_verify_width, + KDA=True, + LOWER=state_gate_lower_bound, + BQK=max(16, 2 * padded_verify_width), + BR=triton.next_power_of_2(max(16, capacity + padded_verify_width)), + NV=triton.cdiv(vd, bv), + num_warps=run_config["num_warps"], + num_stages=run_config["num_stages"], + ) + return out diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py index 8b1d59fc63..878e2296f4 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/index.py @@ -21,6 +21,7 @@ def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: @tensor_cache def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: + # TODO: tolist() 会把 chunk 数同步回 CPU。以后改成在设备上构造索引，去掉这次同步。 indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()]) return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py new file mode 100644 index 0000000000..9b1ba2fc37 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Update the accepted-prefix cursor after replay has returned. + +Replay itself does not touch the cursor: every layer must still see the length +from the start of this step. After the last layer returns, advance the cursor +over the replayed window. MTP verify then trims the rejected tail. +""" + +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["N"]) +def _advance_cursor( + cursors, + reqs, + cu, + N, + HOLD: tl.constexpr, + CAPACITY: tl.constexpr, + BLOCK: tl.constexpr, +): + # 先按整窗全接受。放不下时新 token 占 [0, window)，cursor 记为 window；否则加上 window。 + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(reqs + row, row < N, HOLD) + valid = (row < N) & (req != HOLD) + cursor = tl.load(cursors + req, valid, 0) + window = tl.load(cu + row + 1, valid, 0) - tl.load(cu + row, valid, 0) + updated = tl.where(cursor + window > CAPACITY, window, cursor + window) + tl.store(cursors + req, updated, valid) + + +def advance_cursor(state, reqs, cu_seqlens) -> None: + """最后一层返回后，把 cursor 推进到整窗都被接受的位置。 + + cu_seqlens 与 reqs 逐行对齐，长度是请求数加一。窗口长度为 0 时 cursor 保持不变。HOLD 请求跳过。 + """ + + if reqs.numel() == 0: + return + assert reqs.is_cuda and cu_seqlens.is_cuda, "advance_cursor requires CUDA inputs" + assert cu_seqlens.numel() == reqs.numel() + 1, "cu_seqlens must cover every request" + _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( + state.cursors, + reqs, + cu_seqlens, + reqs.numel(), + state.hold_req_id, + state.capacity, + 256, + num_warps=4, + ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py new file mode 100644 index 0000000000..f3abbeb020 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -0,0 +1,820 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM decode kernels. + +`_decode_replay` is the entry. When the accepted prefix does not fit, it folds +that prefix into the SSM and then replays the new tokens from the rounded +state. The output is either a precomputed GDN or KDA window, or a token walk. +""" + +import torch +import triton +import triton.language as tl + +from .replayssm_delta import _apply_accepted_prefix + + +def check_replay_sync_mark( + sync_mark: torch.Tensor, + reqs: torch.Tensor, + value_heads: int, + key_heads: int, +) -> None: + """确认 SyncMark 按 [launch 行, key head] 连续存放，并且盖住这次启动的每一行。""" + + assert key_heads > 0 and value_heads % key_heads == 0, "value heads must be a multiple of key heads" + assert sync_mark.dtype == torch.int32 and sync_mark.dim() == 2, "sync mark must be a 2D int32 tensor" + assert sync_mark.shape[1] == key_heads, "sync mark columns must equal key heads" + assert sync_mark.shape[0] >= reqs.numel(), "sync mark rows must cover every launch row" + assert sync_mark.stride(1) == 1 and sync_mark.stride(0) == key_heads, "sync mark must be contiguous by key head" + + +@triton.jit +def _arrive_value_tiles(SyncMark, seq, hv, HV: tl.constexpr, KH: tl.constexpr, NV: tl.constexpr): + # 同一 launch 行、同一 key head 的全部 value head 和 V tile 都读完旧 RawKeys 后，才允许写新窗口。 + # 行号不是请求槽。HOLD 会在一批里重复出现，用请求槽做下标会撞到同一格。 + # SyncMark 每行长度是 KH。最后到达的 tile 把计数写回 0，这个 0 就是放行信号。 + group: tl.constexpr = HV // KH + key_head = hv // group + idx = seq * KH + key_head + arrived = tl.atomic_add(SyncMark + idx, 1, sem="acq_rel", scope="gpu") + if arrived == NV * group - 1: + tl.atomic_xchg(SyncMark + idx, 0, sem="release", scope="gpu") + else: + seen = arrived + 1 + while seen != 0: + seen = tl.atomic_add(SyncMark + idx, 0, sem="acquire", scope="gpu") + + +@triton.jit +def _history_coeff(rows, vec, BK: tl.constexpr): + # rows [R, BK] · vec [BK] -> [R]. tf32x3 是 tensor core 上最接近 fp32 的输出重建。 + if BK >= 16: + return tl.sum(tl.dot(rows, vec[:, None], input_precision="tf32x3"), 1) + return tl.sum(rows * vec[None, :], 1) + + +@triton.jit +def _apply_deltas(coeff, deltas, R: tl.constexpr, BV: tl.constexpr): + # coeff [R] · deltas [R, BV] -> [BV]. R 小于 16 时 tensor core 不能做这次点积。 + if R >= 16 and BV >= 16: + return tl.sum(tl.dot(coeff[None, :], deltas, input_precision="tf32x3"), 0) + return tl.sum(deltas * coeff[:, None], 0) + + +@triton.jit +def _decode_replay( + Q, + StepK, + StepV, + A, + B, + Alog, + Bias, + State, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Reqs, + Cursors, + SyncMark, + Cu, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + PRECOMPUTE_PROJECTIONS: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BQK: tl.constexpr, + BR: tl.constexpr, + NV: tl.constexpr, +): + """执行一个请求窗口的 replay decode，并写入 history。 + + 前缀放不下本步 token 时，先在本 kernel 里把已有前缀从原始 k/v 折进 checkpoint， + 舍入成 SSM 的 dtype，再从这份状态重放本步。cursor 本 kernel 不修改，留给最后一层返回后更新。 + """ + # grid 是 (V tile, value head, launch 行)。同一 V tile 上相邻 program 对应连续的 value head。 + iv, hv, seq = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + seq).to(tl.int64) + start, end = tl.load(Cu + seq), tl.load(Cu + seq + 1) + vv = iv * BV + tl.arange(0, BV) + if req == HOLD: + for t in range(start, end): + tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) + return + if start == end: + # 本步没有 token。cursor 保持原值，接受 0 个 token 时长度不变。 + return + n = tl.load(Cursors + req) + + kk = tl.arange(0, BK) + ll = tl.arange(0, L) + slot = req * HV + hv + key_slot = req * KH + hv // (HV // KH) + sp = State + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + # 放不下时先折回，再在同一 kernel 里重放新 token。各层只读 cursor。 + if n + (end - start) > L: + state = _apply_accepted_prefix(state, RawKeys, RawValues, Gates, Betas, key_slot, slot, kk, vv, n, K, V, L, KDA) + state = state.to(State.dtype.element_ty).to(tl.float32) + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + n = 0 + # 新 token 会从槽 0 覆盖刚读过的 RawKeys。等同一 key head 的全部 value tile 都读完再往下走。 + _arrive_value_tiles(SyncMark, seq, hv, HV, KH, NV) + + if PRECOMPUTE_PROJECTIONS: + # 输出用 tf32x3 tensor core 从冻结的 checkpoint 和 history 重建。 + # 状态不在这里更新；压缩仍用原始 k/v 做 fp32 串行折回。 + projection_dtype: tl.constexpr = tl.float32 + projection_precision: tl.constexpr = "tf32x3" + qk, projections, tt, tokens, valid, h = _project_qk( + Q, + StepK, + A, + Alog, + Bias, + Gates, + SQ, + SK, + SA, + H, + HV, + K, + L, + BW, + KDA, + LOWER, + BQK, + end, + hv, + kk, + ll, + n, + projection_dtype, + projection_precision, + slot, + start, + state, + ) + + if PRECOMPUTE_PROJECTIONS and KDA: + _kda_precomputed_window( + StepK, + StepV, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + BR, + BK, + LOWER, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projections, + qk, + slot, + start, + tt, + vv, + ) + elif PRECOMPUTE_PROJECTIONS: + _gdn_precomputed_window( + StepK, + StepV, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + BR, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype, + projection_precision, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, + ) + else: + _replay_each_token( + Q, + StepK, + StepV, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ, + SK, + SV, + SA, + SB, + H, + HV, + KH, + K, + V, + L, + KDA, + LOWER, + BK, + BV, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + slot, + start, + state, + vv, + ) + + +@triton.jit +def _project_qk( + Q, + StepK, + A, + Alog, + Bias, + Gates, + SQ: tl.constexpr, + SK: tl.constexpr, + SA: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + L: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BQK: tl.constexpr, + end, + hv, + kk, + ll, + n, + projection_dtype: tl.constexpr, + projection_precision: tl.constexpr, + slot, + start, + state, +): + # 预计算 verify window 的 q/k 投影。 + tt = tl.arange(0, BQK) + tokens = start + tt % BW + valid = (tt < 2 * BW) & (tokens < end) + safe_tokens = tl.where(valid, tokens, start) + h = hv // (HV // H) + q_rows = tl.load( + Q + safe_tokens[:, None] * SQ + h * K + kk[None, :], + (valid & (tt < BW))[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + k_rows = tl.load( + StepK + safe_tokens[:, None] * SK + h * K + kk[None, :], + (valid & (tt >= BW))[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + # Padded query/key rows share the MMA tile. Keep them at zero so they + # cannot change the valid rows. + row_valid = valid[:, None] + qk = tl.where(row_valid, q_rows + k_rows, 0) + qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] + qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] + qk = tl.where(row_valid, qk, 0) + if KDA: + history = ll < n + gates = tl.load( + Gates + (slot * L + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + total_g = tl.sum(gates, axis=0) + wi = tl.arange(0, BW) + gate_ok = start + wi < end + safe_wi = tl.where(gate_ok, wi, 0) + gate_inputs = tl.load( + A + (start + safe_wi[:, None]) * SA + hv * K + kk[None, :], + gate_ok[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] + fresh_gates = LOWER * tl.sigmoid(tl.exp(tl.load(Alog + hv)) * gate_inputs) + fresh_gates = tl.where((start + wi < end)[:, None], fresh_gates, 0) + query_gates = total_g[None, :] + tl.gather( + tl.cumsum(fresh_gates, axis=0), + tl.broadcast_to((tt % BW)[:, None], qk.shape), + 0, + ) + # All exponents are forward decay. No inverse prefix scaling can + # overflow when a long history contains strongly negative gates. + # Padded rows stay zero: 0 * exp(garbage) is NaN and would enter the dot. + scaled = tl.where(row_valid, qk * tl.exp(query_gates), 0) + projections = tl.dot( + scaled.to(projection_dtype), + state.to(projection_dtype), + input_precision=projection_precision, + ) + else: + projections = tl.dot(qk.to(projection_dtype), state.to(projection_dtype), input_precision=projection_precision) + return qk, projections, tt, tokens, valid, h + + +@triton.jit +def _gdn_precomputed_window( + StepK, + StepV, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BV: tl.constexpr, + BW: tl.constexpr, + BR: tl.constexpr, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype: tl.constexpr, + projection_precision: tl.constexpr, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, +): + # GDN precompute 路径：一次投影整个窗口，再按 token 顺序写入新 history。 + # 只有 corrected-value solve 依赖 token 顺序，避免重复计算历史投影。 + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + # BR is rounded above the history length. Masked lanes must still address + # a real slot; otherwise they read uninitialized neighbor rows into the dot. + safe_jj = tl.where(jj < L, jj, 0) + safe_fresh = tl.where(fresh, fresh_tokens, start) + cached_keys = tl.load( + Keys + (key_slot * L) * K + safe_jj[:, None] * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + raw_keys = tl.load( + StepK + safe_fresh[:, None] * SK + h * K + kk[None, :], + fresh[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + fresh_keys = raw_keys * tl.rsqrt(tl.sum(raw_keys * raw_keys, axis=1) + 1.0e-6)[:, None] + fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) + all_keys = cached_keys + fresh_keys + values = tl.load( + StepV + safe_fresh[:, None] * SV + hv * V + vv[None, :], + fresh[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + log_a, bias = tl.load(Alog + hv).to(tl.float32), tl.load(Bias + hv).to(tl.float32) + x = tl.load(A + safe_fresh * SA + hv, fresh, 0).to(tl.float32) + bias + fresh_gates = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + gates = tl.load(Gates + slot * L + safe_jj, old, 0).to(tl.float32) + gates += tl.where(fresh, fresh_gates, 0.0) + prefix = tl.cumsum(gates, axis=0) + query_gates = tl.sum(tl.where(jj[None, :] <= L + (tt % BW)[:, None], gates[None, :], 0), 1) + coefficients = tl.dot( + qk.to(projection_dtype), tl.trans(all_keys).to(projection_dtype), input_precision=projection_precision + ) + causal = old[None, :] | (fresh[None, :] & (jj[None, :] <= L + (tt % BW)[:, None])) + coefficients *= tl.where(causal, tl.exp(tl.minimum(query_gates[:, None] - prefix[None, :], 0.0)), 0.0) + projections *= tl.exp(query_gates)[:, None] + deltas = tl.load( + Deltas + (slot * L) * V + safe_jj[:, None] * V + vv[None, :], + old[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + betas = tl.sigmoid(tl.load(B + safe_fresh * SB + hv, fresh, 0).to(tl.float32)) + for t in range(end - start): + checkpoint_sk = tl.sum(tl.where((tt == BW + t)[:, None], projections, 0), 0) + weights = tl.sum(tl.where((tt == BW + t)[:, None], coefficients, 0), 0) + sk = checkpoint_sk + _apply_deltas(weights, deltas, BR, BV) + value = tl.sum(tl.where((jj == L + t)[:, None], values, 0), 0) + beta = tl.sum(tl.where(jj == L + t, betas, 0)) + delta = beta * (value - sk) + deltas = tl.where((jj == L + t)[:, None], delta[None, :], deltas) + outputs = projections + tl.dot( + coefficients.to(projection_dtype), deltas.to(projection_dtype), input_precision=projection_precision + ) + out_tokens = tl.where(valid, tokens, start) + tl.store( + Out + (out_tokens[:, None] * HV + hv) * V + vv[None, :], + outputs, + (valid & (tt < BW))[:, None] & (vv[None, :] < V), + ) + records = n + jj - L + in_prefix = fresh & (records < L) + safe_records = tl.where(in_prefix, records, 0) + tl.store( + Deltas + (slot * L + safe_records[:, None]) * V + vv[None, :], + deltas, + in_prefix[:, None] & (vv[None, :] < V), + ) + tl.store( + RawValues + (slot * L + safe_records[:, None]) * V + vv[None, :], + values, + in_prefix[:, None] & (vv[None, :] < V), + ) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store( + Keys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + fresh_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + RawKeys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + raw_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store(Gates + slot * L + safe_records, fresh_gates, in_prefix) + tl.store(Betas + slot * L + safe_records, betas, in_prefix) + return + + +@triton.jit +def _kda_precomputed_window( + StepK, + StepV, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BV: tl.constexpr, + BW: tl.constexpr, + BR: tl.constexpr, + BK: tl.constexpr, + LOWER: tl.constexpr, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projections, + qk, + slot, + start, + tt, + vv, +): + # KDA 的衰减在每个 key 通道上，不能把标量 exp 提到点积外面。 + # 每个 token 用冻结 checkpoint 的投影，加上 history 的通道衰减点积重建输出。 + # exp(G_t - G_j) 对已可见的 token 不大于 1。新 delta 仍按 token 解。 + # 状态压缩继续用原始 k/v 在 fp32 里折回。 + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + safe_jj = tl.where(jj < L, jj, 0) + safe_fresh = tl.where(fresh, fresh_tokens, start) + cached_keys = tl.load( + Keys + (key_slot * L) * K + safe_jj[:, None] * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + raw_keys = tl.load( + StepK + safe_fresh[:, None] * SK + h * K + kk[None, :], + fresh[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + fresh_keys = raw_keys * tl.rsqrt(tl.sum(raw_keys * raw_keys, axis=1) + 1.0e-6)[:, None] + fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) + all_keys = tl.where(old[:, None], cached_keys, 0) + fresh_keys + values = tl.load( + StepV + safe_fresh[:, None] * SV + hv * V + vv[None, :], + fresh[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + hist_g = tl.load( + Gates + (slot * L + safe_jj[:, None]) * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + total_hist = tl.sum(hist_g, axis=0) + hist_prefix = tl.cumsum(hist_g, axis=0) + wi = tl.arange(0, BW) + gate_ok = start + wi < end + safe_wi = tl.where(gate_ok, wi, 0) + gate_inputs = tl.load( + A + (start + safe_wi[:, None]) * SA + hv * K + kk[None, :], + gate_ok[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] + fresh_g_rows = LOWER * tl.sigmoid(tl.exp(tl.load(Alog + hv).to(tl.float32)) * gate_inputs) + fresh_g_rows = tl.where(gate_ok[:, None], fresh_g_rows, 0) + fresh_prefix = tl.cumsum(fresh_g_rows, axis=0) + fresh_index = tl.broadcast_to(tl.where(fresh, jj - L, 0)[:, None], (BR, BK)) + G_pos = tl.where(old[:, None], hist_prefix, 0.0) + G_pos = tl.where(fresh[:, None], total_hist[None, :] + tl.gather(fresh_prefix, fresh_index, 0), G_pos) + step_g = tl.where(fresh[:, None], tl.gather(fresh_g_rows, fresh_index, 0), 0.0) + G_query = total_hist[None, :] + tl.gather(fresh_prefix, tl.broadcast_to((tt % BW)[:, None], qk.shape), 0) + deltas = tl.load( + Deltas + (slot * L) * V + safe_jj[:, None] * V + vv[None, :], + old[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + betas = tl.sigmoid(tl.load(B + safe_fresh * SB + hv, fresh, 0).to(tl.float32)) + for t in range(start, end): + local = t - start + q = tl.sum(tl.where((tt == local)[:, None], qk, 0), 0) + k = tl.sum(tl.where((tt == BW + local)[:, None], qk, 0), 0) + g_now = tl.sum(tl.where((tt == local)[:, None], G_query, 0), 0) + causal = old | (fresh & (jj <= L + local)) + decay = tl.exp(tl.where(causal[:, None], g_now[None, :] - G_pos, 0.0)) + scaled = tl.where(causal[:, None], all_keys * decay, 0) + hk = _history_coeff(scaled, k, BK) + proj_k = tl.sum(tl.where((tt == BW + local)[:, None], projections, 0), 0) + sk = proj_k + _apply_deltas(hk, deltas, BR, BV) + value = tl.sum(tl.where((jj == L + local)[:, None], values, 0), 0) + beta = tl.sum(tl.where(jj == L + local, betas, 0)) + delta = beta * (value - sk) + deltas = tl.where((jj == L + local)[:, None], delta[None, :], deltas) + hq = _history_coeff(scaled, q, BK) + proj_q = tl.sum(tl.where((tt == local)[:, None], projections, 0), 0) + sq = proj_q + _apply_deltas(hq, deltas, BR, BV) + tl.store(Out + (t * HV + hv) * V + vv, sq, vv < V) + records = n + jj - L + in_prefix = fresh & (records < L) + safe_records = tl.where(in_prefix, records, 0) + tl.store( + Deltas + (slot * L + safe_records[:, None]) * V + vv[None, :], + deltas, + in_prefix[:, None] & (vv[None, :] < V), + ) + tl.store( + RawValues + (slot * L + safe_records[:, None]) * V + vv[None, :], + values, + in_prefix[:, None] & (vv[None, :] < V), + ) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store( + Keys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + fresh_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + RawKeys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + raw_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + Gates + (slot * L + safe_records[:, None]) * K + kk[None, :], + step_g, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store(Betas + slot * L + safe_records, betas, in_prefix) + return + + +@triton.jit +def _replay_each_token( + Q, + StepK, + StepV, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + slot, + start, + state, + vv, +): + # 逐 token 递推。precompute 走单独的窗口 kernel，不会进这里。 + history = ll < n + keys = tl.load( + Keys + (key_slot * L) * K + ll[:, None] * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + deltas = tl.load( + Deltas + (slot * L) * V + ll[:, None] * V + vv[None, :], + history[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + if KDA: + gates = tl.load( + Gates + (slot * L + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + else: + gates = tl.load(Gates + slot * L + ll, history, 0).to(tl.float32) + gate_prefix = tl.cumsum(gates, axis=0) + total_g = tl.sum(gates, axis=0) + if KDA: + keys *= tl.where(history[:, None], tl.exp(total_g[None, :] - gate_prefix), 0.0) + h = hv // (HV // H) + log_a = tl.load(Alog + hv).to(tl.float32) + bias = tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32) if KDA else tl.load(Bias + hv).to(tl.float32) + + for t in range(start, end): + q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) + raw_k = tl.load(StepK + t * SK + h * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(StepV + t * SV + hv * V + vv, vv < V, 0).to(tl.float32) + q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K ** -0.5) + k = raw_k / tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + if KDA: + x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias + g = LOWER * tl.sigmoid(tl.exp(log_a) * x) + else: + x = tl.load(A + t * SA + hv).to(tl.float32) + bias + g = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + beta = tl.sigmoid(tl.load(B + t * SB + hv).to(tl.float32)) + total_g += g + if KDA: + # Advance the decayed keys once per token, rather than + # evaluating L independent vector exponentials each time. + keys *= tl.exp(g)[None, :] + hk = _history_coeff(keys, k, BK) + hq = _history_coeff(keys, q, BK) + checkpoint_sk = tl.sum(state * (k * tl.exp(total_g))[:, None], axis=0) + checkpoint_sq = tl.sum(state * (q * tl.exp(total_g))[:, None], axis=0) + sk = checkpoint_sk + _apply_deltas(hk, deltas, L, BV) + sq = checkpoint_sq + _apply_deltas(hq, deltas, L, BV) + else: + weights = tl.where(history, tl.exp(total_g - gate_prefix), 0.0) + hk = _history_coeff(keys, k, BK) * weights + hq = _history_coeff(keys, q, BK) * weights + checkpoint_sk = tl.sum(state * k[:, None], axis=0) + checkpoint_sq = tl.sum(state * q[:, None], axis=0) + sk = checkpoint_sk * tl.exp(total_g) + _apply_deltas(hk, deltas, L, BV) + sq = checkpoint_sq * tl.exp(total_g) + _apply_deltas(hq, deltas, L, BV) + d = beta * (v - sk) + tl.store(Out + (t * HV + hv) * V + vv, sq + d * tl.sum(k * q), vv < V) + record = n + in_prefix = record < L + # 被 mask 掉的 lane 仍会计算地址。cursor 已经顶到容量时必须改写 0 号槽，避免非法访问。 + safe_record = tl.where(in_prefix, record, 0) + tl.store(Deltas + (slot * L + safe_record) * V + vv, d, in_prefix & (vv < V)) + tl.store(RawValues + (slot * L + safe_record) * V + vv, v, in_prefix & (vv < V)) + if iv == 0: + # Only the group owner writes shared K. Every CTA computes this + # round's K locally; readers only consume earlier accepted records. + if hv % (HV // KH) == 0: + tl.store(Keys + (key_slot * L + safe_record) * K + kk, k, in_prefix & (kk < K)) + tl.store(RawKeys + (key_slot * L + safe_record) * K + kk, raw_k, in_prefix & (kk < K)) + if KDA: + tl.store(Gates + (slot * L + safe_record) * K + kk, g, in_prefix & (kk < K)) + else: + tl.store(Gates + slot * L + safe_record, g, in_prefix) + tl.store(Betas + slot * L + safe_record, beta, in_prefix) + keys = tl.where((ll == n)[:, None], k[None, :], keys) + deltas = tl.where((ll == n)[:, None], d[None, :], deltas) + if not KDA: + gate_prefix = tl.where(ll >= n, total_g, gate_prefix) + history |= ll == n + n += 1 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py new file mode 100644 index 0000000000..e832c6f3f3 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +"""One gated-delta step shared by prefix folding.""" + +import triton +import triton.language as tl + + +@triton.jit +def _delta_rule_step(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): + # Folding applies this same update once per accepted history token. + raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + decay = tl.exp(g) + if KDA: + state *= decay[:, None] + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + else: + sk = tl.sum(state * raw_k[:, None], 0) + d = beta * (raw_v - decay * sk) + state = state * decay + raw_k[:, None] * d[None, :] + return state + + +@triton.jit +def _apply_accepted_prefix( + state, + RawKeys, + RawValues, + Gates, + Betas, + key_slot, + slot, + kk, + vv, + prefix_len, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + KDA: tl.constexpr, +): + # 按 history 槽位顺序把已接受前缀折进 state。j >= prefix_len 的槽不读、不更新。 + for j in range(L): + active = j < prefix_len + raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) + if KDA: + g = tl.load(Gates + (slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + else: + g = tl.load(Gates + slot * L + j, active, 0).to(tl.float32) + beta = tl.load(Betas + slot * L + j, active, 0).to(tl.float32) + updated = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) + state = tl.where(active, updated, state) + return state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py new file mode 100644 index 0000000000..42aad6edc2 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fold an accepted prefix into the canonical SSM.""" + +import triton +import triton.language as tl + +from .replayssm_delta import _apply_accepted_prefix + + +@triton.jit +def _fold_prefix( + State, + RawKeys, + RawValues, + Gates, + Betas, + Cursors, + Reqs, + SLOTS: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + KDA: tl.constexpr, + STRIDE_LAYER: tl.constexpr, + STRIDE_SLOT: tl.constexpr, + STRIDE_HEAD: tl.constexpr, + STRIDE_K: tl.constexpr, + STRIDE_V: tl.constexpr, +): + # 将 cursor 指向的 accepted history 写回 canonical state。 + iv, row, lh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + row).to(tl.int64) + if req == HOLD: + return + prefix_len = tl.load(Cursors + req) + if prefix_len == 0: + return + linear_layer_idx, head = lh // HV, lh % HV + # history 是连续分配的。canonical SSM 可能带大于 int32 的层距，地址必须用真实 stride 的 int64。 + layer64 = linear_layer_idx.to(tl.int64) + head64 = head.to(tl.int64) + slot = (linear_layer_idx * SLOTS + req) * HV + head + key_slot = (linear_layer_idx * SLOTS + req) * KH + head // (HV // KH) + kk = tl.arange(0, BK) + vv = iv * BV + tl.arange(0, BV) + sp = ( + State + + layer64 * STRIDE_LAYER + + req * STRIDE_SLOT + + head64 * STRIDE_HEAD + + kk[:, None].to(tl.int64) * STRIDE_K + + vv[None, :].to(tl.int64) * STRIDE_V + ) + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + # cursor 可能暂时超过容量。循环只走容量内的槽，越界步不读 history，也不更新 state。 + state = _apply_accepted_prefix( + state, RawKeys, RawValues, Gates, Betas, key_slot, slot, kk, vv, prefix_len, K, V, L, KDA + ) + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + # cursor 不在这里清。和折回并行的其他线程块还要读原来的长度，清零放到 kernel 结束之后。 + + +def fold_prefix(state, reqs, run_config=None) -> None: + """把 cursor 覆盖的已接受前缀折进 canonical SSM。 + + 没传入 run_config 时用 BV=32、4 个 warp、3 个 stage。 + """ + + if run_config is None: + run_config = {"BV": 32, "num_warps": 4, "num_stages": 3} + bv = run_config["BV"] + num_warps = run_config["num_warps"] + num_stages = run_config["num_stages"] + ssm = state.ssm_state + layers, slots, hv, k, v = ssm.shape + layer_stride, slot_stride, head_stride, key_stride, value_stride = ssm.stride() + _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + ssm, + state.raw_keys, + state.raw_values, + state.gates, + state.betas, + state.cursors, + reqs, + slots, + hv, + state.num_key_heads, + k, + v, + state.capacity, + state.hold_req_id, + triton.next_power_of_2(k), + bv, + state.kda, + STRIDE_LAYER=layer_stride, + STRIDE_SLOT=slot_stride, + STRIDE_HEAD=head_stride, + STRIDE_K=key_stride, + STRIDE_V=value_stride, + num_warps=num_warps, + num_stages=num_stages, + ) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py new file mode 100644 index 0000000000..bf4a61e95a --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +def commit_ssm_prefix( + mtp_state_index: torch.Tensor, + reqs: torch.Tensor, + accepted_lengths: torch.Tensor, + req_verify_width: torch.Tensor, + replay_cursors: torch.Tensor, +) -> None: + """写入 MTP 索引，并从已按验证宽度全接受推进的 cursor 上剪掉未接受尾部。 + + accepted_lengths 含基准 token。req_verify_width 与 reqs 逐行对齐，是每个请求这次送去验证的 token 数。 + """ + + assert reqs.is_cuda, "commit_ssm_prefix requires CUDA requests" + if reqs.numel(): + _commit_prefix[(triton.cdiv(reqs.numel(), 256),)]( + mtp_state_index, + reqs, + accepted_lengths, + replay_cursors, + req_verify_width, + reqs.numel(), + mtp_state_index.numel() - 1, + reqs.stride(0), + accepted_lengths.stride(0), + req_verify_width.stride(0), + 256, + num_warps=4, + ) + + +@triton.jit(do_not_specialize=["N"]) +def _commit_prefix( + MtpIndex, + Reqs, + Accepted, + Cursors, + ReqVerifyWidth, + N, + HOLD: tl.constexpr, + REQ_STRIDE: tl.constexpr, + ACCEPTED_STRIDE: tl.constexpr, + VERIFY_STRIDE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(Reqs + row * REQ_STRIDE, row < N, 0) + accepted = tl.load(Accepted + row * ACCEPTED_STRIDE, row < N, 0) + tl.store(MtpIndex + req, accepted - 1, row < N) + valid = (row < N) & (req != HOLD) + verify_tokens = tl.load(ReqVerifyWidth + row * VERIFY_STRIDE, valid, 0) + cursor = tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, cursor - (verify_tokens - accepted), valid) diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 0e52e736a9..31d79697f6 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -14,15 +14,17 @@ @dataclasses.dataclass class Glm5NextResourceConfig(Qwen3NextResourceConfig): - """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" + """描述复制式 MLA/index KV 与 TP 分片 KDA checkpoint 的资源配置。""" # Indexer 池化步长，单位为原始 token；同时决定 K-pool tail 的基准 ring 长度和恢复边界。 + gate_lower_bound: float = -5.0 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": + """从模型配置和启动参数构造 GLM5.3 Flash 资源配置。""" config = config.get("text_config", config) linear = config["linear_attn_config"] tp = args.tp // args.dp @@ -45,6 +47,7 @@ def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), full_attention_interval=4, all_layer_num=layers, + gate_lower_bound=linear.get("gate_lower_bound", -5.0), index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], ) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index e15efe0143..3f20ad97ab 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -6,7 +6,6 @@ import triton 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.models.glm5_next.config import Glm5NextResourceConfig from lightllm.distributed.communication_op import dist_group_manager @@ -79,6 +78,9 @@ def create_bounded_att_states( request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ssm_state_mode=args.ssm_state_mode, + ssm_replay_config=args.ssm_replay_config, + gate_lower_bound=cfg.gate_lower_bound, ) ) full_layers = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) @@ -114,37 +116,9 @@ def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, ob config["autotune_layer_num"] = 4 return config - def _init_config(self): - with open(os.path.join(self.weight_dir_, "config.json")) as f: - 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 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 - def _init_some_value(self): - self.layers_num = self.config["n_layer"] - self.vocab_size = self.config["vocab_size"] - # MLA stores one replicated latent KV vector per token on every TP rank. - self.tp_k_head_num_ = 1 - self.tp_v_head_num_ = 0 - self.qk_nope_head_dim = self.config["qk_nope_head_dim"] - self.qk_rope_head_dim = self.config["qk_rope_head_dim"] - self.q_lora_rank = self.config["q_lora_rank"] - self.kv_lora_rank = self.config["kv_lora_rank"] - 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 - @classmethod def create_full_att_states( cls, @@ -155,7 +129,10 @@ def create_full_att_states( 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.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) @@ -188,8 +165,38 @@ def prepare_attention_states(self) -> None: # Reserve fixed scratch space before the shared KV capacity is profiled. Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) + def _init_config(self): + with open(os.path.join(self.weight_dir_, "config.json")) as f: + 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 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 _init_some_value(self): + self.layers_num = self.config["n_layer"] + self.vocab_size = self.config["vocab_size"] + # MLA stores one replicated latent KV vector per token on every TP rank. + self.tp_k_head_num_ = 1 + self.tp_v_head_num_ = 0 + self.qk_nope_head_dim = self.config["qk_nope_head_dim"] + self.qk_rope_head_dim = self.config["qk_rope_head_dim"] + self.q_lora_rank = self.config["q_lora_rank"] + self.kv_lora_rank = self.config["kv_lora_rank"] + 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_att_backend(self): - from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( + Fp8Glm5NextSparseAttBackend, + ) backend_class = ( Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend @@ -198,7 +205,12 @@ def _init_att_backend(self): self.decode_att_backend = self.prefill_att_backend def _init_att_backend1(self): - self.prefill_att_backend1 = KDALinearAttBackend(model=self) + from lightllm.common.basemodel.attention.create_linear_utils import ( + get_kda_linear_att_backend_class, + ) + + backend_cls = get_kda_linear_att_backend_class() + self.prefill_att_backend1 = backend_cls(model=self) self.decode_att_backend1 = self.prefill_att_backend1 def _init_custom(self): diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index f1a3a65748..57de61241b 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -81,26 +81,14 @@ def create_bounded_att_states( request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ssm_state_mode=args.ssm_state_mode, + ssm_replay_config=args.ssm_replay_config, ) 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) - - # Set Triton allocator for TMA descriptors - # This is required for kernels in common/.../linear_att/fla/ops/solve_tril.py - triton.set_allocator(_triton_allocator) - logger.info("Triton allocator set for Qwen3Next model") - return - def autotune_layers(self): return self.config["full_attention_interval"] - def _init_config(self): - super()._init_config() - self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - @classmethod def create_full_att_states( cls, @@ -135,6 +123,20 @@ def create_full_att_states( ) 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) + + # Set Triton allocator for TMA descriptors + # This is required for kernels in common/.../linear_att/fla/ops/solve_tril.py + triton.set_allocator(_triton_allocator) + logger.info("Triton allocator set for Qwen3Next model") + return + + 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_att_backend1(self): role = self.att_state_role if not self.att_state_manager.get_bounded_attention_states(role): @@ -142,12 +144,12 @@ def _init_att_backend1(self): self.decode_att_backend1 = None return from lightllm.common.basemodel.attention.create_linear_utils import ( - get_qwen35_linear_prefill_att_backend_class, - get_qwen35_linear_decode_att_backend_class, + get_gdn_linear_prefill_att_backend_class, + get_gdn_linear_decode_att_backend_class, ) - linear_prefill_att_backend_class = get_qwen35_linear_prefill_att_backend_class(index=1) - linear_decode_att_backend_class = get_qwen35_linear_decode_att_backend_class(index=1) + linear_prefill_att_backend_class = get_gdn_linear_prefill_att_backend_class(index=1) + linear_decode_att_backend_class = get_gdn_linear_decode_att_backend_class(index=1) self.prefill_att_backend1 = linear_prefill_att_backend_class(model=self) self.decode_att_backend1 = linear_decode_att_backend_class(model=self) return diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 41ae769ae6..16a8735a7c 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -1,4 +1,5 @@ import argparse +import json def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: @@ -984,6 +985,21 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: If radix cache hit rate is low under high load due to limited small-page capacity and LRU eviction, increasing linear_att_cache_size can improve hit rate at the cost of more memory.""", ) + parser.add_argument( + "--ssm_state_mode", + choices=["native", "replay"], + default="native", + help="SSM state storage: native snapshots or deferred GDN/KDA ReplaySSM (FP32/BF16).", + ) + parser.add_argument( + "--ssm_replay_config", + type=json.loads, + default={}, + help='Replay-only JSON object: {"cache_len":16,"compute_mode":"precompute"}. ' + "cache_len: 4/8/16/32/64, at least the verify width; compute_mode: inline/precompute. " + "mtp_step=0 decode uses the same replay history: explicit precompute projects that one token, " + "and the default is inline. MTP defaults to precompute.", + ) parser.add_argument( "--linear_att_ssm_data_type", type=str, diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index a42acef884..a4c990c2be 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -355,6 +355,8 @@ def _launch_subprocesses(args: StartArgs): s3_model_prepare(args.model_dir) + _validate_and_normalize_ssm_config(args) + # 如果args.eos_id 是 None, 从 config.json 中读取 eos_token_id 相关的信息，赋值给 args if args.eos_id is None: from lightllm.utils.config_utils import get_eos_token_ids @@ -744,3 +746,54 @@ def _auto_set_prefill_queue_strategy(args: StartArgs) -> None: if args.prefill_queue_strategy == "auto": args.prefill_queue_strategy = "hrrn" if args.run_mode == "prefill" else "default" logger.info(f"run_mode={args.run_mode}, auto set prefill_queue_strategy to {args.prefill_queue_strategy}") + + +def _validate_and_normalize_ssm_config(args: StartArgs): + assert args.ssm_state_mode in ("native", "replay"), "Invalid ssm_state_mode" + config = args.ssm_replay_config + assert isinstance(config, dict), "ssm_replay_config must be a JSON object" + assert args.ssm_state_mode == "replay" or not config, "ssm_replay_config requires ssm_state_mode=replay" + if args.ssm_state_mode == "replay": + defaults = {"cache_len": 16, "compute_mode": "precompute" if args.mtp_step > 0 else "inline"} + unknown = config.keys() - defaults.keys() + assert not unknown, f"Unknown ssm_replay_config fields: {unknown}" + config = {**defaults, **config} + assert type(config["cache_len"]) is int and config["cache_len"] in ( + 4, + 8, + 16, + 32, + 64, + ), "ReplaySSM capacity must be a power of two and at least 4, up to 64" + assert config["compute_mode"] in ("inline", "precompute"), "Invalid replay compute_mode" + args.ssm_replay_config = config + logger.info(f"SSM mode: replay, config: {config}") + else: + logger.info(f"SSM mode: {args.ssm_state_mode}") + from lightllm.utils.config_utils import get_model_type + + model_type = get_model_type(args.model_dir) + is_gdn_or_kda = model_type in ( + "glm5_next", + "glm5_next_text", + "qwen3_next", + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "sensenova_loop", + ) + if is_gdn_or_kda: + assert args.mtp_step <= 15, ( + "GDN/KDA requires mtp_step <= 15: " "linear attention MTP state acceptance supports verify width <= 16" + ) + if args.ssm_state_mode == "native": + return + assert args.linear_att_ssm_data_type in ( + "float32", + "bfloat16", + ), "Non-native SSM state modes require FP32 or BF16 state" + assert is_gdn_or_kda, "Non-native SSM state modes currently require a GDN or KDA model" + assert args.mtp_step >= 0, "ssm_state_mode=replay requires mtp_step >= 0" + capacity = config["cache_len"] + assert capacity >= args.mtp_step + 1, "ReplaySSM capacity must cover the verify width" diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0d6468f397..0609132e62 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -255,6 +255,8 @@ class StartArgs: weight_version: str = "default" # hybrid attention model (Qwen3Next) + ssm_state_mode: str = "native" + ssm_replay_config: dict[str, object] = field(default_factory=dict) linear_att_hash_page_size: int = field(default=512) linear_att_page_block_num: int = field(default=10000000) disable_linear_att_small_page_cpu_cache: bool = field(default=False) diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index a32cedd17d..19d4656f14 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -32,12 +32,32 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) + # 最后一层已经按这个窗口全接受推进了 cursor。这里剪掉没通过验证的尾部。 backend.model.att_state_manager.commit_bounded_att_state( - b_req_idx[b_req_mtp_start_loc], accept_lengths, role="main" + b_req_idx[b_req_mtp_start_loc], + accept_lengths, + role="main", + req_verify_width=_req_verify_width(b_req_mtp_start_loc, b_req_idx.shape[0]), ) return accept_lengths, accepted_index +def _req_verify_width(starts: torch.Tensor, num_verify_tokens: int) -> torch.Tensor: + """每个请求这次送去验证的 token 数，与 b_req_mtp_start_loc 逐行对齐。 + + num_verify_tokens 是验证 batch 的 token 数，用作最后一个请求的结束位置。 + TODO: 并进 MTP verify 的同一个算子，直接用 b_req_mtp_start_loc 裁剪 cursor，省掉这次单独的宽度计算。 + """ + + req_verify_width = torch.empty_like(starts) + if starts.numel() == 0: + return req_verify_width + if starts.numel() > 1: + req_verify_width[:-1].copy_(starts[1:] - starts[:-1]) + req_verify_width[-1] = num_verify_tokens - starts[-1] + return req_verify_width + + def scatter_mtp_next_tokens( backend: ModeBackend, proposal: SpecProposal, # proposal.token_ids: [req_num, draft_step] diff --git a/test.md b/test.md new file mode 100644 index 0000000000..f1a6298685 --- /dev/null +++ b/test.md @@ -0,0 +1,58 @@ +# ReplaySSM 折回：精度确认与两种时机的对比 + +测量对象是当前工作区的实现：普通 decode 放不下时，折回发生在每一层正在算线性注意力的 `replay` kernel 开头。容量够时只重放，不写 canonical SSM。Qwen3.5-0.8B，思考关闭，GSM8K 1000 题，`max_tokens=512`。 + +## 1. 1000 题上，现在这条 kernel 没有原先的精度崩溃 + +崩溃的样子是答案尾部刷出长串 `!` 后被 512 token 截断。同一套协议下，现在三条 replay 配置的 `!` 循环都是 0，strict 正确数贴着 native fp32。 + +| 配置 | strict | flexible | 截断 | `!` 循环 | 与 native 最终数字一致 | +|---|---:|---:|---:|---:|---:| +| native fp32 | 519 | 542 | 74 | 0 | 1000 | +| inline bf16，cache 4 | 511 | 534 | 81 | 0 | 807 | +| inline bf16，cache 16 | 515 | 536 | 63 | 0 | 810 | +| inline fp32，cache 16 | 518 | 533 | 67 | 0 | 803 | + +同一批题上，改层内折回之前： + +| 配置 | strict | 截断 | `!` 循环 | +|---|---:|---:|---:| +| inline bf16，cache 4 | 336 | 403 | 383 | +| inline bf16，cache 16 | 444 | 134 | 60 | +| inline fp32，cache 16 | 481 | 126 | 58 | + +和 native 对不上的大约两成，是回放和 fused kernel 的正常差别。cache 4 的正确数从 336 回到 511，已经贴着 native 的 519。这次没有重跑 precompute 和 KDA。 + +## 2. 整步结束后再折回，和层内折回的差别 + +「整步后再折回」指旧路径：每层 replay 先把 token 算完，最后一层 `maybe_compress_prefix` 再用独立 kernel `_fold_prefix` 折回，并在该 kernel 内把 cursor 写成 0。「层内折回」指现在的路径：放不下时，折回在每一层的 replay kernel 开头完成，共享 cursor 等全部层结束才在外面更新。 + +### 精度 + +独立折回如果等全部线程块写完再改 cursor，和层内折回是同一个结果。bf16 和 fp32、cache 4 和 cache 16，一步边界 token 的输出和 SSM 最大误差都是 0。 + +旧路径不是这样。`WHEN_FULL` 由其中一个线程块在 kernel 内部把 cursor 写成 0，其他线程块可能还没读到原来的长度，对应的层或头就跳过折回。和正确折回相比，状态最大误差： + +| 批量 | cache | 20 次里对不上 | 状态最大误差 | +|---|---:|---:|---:| +| 1 | 4 | 1 | 0.125 | +| 1 | 16 | 11 | 0.132 | +| 16 | 4 | 20 | 0.218 | +| 16 | 16 | 20 | 0.268 | + +这个误差会留在 SSM 里。cache 4 的 1000 题里有 530 题最终数字因此改变，其中 379 题改到与 native 一致，44 题离开 native。 + +### 性能 + +形状是 Qwen3.5-0.8B 的 18 个线性层、16 头、维度 128、bf16，CUDA graph。时间是 18 层加在一起。普通步是前缀已满但还没触发折回；折回步是刚好放不下的那一步。 + +| 批量 | cache | 普通回放 | 层内折回再回放 | 整步后的折回 kernel | 每步都有的空折回 | +|---|---:|---:|---:|---:|---:| +| 1 | 4 | 106 µs | 178 µs | 12 µs | 3.3 µs | +| 1 | 16 | 140 µs | 361 µs | 31 µs | 3.4 µs | +| 16 | 4 | 180 µs | 309 µs | 161 µs | 13 µs | +| 16 | 16 | 283 µs | 732 µs | 400 µs | 13 µs | + +单条请求时，折回串在每一层的 replay 里。cache 4 的那一步大约多 70 µs；独立 kernel 把 18 层一起折完只要 12 µs。摊到每个 token 大约多 10 µs。批量 16 时，独立折回本身要 160–400 µs，放进已经要启动的层 kernel 里，每个 token 大约少 10 µs。 + +旧路径在容量够的时候仍会每步启动一次空折回，批量 1 约 3 µs，批量 16 约 13 µs。现在这条路径不再发这个 kernel。相对完整 decode（还有 full attention 和 MLP），这几十微秒不是端到端的大头。1000 题墙钟时间现在是 55–61 秒，native 是 54 秒。 diff --git a/test/att_resources/validate_model_cache.py b/test/att_resources/validate_model_cache.py index 4bf1514704..569d4f99a5 100644 --- a/test/att_resources/validate_model_cache.py +++ b/test/att_resources/validate_model_cache.py @@ -236,10 +236,12 @@ def decode(request: torch.Tensor): 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) + req_verify_width = torch.full_like(cuda_req_indices, step_width, dtype=torch.int32) for item in (model, *draft_models): resources.commit_bounded_att_state( req_indices=cuda_req_indices, accepted_steps=accepted_steps, + req_verify_width=req_verify_width, role=item.att_state_role, ) return logits diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py new file mode 100644 index 0000000000..c60038af40 --- /dev/null +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -0,0 +1,176 @@ +"""Measure a full KDA state cycle: all layers, acceptance, folds and export. + +This synthetic kernel benchmark excludes convolutions, model projections and +HTTP scheduling. Use a matched service workload for end-to-end conclusions. +""" + +import argparse +import gc +import json +import statistics +from pathlib import Path + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda +from lightllm.common.basemodel.attention.linear.replay_kda_decode import get_configs, replay +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, + attach_replay_history as ReplaySSMExecutor, +) + + +def workload(args, batch, width, dtype, mode, acceptance): + layers, heads, dim = args.layers, args.heads, 128 + packed = torch.randn(batch * width, heads * dim * 3, dtype=torch.bfloat16, device="cuda") + q, k, v = [x.view(batch * width, heads, dim) for x in packed.split(heads * dim, -1)] + native_shape = (1, batch * width) if width > 1 else (batch, 1) + gate = torch.randn(batch * width, heads * dim, dtype=q.dtype, device="cuda") - 3 + beta = torch.randn(batch * width, heads, dtype=q.dtype, device="cuda") + alog = torch.zeros(heads, device="cuda") + bias = torch.zeros(heads * dim, device="cuda") + reqs = torch.arange(batch, dtype=torch.int32, device="cuda") + # 原生 decode 在宽度为 1 时不收前缀和。replay 始终要，一行一个 token 时就是 0, 1, 2, ...。 + replay_cu = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * width + cu = replay_cu if width > 1 else None + counts = reqs % width + 1 if acceptance == "partial" else torch.full_like(reqs, width) + accepted = torch.cat((counts - 1, counts.new_zeros(1))) + snapshot = torch.empty(layers, heads, dim, dim, dtype=dtype, device="cuda") + slots = (batch + 1) * (width if mode == "native" else 1) + state = torch.randn(layers, slots, heads, dim, dim, dtype=dtype, device="cuda") * 0.01 + ssm_backend = None + if mode != "native": + capacity = max(8, triton.next_power_of_2(width)) + ssm_backend = ReplaySSMExecutor(state, capacity, width, q.dtype, kda=True, compute_mode=mode) + indexes = reqs[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) + if width == 1: + indexes = reqs + bytes_used = ( + sum(x.numel() * x.element_size() for x in vars(ssm_backend).values() if isinstance(x, torch.Tensor)) + if ssm_backend + else state.numel() * state.element_size() + ) + + native_index = 0 if acceptance == "partial" else width - 1 + pinned_config = {"value": None} + + def cycle(): + for _ in range(args.rounds): + for layer in range(layers): + if ssm_backend is None: + fused_recurrent_kda( + q.view(*native_shape, heads, dim), + k.view(*native_shape, heads, dim), + v.view(*native_shape, heads, dim), + gate.view(*native_shape, -1), + beta.view(*native_shape, heads), + alog, + bias, + state[layer], + indexes, + cu_seqlens=cu, + num_accepted_tokens=counts if cu is not None else None, + ) + else: + replay( + state_ssm=ssm_backend.ssm_state[layer], + state_keys=ssm_backend.keys[layer], + state_deltas=ssm_backend.deltas[layer], + state_gates=ssm_backend.gates[layer], + state_raw_keys=ssm_backend.raw_keys[layer], + state_raw_values=ssm_backend.raw_values[layer], + state_betas=ssm_backend.betas[layer], + state_cursors=ssm_backend.cursors, + state_tile_sync_mark=ssm_backend.tile_sync_mark, + q=q, + k=k, + v=v, + a=gate, + b=beta, + a_log=alog, + bias=bias, + reqs=reqs, + cu_seqlens=replay_cu, + state_verify_width=ssm_backend.verify_width, + state_compute_mode=ssm_backend.compute_mode, + state_gate_lower_bound=ssm_backend.gate_lower_bound, + run_config=pinned_config["value"], + ) + if ssm_backend is not None: + accept_replay_window(ssm_backend, reqs, accepted, replay_cu) + if isinstance(ssm_backend, ReplaySSMExecutor): + ssm_backend._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) + ssm_backend.fold_accepted_prefix(reqs) + else: + snapshot.copy_(state[:, native_index] if ssm_backend is None else state[:, 0]) + + return ssm_backend, cycle, bytes_used, pinned_config + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--batch", type=int, nargs="+", default=[16, 64, 192]) + p.add_argument("--width", type=int, nargs="+", default=[1, 3]) + p.add_argument("--dtype", nargs="+", choices=["float32", "bfloat16"], default=["float32", "bfloat16"]) + p.add_argument("--layers", type=int, default=34) + p.add_argument("--heads", type=int, default=8) + p.add_argument("--rounds", type=int, default=8) + p.add_argument("--repeats", type=int, default=5) + p.add_argument("--tune", action="store_true") + p.add_argument("--output", type=Path, required=True) + args = p.parse_args() + torch.manual_seed(53) + rows = [] + report = dict( + settings={**vars(args), "output": str(args.output)}, + device=torch.cuda.get_device_name(), + torch=torch.__version__, + triton=triton.__version__, + results=rows, + ) + for dtype_name in args.dtype: + for batch in args.batch: + for width in args.width: + for acceptance in ["full", "partial"] if width > 1 else ["full"]: + for mode in ["native", "inline", "precompute"] if width > 1 else ["native", "inline"]: + ssm_backend, cycle, size, pinned_config = workload( + args, batch, width, getattr(torch, dtype_name), mode, acceptance + ) + candidates = [] + if ssm_backend is not None: + for config in get_configs() if args.tune else [{"BV": 32, "num_warps": 2, "num_stages": 1}]: + if mode == "precompute" and config["BV"] < 16: + continue + pinned_config["value"] = config + cycle() + ms = triton.testing.do_bench_cudagraph(cycle, rep=50) + candidates.append(dict(config=config, ms=ms)) + pinned_config["value"] = min(candidates, key=lambda r: r["ms"])["config"] + cycle() + samples = [ + triton.testing.do_bench_cudagraph(cycle, rep=100) / args.rounds for _ in range(args.repeats) + ] + row = dict( + batch=batch, + width=width, + dtype=dtype_name, + acceptance=acceptance, + mode=mode, + cache_len=getattr(ssm_backend, "capacity", None), + state_history_bytes=size, + ms_per_round=statistics.median(samples), + samples_ms=samples, + config=dict(pinned_config["value"]) if pinned_config["value"] else None, + candidates=candidates, + ) + rows.append(row) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(row), flush=True) + del ssm_backend, cycle + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/kernels/benchmark_replayssm_service.py b/test/benchmark/kernels/benchmark_replayssm_service.py new file mode 100644 index 0000000000..a1cee4826a --- /dev/null +++ b/test/benchmark/kernels/benchmark_replayssm_service.py @@ -0,0 +1,119 @@ +"""Compare streamed requests on baseline and optimized services with identical warmup. + +Reports request timings and token IDs for an exact cross-revision comparison. +""" + +import argparse +import concurrent.futures +import json +import statistics +import time +import threading +from pathlib import Path + +import requests +from transformers import AutoTokenizer + +p = argparse.ArgumentParser() +p.add_argument("--output", required=True) +p.add_argument("--batch", type=int, nargs="+", default=[1, 8, 16]) +p.add_argument("--repeats", type=int, default=3) +p.add_argument("--tokens", type=int, default=64) +p.add_argument("--warmups", type=int, default=1) +p.add_argument("--model-dir", required=True) +p.add_argument("--url", default="http://127.0.0.1:18785/generate_stream") +args = p.parse_args() +url = args.url +tok = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=True) +prefix = tok.encode( + "Describe how a computer executes a program, from source code to machine instructions. " + "Explain the role of memory, processors, and operating systems. " * 100 +)[:512] + + +def prepare_prompt(index, rep): + suffix = tok.encode( + f"Example number {rep * 1000 + index}. Discuss this scenario in detail. " + + ("An engineer measures performance and validates numerical correctness. " * 100) + )[:512] + return prefix + suffix + + +def stream_request(index, rep, ids, barrier, cohort_start): + payload = { + "inputs": ids, + "parameters": { + "do_sample": False, + "disable_prompt_cache": True, + "ignore_eos": True, + "max_new_tokens": args.tokens, + }, + } + barrier.wait() + start = time.perf_counter() + times = [] + tokens = [] + texts = [] + finished = False + with requests.post(url, json=payload, stream=True, timeout=300) as response: + response.raise_for_status() + # Consume each HTTP chunk; byte-sized iteration can saturate the client CPU. + for line in response.iter_lines(chunk_size=None): + if not line.startswith(b"data:"): + continue + item = json.loads(line[5:]) + assert "error" not in item, item + token = item.get("token", {}) + if token.get("id") is not None: + times.append(time.perf_counter() - cohort_start) + tokens.append(token["id"]) + texts.append(token.get("text", "")) + finished = item.get("finished", finished) + assert finished and len(tokens) == args.tokens, (finished, len(tokens)) + return dict( + index=index, + rep=rep, + prompt_len=len(ids), + token_ids=tokens, + text="".join(texts), + token_times_s=times, + ttft_ms=(times[0] - (start - cohort_start)) * 1000, + tpot_ms=(times[-1] - times[0]) * 1000 / max(1, len(times) - 1), + elapsed_s=time.perf_counter() - start, + ) + + +def run_concurrent_requests(batch, rep): + prompts = [prepare_prompt(i, rep) for i in range(batch)] + barrier = threading.Barrier(batch) + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=batch) as pool: + samples = list(pool.map(lambda i: stream_request(i, rep, prompts[i], barrier, start), range(batch))) + elapsed = time.perf_counter() - start + begin = max(r["token_times_s"][0] for r in samples) + 0.5 + end = min(r["token_times_s"][-1] for r in samples) - 0.5 + count = sum(sum(begin <= t <= end for t in r["token_times_s"]) for r in samples) + return dict( + batch=batch, + rep=rep, + elapsed_s=elapsed, + tokens_per_s=sum(len(r["token_ids"]) for r in samples) / elapsed, + steady_tokens_per_s=count / (end - begin) if end > begin else None, + steady_window_s=max(0.0, end - begin), + ttft_ms=statistics.median(r["ttft_ms"] for r in samples), + tpot_ms=statistics.median(r["tpot_ms"] for r in samples), + requests=samples, + ) + + +# Warm every measured graph bucket before timing, with the same requests on both revisions. +for batch in args.batch: + for warm in range(args.warmups): + run_concurrent_requests(batch, 99 + warm) +rows = [] +for batch in args.batch: + for rep in range(args.repeats): + row = run_concurrent_requests(batch, rep) + rows.append(row) + Path(args.output).write_text(json.dumps(rows, indent=2)) + print(json.dumps({k: v for k, v in row.items() if k != "requests"}), flush=True) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 10c52c8cea..86e62cb671 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -253,6 +253,8 @@ def build_hybrid_resources( small_page_capacity: int, draft_full_state: FullAttentionState | None = None, tp_rank: int = 0, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, ) -> AttentionStateManager: """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" from lightllm.common.attention_state import AttentionStateManager @@ -282,6 +284,9 @@ def build_hybrid_resources( request_capacity=request_capacity, mtp_step=mtp_step, device="cuda", + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=config.gate_lower_bound if glm else -5.0, ) bounded_state.initialize_page_pools( big_page_capacity=big_page_capacity, diff --git a/unit_tests/common/att_resources/test_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py index e3a09adb8a..457fdfa1c3 100644 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -298,6 +298,7 @@ def start_runtime(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([draft_count + 1], 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) @@ -308,6 +309,7 @@ def start_runtime(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([i + 2], 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 diff --git a/unit_tests/common/att_resources/test_bounded_state.py b/unit_tests/common/att_resources/test_bounded_state.py index ffd59db4aa..e2c08b1461 100644 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -254,7 +254,11 @@ def test_reset_requests_preserves_pages(cls: type[KDAState] | type[GDNState]): 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)) + state.commit_bounded_att_state( + torch.tensor([0]), + torch.tensor([1], dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + ) @pytest.mark.parametrize("cls", [KDAState, GDNState]) diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index 15d29bd85d..fb9ce7ba5b 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -9,6 +9,7 @@ from lightllm.server.core.objs.start_args_type import StartArgs +@pytest.mark.parametrize("mode", ["native", "replay"]) @pytest.mark.parametrize("draft_layers", [0, 1, 3]) @pytest.mark.parametrize( "layer_types,expected_indices,linear_layers", @@ -41,6 +42,7 @@ def test_hybrid_layer_mapping_and_mtp_states( expected_indices: tuple[int, ...], linear_layers: int, draft_layers: int, + mode: str, ) -> None: model_config = dict( num_hidden_layers=len(layer_types), @@ -56,7 +58,13 @@ def test_hybrid_layer_mapping_and_mtp_states( full_attention_interval=4, layer_types=layer_types, ) - args = StartArgs(data_type="bfloat16", linear_att_ssm_data_type="float32") + args = StartArgs( + data_type="bfloat16", + linear_att_ssm_data_type="float32", + mtp_step=2, + ssm_state_mode=mode, + ssm_replay_config={"cache_len": 8} if mode == "replay" else {}, + ) config = Qwen3NextResourceConfig.from_model_config(model_config, args) logical_layer_ids = tuple(range(len(layer_types))) (full_state,) = Qwen3NextTpPartModel.create_full_att_states( @@ -70,6 +78,11 @@ def test_hybrid_layer_mapping_and_mtp_states( model_logical_layer_ids=logical_layer_ids, device="meta", ) + if mode == "native": + stride = bounded_state.mtp_step + 1 + assert bounded_state.ssm_state.shape[1] == bounded_state.request_capacity * stride + if mode == "replay": + assert bounded_state.capacity == 8 manager = AttentionStateManager(request_capacity=1) manager.register_state(model_role="main", state=full_state) manager.register_state(model_role="main", state=bounded_state) diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py index af8dbfcde0..09e88b4a93 100644 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -162,7 +162,9 @@ def check_runtime_before_profile(device: str | torch.device | int): 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) + linear_state = main_linear_state(resources) + assert linear_state.mtp_state_index is not None + assert linear_state.mtp_enabled == 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() diff --git a/unit_tests/common/att_resources/test_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index a50ff4fc71..ea114d2ba9 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -36,6 +36,7 @@ def test_commit_selects_state_for_one_role(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], 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): @@ -46,6 +47,7 @@ def test_commit_selects_state_for_one_role(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([1], 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] @@ -60,17 +62,20 @@ def test_commit_without_recurrent_state_and_invalid_role(): manager.commit_bounded_att_state( torch.tensor([0], dtype=torch.int32, device="cuda"), torch.tensor([3], 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"), + 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"), + torch.tensor([3], dtype=torch.int32, device="cuda"), role="target", ) diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 8ec77c6f4d..344bfea57f 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -154,6 +154,7 @@ def auto_linear_backend_args(monkeypatch): lambda: SimpleNamespace( llm_prefill_att_backend=["fa3"], llm_decode_att_backend=["flashinfer"], + ssm_state_mode="native", ), ) @@ -169,7 +170,7 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear lambda name: validate_calls.append(name) or True, ) - backend_class = linear_create_utils.get_qwen35_linear_prefill_att_backend_class(index=1) + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) assert backend_class is FlashQlaLinearAttBackend backend = object.__new__(backend_class) @@ -177,6 +178,53 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear assert validate_calls == ["flashqla"] +def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace( + llm_prefill_att_backend=["fa3", "flashqla"], + llm_decode_att_backend=["flashinfer", "triton"], + ssm_state_mode="replay", + ), + ) + from lightllm.common.basemodel.attention.base_att import BaseAttBackend + from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend + from lightllm.common.basemodel.attention.linear.replay_gdn import ( + ReplayGDNFlashQlaAttBackend, + ReplayGDNTritonAttBackend, + ) + + prefill_backend = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) + decode_backend = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) + + assert prefill_backend is ReplayGDNFlashQlaAttBackend + assert decode_backend is ReplayGDNTritonAttBackend + assert prefill_backend.__bases__ == (BaseAttBackend,) + assert decode_backend.__bases__ == (BaseAttBackend,) + assert not issubclass(prefill_backend, LinearAttBackend) + assert not issubclass(decode_backend, LinearAttBackend) + + +def test_kda_entry_selects_native_or_replay(monkeypatch): + from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend + from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend + + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace(ssm_state_mode="native"), + ) + assert linear_create_utils.get_kda_linear_att_backend_class() is KDALinearAttBackend + + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace(ssm_state_mode="replay"), + ) + assert linear_create_utils.get_kda_linear_att_backend_class() is ReplayKDALinearAttBackend + + def test_flashqla_validation_loads_linear_cache_config(monkeypatch): linear_config = SimpleNamespace( tp_world_size=2, @@ -246,6 +294,7 @@ def test_explicit_linear_backend_arg_skips_auto_selection(monkeypatch): lambda: SimpleNamespace( llm_prefill_att_backend=["fa3", "triton"], llm_decode_att_backend=["flashinfer", "triton"], + ssm_state_mode="native", ), ) monkeypatch.setattr( @@ -254,8 +303,8 @@ def test_explicit_linear_backend_arg_skips_auto_selection(monkeypatch): lambda *args: pytest.fail("an explicit linear backend must not be auto-validated"), ) - prefill_backend_class = linear_create_utils.get_qwen35_linear_prefill_att_backend_class(index=1) - decode_backend_class = linear_create_utils.get_qwen35_linear_decode_att_backend_class(index=1) + prefill_backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) + decode_backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) assert prefill_backend_class is TritonLinearAttBackend assert decode_backend_class is TritonLinearAttBackend @@ -280,8 +329,8 @@ def test_cli_accepts_linear_backend_at_index_one(): @pytest.mark.parametrize( "get_backend_class", [ - linear_create_utils.get_qwen35_linear_prefill_att_backend_class, - linear_create_utils.get_qwen35_linear_decode_att_backend_class, + linear_create_utils.get_gdn_linear_prefill_att_backend_class, + linear_create_utils.get_gdn_linear_decode_att_backend_class, ], ) def test_linear_backend_rejects_non_linear_index(auto_linear_backend_args, get_backend_class): @@ -290,7 +339,7 @@ def test_linear_backend_rejects_non_linear_index(auto_linear_backend_args, get_b def test_missing_linear_decode_backend_arg_defaults_to_triton(auto_linear_backend_args): - backend_class = linear_create_utils.get_qwen35_linear_decode_att_backend_class(index=1) + backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) assert backend_class is TritonLinearAttBackend @@ -299,8 +348,8 @@ def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_ validate_calls = [] monkeypatch.setattr( linear_create_utils, - "linear_decode_att_backend_classes", - {"optimized": "linear.flashqla:FlashQlaLinearAttBackend", "triton": "linear.triton:TritonLinearAttBackend"}, + "gdn_decode_backends", + {"optimized": FlashQlaLinearAttBackend, "triton": TritonLinearAttBackend}, ) monkeypatch.setattr( linear_create_utils, @@ -308,7 +357,7 @@ def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_ lambda name: validate_calls.append(name) or True, ) - backend_class = linear_create_utils.get_qwen35_linear_decode_att_backend_class( + backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class( index=1, priority_list=("optimized", "triton") ) @@ -319,7 +368,7 @@ def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_ def test_gdn_prefill_backend_falls_back_when_flashqla_validation_fails(monkeypatch, auto_linear_backend_args): monkeypatch.setattr(linear_create_utils, "validate", lambda *args: False) - backend_class = linear_create_utils.get_qwen35_linear_prefill_att_backend_class(index=1) + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) assert backend_class is TritonLinearAttBackend @@ -328,8 +377,8 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ validate_calls = [] monkeypatch.setattr( linear_create_utils, - "linear_prefill_att_backend_classes", - {"flashqla2": "linear.triton:TritonLinearAttBackend", "flashqla3": "linear.flashqla:FlashQlaLinearAttBackend"}, + "gdn_prefill_backends", + {"flashqla2": TritonLinearAttBackend, "flashqla3": FlashQlaLinearAttBackend}, ) monkeypatch.setattr( linear_create_utils, @@ -337,7 +386,7 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ lambda name, *args: validate_calls.append(name) or name == "flashqla3", ) - backend_class = linear_create_utils.get_qwen35_linear_prefill_att_backend_class( + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class( index=1, priority_list=("flashqla2", "flashqla3") ) diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py index 1efed8423e..e7423d30f9 100644 --- a/unit_tests/common/basemodel/attention/test_lazy_import.py +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -67,6 +67,7 @@ def setUp(self): llm_decode_att_backend=["triton"], enable_ep_moe=False, mtp_step=0, + ssm_state_mode="native", ) self.validate = Mock(return_value=True) dependencies = { @@ -89,7 +90,6 @@ def setUp(self): self.attention = importlib.import_module(ATTENTION) def test_importing_selection_apis_loads_no_implementations(self): - importlib.import_module(f"{ATTENTION}.create_linear_utils") importlib.import_module(f"{ATTENTION}.linear") importlib.import_module(f"{ATTENTION}.nsa") self.assertEqual(self.imports.imported, []) @@ -184,24 +184,32 @@ def test_neo_auto_selection_uses_supported_fa3(self): self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.neo"]) self.validate.assert_called_once_with("fa3") - def test_linear_selection_imports_only_validated_backend(self): - backend = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") + def _allow_linear_backend_classes(self): + flashqla = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") + triton = self.imports.allow("linear.triton", "TritonLinearAttBackend") + self.imports.allow("linear.replay_gdn", "ReplayGDNTritonAttBackend") + replay_gdn = self.imports.allowed[f"{ATTENTION}.linear.replay_gdn"] + replay_gdn.ReplayGDNFlashQlaAttBackend = type("ReplayGDNFlashQlaAttBackend", (), {}) + self.imports.allow("linear.kda", "KDALinearAttBackend") + self.imports.allow("linear.replay_kda", "ReplayKDALinearAttBackend") + return flashqla, triton + + def test_linear_prefill_auto_selects_validated_flashqla(self): + flashqla, _triton = self._allow_linear_backend_classes() selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") - self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.linear.flashqla"]) + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), flashqla) self.validate.assert_called_once_with("flashqla") - def test_linear_selection_falls_back_without_importing_flashqla(self): - backend = self.imports.allow("linear.triton", "TritonLinearAttBackend") + def test_linear_selection_falls_back_to_triton(self): + _flashqla, triton = self._allow_linear_backend_classes() selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") self.validate.return_value = False - self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) - self.assertIs(selection.get_qwen35_linear_decode_att_backend_class(), backend) + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) + self.assertIs(selection.get_gdn_linear_decode_att_backend_class(), triton) self.args.llm_prefill_att_backend.append("triton") - self.assertIs(selection.get_qwen35_linear_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.linear.triton"]) + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) def test_vit_selection_loads_only_requested_implementation(self): package = "lightllm.common.basemodel.attention_vit" diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py new file mode 100644 index 0000000000..956bf5e173 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/replay_history.py @@ -0,0 +1,176 @@ +"""测试用的 replay history。正式路径走 ReplayGDNState / ReplayKDAState 的构造函数。""" + +import torch +import triton +import triton.language as tl + +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor + + +def attach_replay_history( + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + num_key_heads: int | None = None, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + *, + kda: bool = False, +): + """把连续前缀挂到调用方已有的 SSM 上，保留它原来的 stride。 + + kernel 测试需要非连续的层距，不能再走会重新分配状态的构造函数。 + """ + + # capacity 是 2 的幂，并且不小于一次验证窗口。 + assert ssm.dtype in (torch.float32, torch.bfloat16), "replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + state_cls = ReplayKDAState if kda else ReplayGDNState + owner = state_cls.__new__(state_cls) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner.mtp_enabled = owner.mtp_step > 0 + owner._algorithm = "kda" if kda else "gdn" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = torch.zeros(owner.request_capacity, dtype=torch.int32, device=ssm.device) + owner.big_pages_pool = None + owner.small_pages_pool = None + owner.ssm_state = ssm + owner.kda = kda + owner.gate_lower_bound = gate_lower_bound + owner.capacity = capacity + owner.verify_width = verify_width + owner.compute_mode = compute_mode + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + # 最后一个 slot 是 HOLD。cursor 是尚未折回的已接受长度。 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + owner.tile_sync_mark = torch.zeros( + (slots * (owner.mtp_step + 6), owner.num_key_heads), dtype=torch.int32, device=ssm.device + ) + # 重建用归一化 k：bf16 激活用 fp16；d 使用 fp32。折回仍用原始 k/v。KDA 的 g 按 key 通道存。 + reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, owner.num_key_heads, capacity, k) + gate_shape = (layers, slots, hv, capacity, k) if kda else (layers, slots, hv, capacity) + owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=ssm.device) + owner.gates = torch.empty(gate_shape, dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) + return owner + + +@triton.jit(do_not_specialize=["N"]) +def _trim_rejected_index( + cursors, + reqs, + cu, + accepted_index, + N, + HOLD: tl.constexpr, + BLOCK: tl.constexpr, +): + # accepted_index 按请求槽存放，接受个数是 index + 1。窗口长度与 reqs 逐行对齐。 + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + r = tl.load(reqs + i, i < N, HOLD) + valid = (i < N) & (r != HOLD) + window = tl.load(cu + i + 1, valid, 0) - tl.load(cu + i, valid, 0) + accepted = tl.load(accepted_index + r, valid, 0) + 1 + cursor = tl.load(cursors + r, valid, 0) + tl.store(cursors + r, cursor - (window - accepted), valid) + + +def replay_layer( + state: ReplayGDNState | ReplayKDAState, + layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor, + run_config: dict[str, int] | None = None, +): + """按测试 state 的类型调用对应的 replay decode。""" + + if state.kda: + from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay + + return replay( + state_ssm=state.ssm_state[layer_idx], + state_keys=state.keys[layer_idx], + state_deltas=state.deltas[layer_idx], + state_gates=state.gates[layer_idx], + state_raw_keys=state.raw_keys[layer_idx], + state_raw_values=state.raw_values[layer_idx], + state_betas=state.betas[layer_idx], + state_cursors=state.cursors, + state_tile_sync_mark=state.tile_sync_mark, + q=q, + k=k, + v=v, + a=a, + b=b, + a_log=a_log, + bias=bias, + reqs=reqs, + cu_seqlens=cu_seqlens, + state_verify_width=state.verify_width, + state_compute_mode=state.compute_mode, + state_gate_lower_bound=state.gate_lower_bound, + run_config=run_config, + ) + from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + + return replay( + state=state, + linear_layer_idx=layer_idx, + q=q, + k=k, + v=v, + a=a, + b=b, + a_log=a_log, + bias=bias, + reqs=reqs, + cu_seqlens=cu_seqlens, + run_config=run_config, + ) + + +def accept_replay_window(history, reqs, accepted_state_indices=None, cu_seqlens=None) -> None: + """测试里一次结算 cursor：先按整窗全接受，再按槽位 index 剪掉未接受尾部。 + + 没传 cu_seqlens 时每个请求的窗口长度是 1。正式路径先 advance_cursor，开了 MTP 再 commit_ssm_prefix。 + """ + + if reqs.numel() == 0: + return + if cu_seqlens is None: + cu_seqlens = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + advance_cursor(history, reqs, cu_seqlens) + if cu_seqlens is None or accepted_state_indices is None: + return + _trim_rejected_index[(triton.cdiv(reqs.numel(), 256),)]( + history.cursors, + reqs, + cu_seqlens, + accepted_state_indices, + reqs.numel(), + history.hold_req_id, + 256, + num_warps=4, + ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py new file mode 100644 index 0000000000..7fbb63df5d --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -0,0 +1,272 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, + attach_replay_history as ReplaySSMExecutor, + replay_layer, +) + + +def ssm_snapshot(backend, req): + """折进一份 SSM 副本再读出这个请求。不改 canonical SSM，也不动 cursor。""" + + saved = backend.ssm_state + clone = saved.clone() + backend.ssm_state = clone + try: + fold_prefix(backend, torch.tensor([req], dtype=torch.int32, device=saved.device)) + finally: + backend.ssm_state = saved + return clone[:, req] + + +from types import SimpleNamespace +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + + +def make(n, width): + hold = n + 3 + ids = list(reversed(range(n))) + [hold] + lengths = [1 + i % width for i in range(n)] + [1] + starts = [] + reqs = [] + mtps = [] + for r, l in zip(ids, lengths): + starts.append(len(reqs)) + reqs.extend([r] * l) + mtps.extend(range(l)) + tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) + args = [tensor(starts), tensor(reqs), tensor(mtps), tensor([1] * len(reqs)), width] + ssm_backend = ReplayGDNState.__new__(ReplayGDNState) + ssm_backend.cursors = tensor([7] * (hold + 1)) + ssm_backend.hold_req_id = hold + # 这个用例不触发折回，容量要盖住多步累加后的 cursor。 + ssm_backend.capacity = 4096 + ssm_backend.verify_width = width + manager = SimpleNamespace() + manager.mtp_state_index = tensor([-9] * (hold + 1)) + manager._ssm_executor = ssm_backend + return manager, args, (ids, lengths, starts) + + +def _replay_cu(args): + """args[0] 是每个请求的起点，最后一个请求接到本批 token 末尾。""" + + starts = args[0] + return torch.cat((starts, starts.new_full((1,), args[3].numel()))) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_request_counts_share_compiled_kernel(monkeypatch): + from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix + + reqs = torch.arange(512, dtype=torch.int32, device="cuda") + lengths = reqs % 4 + 1 + window = lengths + 4 + indices = torch.zeros(513, dtype=torch.int32, device="cuda") + cursors = torch.zeros_like(indices) + hashes = set() + run = _commit_prefix.run + + def record(*args, **kwargs): + kernel = run(*args, **kwargs) + hashes.add(kernel.hash) + return kernel + + monkeypatch.setattr(_commit_prefix, "run", record) + for count in (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511): + indices.fill_(-1) + cursors.zero_() + cursors[:count] = window[:count] + commit_ssm_prefix(indices, reqs[:count], lengths[:count], window[:count], cursors) + assert torch.equal(indices[:count], lengths[:count] - 1) + assert (indices[count:] == -1).all() + assert torch.equal(cursors[:count], lengths[:count]) + assert not cursors[count:].any() + assert len(hashes) == 1 + + +def accept_unfused(m, args): + update(m.mtp_state_index, *args) + reqs = args[1][args[0].long()] + if m._ssm_executor.verify_width > 1: + accept_replay_window(m._ssm_executor, reqs, m.mtp_state_index, _replay_cu(args)) + else: + accept_replay_window(m._ssm_executor, reqs) + + +def accept_fused(m, args): + reqs = args[1][args[0].long()] + ends = torch.cat((args[0][1:], args[0].new_full((1,), args[3].numel()))) + cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) + lengths = cumulative[ends.long()] - cumulative[args[0].long()] + executor = m._ssm_executor + if executor.verify_width > 1: + cu = _replay_cu(args) + advance_cursor(executor, reqs, cu) + commit_ssm_prefix(m.mtp_state_index, reqs, lengths, cu[1:] - cu[:-1], executor.cursors) + else: + # 宽度为 1 时正式路径不提交。对照只写下标，并把每个请求的 cursor 加一。 + m.mtp_state_index[reqs] = (lengths - 1).to(m.mtp_state_index.dtype) + advance_cursor(executor, reqs, torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_acceptance_fusion_matches_original_and_cpu_under_graph(): + checks = 0 + for width in [1, 2, 4, 8, 16]: + for n in [1, 7, 33, 128]: + m, args, (ids, lengths, starts) = make(n, width) + ref, _, _ = make(n, width) + # Compile both paths before graph capture, then reset before replay. + accept_unfused(ref, args) + accept_fused(m, args) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + accept_fused(m, args) + expected_index = [-9] * len(m.mtp_state_index) + expected_cursor = [7] * len(expected_index) + m.mtp_state_index.fill_(-9) + ref.mtp_state_index.fill_(-9) + m._ssm_executor.cursors.fill_(7) + ref._ssm_executor.cursors.fill_(7) + for step in range(24): + mask = [] + for i, (r, l) in enumerate(zip(ids, lengths)): + count = (step + i) % (l + 1) + mask.extend([1] * count + [0] * (l - count)) + expected_index[r] = count - 1 + if r != m._ssm_executor.hold_req_id: + expected_cursor[r] += count if width > 1 else 1 + args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) + accept_unfused(ref, args) + if step % 2: + graph.replay() + else: + accept_fused(m, args) + assert torch.equal(m.mtp_state_index, ref.mtp_state_index) + assert torch.equal(m._ssm_executor.cursors, ref._ssm_executor.cursors) + assert m.mtp_state_index.tolist() == expected_index + assert m._ssm_executor.cursors.tolist() == expected_cursor + checks += 1 + + print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, compute_mode): + torch.manual_seed(93) + initial = torch.randn(2, 8, 4, 32, 32, device="cuda").to(dtype) * 0.01 + ssm_backends = [ + ReplaySSMExecutor( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + compute_mode=compute_mode, + ) + for _ in range(2) + ] + managers = [] + for ssm_backend in ssm_backends: + manager = SimpleNamespace() + manager._ssm_executor = ssm_backend + manager.mtp_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") + managers.append(manager) + reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32, device="cuda") + token_reqs = torch.full((12,), 7, dtype=torch.int32, device="cuda") + mtp_indexes = torch.arange(12, dtype=torch.int32, device="cuda") % 3 + accepted_mask = torch.ones(12, dtype=torch.int32, device="cuda") + q = torch.randn(12, 2, 32, dtype=torch.bfloat16, device="cuda") + k = torch.randn_like(q) + v = torch.randn(12, 4, 32, dtype=torch.bfloat16, device="cuda") + a = torch.full((12, 4), -3.0, dtype=torch.bfloat16, device="cuda") + b = torch.randn_like(a) + alog = torch.zeros(4, device="cuda") + args = [cu[:-1], token_reqs, mtp_indexes, accepted_mask, 4] + + def cycle(i): + ssm_backend = ssm_backends[i] + outputs = [ + replay_layer( + ssm_backend, + layer, + q, + k, + v, + a, + b, + alog, + alog, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for layer in range(2) + ] + (accept_unfused if i == 0 else accept_fused)(managers[i], args) + return outputs + + graphs, outputs = [], [] + for i in range(2): + cycle(i) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = cycle(i) + graphs.append(graph) + outputs.append(output) + for iteration in range(40): + ids = [1, 3, 6] if iteration % 2 else [6, 1, 3] + lengths = [4, 4, 3] if iteration % 3 == 0 else [3, 3, 3] + lengths.append(12 - sum(lengths)) + ids.append(7) + starts, tokens, indexes, masks = [0], [], [], [] + for i, (req, length) in enumerate(zip(ids, lengths)): + count = (iteration + i) % (length + 1) if req != 7 else 0 + tokens.extend([req] * length) + indexes.extend(range(length)) + masks.extend([1] * count + [0] * (length - count)) + starts.append(len(tokens)) + for tensor, values in [ + (reqs, ids), + (cu, starts), + (token_reqs, tokens), + (mtp_indexes, indexes), + (accepted_mask, masks), + ]: + tensor.copy_(tensor.new_tensor(values)) + for tensor in (q, k, v, b): + tensor.normal_() + for graph in graphs: + graph.replay() + for left, right in zip(outputs[0], outputs[1]): + assert torch.equal(left, right) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + assert torch.equal(managers[0].mtp_state_index, managers[1].mtp_state_index) + if iteration % 7 == 0: + assert torch.equal(ssm_snapshot(ssm_backends[0], 3), ssm_snapshot(ssm_backends[1], 3)) + if iteration % 11 == 10: + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(reqs[:3]) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + if iteration == 19: + for manager in managers: + manager._ssm_executor.cursors[1].zero_() + manager._ssm_executor.ssm_state[:, 1].zero_() + manager.mtp_state_index[1] = 0 + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(reqs[:3]) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + assert torch.equal(ssm_backends[0].ssm_state[:, [0, 2, 4, 5, 7]], initial[:, [0, 2, 4, 5, 7]]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py new file mode 100644 index 0000000000..8252c445a1 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -0,0 +1,65 @@ +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import GDNState, ConvSSMStateConfig + +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("request_capacity,token_bucket", [(4, 3), (4, 16), (128, 512)]) +def test_dynamic_sequences_stay_within_request_capacity_under_graph(request_capacity, token_bucket): + hold = request_capacity + reqs = torch.full((token_bucket,), hold, dtype=torch.int32, device="cuda") + mtp_index = torch.zeros_like(reqs) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + manager = GDNState(config, request_capacity + 1, mtp_step=3) + manager.mtp_state_index.copy_(torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4) + sequences = min(request_capacity, token_bucket) + + def build(): + return build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, manager.mtp_state_index, hold) + + build() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + cu, ids, counts = build() + assert cu.numel() == token_bucket + 1 + assert ids.numel() == counts.numel() == token_bucket + + # Reuse one graph for full request pools, mixed verify lengths, and HOLD-only padding. + for lengths in ([1] * sequences, [min(4, token_bucket // sequences)] * sequences, [2, 1], []): + lengths = lengths[:request_capacity] + if sum(lengths) > token_bucket: + continue + request_ids = list(reversed(range(len(lengths)))) + host_reqs = [req for req, length in zip(request_ids, lengths) for _ in range(length)] + host_mtp = [j for length in lengths for j in range(length)] + tokens = len(host_reqs) + host_reqs += [hold] * (token_bucket - tokens) + host_mtp += [0] * (token_bucket - tokens) + reqs.copy_(torch.tensor(host_reqs, dtype=torch.int32, device="cuda")) + mtp_index.copy_(torch.tensor(host_mtp, dtype=torch.int32, device="cuda")) + graph.replay() + + expected_cu = [0] + for length in lengths: + expected_cu.append(expected_cu[-1] + length) + expected_cu += [tokens] * (token_bucket + 1 - len(expected_cu)) + expected_ids = request_ids + [hold] * (token_bucket - len(request_ids)) + assert cu.tolist() == expected_cu + assert ids.tolist() == expected_ids + expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (token_bucket - len(request_ids)) + assert counts.tolist() == expected_counts diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py new file mode 100644 index 0000000000..786376a362 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -0,0 +1,975 @@ +import pytest +import torch + +from lightllm.common.basemodel.attention.linear.replay_gdn_decode import get_configs +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, + attach_replay_history as ReplaySSMExecutor, + replay_layer, +) + + +def ssm_snapshot(backend, req): + """折进一份 SSM 副本再读出这个请求。不改 canonical SSM，也不动 cursor。""" + + saved = backend.ssm_state + clone = saved.clone() + backend.ssm_state = clone + try: + fold_prefix(backend, torch.tensor([req], dtype=torch.int32, device=saved.device)) + finally: + backend.ssm_state = saved + return clone[:, req] + + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def reference(q, k, v, a, b, a_log, bias, state, kda=False, gate_lower_bound=-5.0): + q, k, v = q.float(), k.float(), v.float() + q = q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt() * (q.shape[-1] ** -0.5) + k = k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt() + q = q.repeat_interleave(v.shape[-2] // q.shape[-2], -2) + k = k.repeat_interleave(v.shape[-2] // k.shape[-2], -2) + if kda: + g = gate_lower_bound * torch.sigmoid(a_log.exp()[:, None] * (a.float().view_as(k) + bias.view_as(k))) + state = state * g.exp()[..., None] + else: + g = -a_log.exp() * torch.nn.functional.softplus(a.float() + bias) + state = state * g.exp()[..., None, None] + d = b.float().sigmoid()[..., None] * (v - torch.einsum("...kv,...k->...v", state, k)) + state = state + k[..., None] * d[..., None, :] + return torch.einsum("...kv,...k->...v", state, q), state + + +@pytest.mark.parametrize("capacity", [4, 8]) +@pytest.mark.parametrize("bv", [8, 32], ids=["small_bv_fallback", "precomputed"]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("activation_dtype", [torch.float16, torch.bfloat16]) +def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state_dtype, kda, activation_dtype): + state = torch.zeros(1, 2, 1, 128, 128, dtype=state_dtype, device="cuda") + ssm_backend = ReplaySSMExecutor( + state, + capacity, + 2, + activation_dtype, + num_key_heads=1, + kda=kda, + compute_mode="precompute", + ) + reqs = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") + accepted = torch.tensor([1, 0], dtype=torch.int32, device="cuda") + q = torch.zeros(4, 1, 128, dtype=activation_dtype, device="cuda") + q[:2, 0, 0] = 1 + k = torch.zeros_like(q) + k[0, 0, 0] = -0.25 + k[1, 0, 0] = 1 + v = torch.zeros_like(q) + v[:2, 0, 0] = 60000 + a = torch.full((4, 128 if kda else 1), -20, dtype=activation_dtype, device="cuda") + b = torch.full((4, 1), 20, dtype=activation_dtype, device="cuda") + alog = torch.zeros(1, device="cuda") + bias = torch.zeros(128 if kda else 1, device="cuda") + expected_state = state[0, 0].float().clone() + for round_idx in range(2): + # The second round must consume the first round's accepted history. + # For capacity 4 the next span exactly fills it, so no fold should hide + # an overflowed history correction by reconstructing from raw inputs. + assert ssm_backend.cursors[0].item() == 2 * round_idx + assert not state.any() + actual = replay_layer( + ssm_backend, + 0, + q, + k, + v, + a, + b, + alog, + bias, + reqs, + cu, + run_config={"BV": bv, "num_warps": 1, "num_stages": 3}, + ) + expected = [] + for token in range(2): + output, expected_state = reference( + q[token], k[token], v[token], a[token], b[token], alog, bias, expected_state, kda=kda + ) + expected.append(output) + assert torch.isfinite(expected_state).all() + assert torch.isfinite(actual).all() + torch.testing.assert_close( + actual[:2].float(), + torch.stack(expected), + atol=32 if activation_dtype == torch.bfloat16 else 4, + rtol=4e-3 if activation_dtype == torch.bfloat16 else 1e-3, + ) + accept_replay_window(ssm_backend, reqs, accepted, cu) + # The correction is larger than FP16 can represent even though the + # recurrent state, inputs, and attention outputs are all finite. + assert ssm_backend.deltas[0, 0, 0, 1, 0].item() > torch.finfo(torch.float16).max + assert torch.isfinite(ssm_backend.deltas[0, 0, :, :4]).all() + assert ssm_backend.cursors[1].item() == 0 + torch.testing.assert_close(ssm_snapshot(ssm_backend, 0)[0].float(), expected_state, atol=128, rtol=0.01) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_flush_export_updates_checkpoint_and_clears_history(dtype): + state = torch.randn(1, 2, 1, 32, 64, device="cuda", dtype=dtype) * 0.01 + backend = ReplaySSMExecutor( + state, + capacity=4, + verify_width=1, + activation_dtype=torch.bfloat16, + ) + backend.raw_keys.normal_(0, 0.1) + backend.raw_values.normal_(0, 0.1) + backend.gates.fill_(-0.1) + backend.betas.fill_(0.5) + backend.cursors[0] = 2 + expected = ssm_snapshot(backend, 0).clone() + target = ssm_snapshot(backend, 0) + reqs = torch.tensor([0], dtype=torch.int32, device="cuda") + backend.fold_accepted_prefix(reqs) + torch.cuda.synchronize() + + assert backend.cursors[0].item() == 0 + torch.testing.assert_close(target, expected, atol=0, rtol=0) + torch.testing.assert_close(state[:, 0], expected, atol=0, rtol=0) + + empty_history_target = ssm_snapshot(backend, 0) + torch.cuda.synchronize() + torch.testing.assert_close(empty_history_target, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("capacity", [16, 32, 64]) +@pytest.mark.parametrize("width", [1, 3, 5, 16]) +@pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_replay_acceptance_flush_and_graph( + width, + dims, + capacity, + dtype, + run_config=None, + compute_mode="inline", + activation_dtype=torch.bfloat16, + gate_shift=-3, + head_geometry=(2, 4), + kda=False, + gate_lower_bound=-5.0, +): + torch.manual_seed(21) + kdim, vdim = dims + h, hv = head_geometry + slots = 5 + state = (torch.randn(2, slots, hv, kdim, vdim, device="cuda") * 0.01).to(dtype) + ssm_backend = ReplaySSMExecutor( + state, + capacity, + width, + num_key_heads=h, + kda=kda, + gate_lower_bound=gate_lower_bound, + compute_mode=compute_mode, + activation_dtype=activation_dtype, + ) + expected = state.float().clone() + pending = [0, 0] + state_atol = 2e-5 if dtype == torch.float32 else 0.002 + # Independent FP32 reductions can straddle a BF16 rounding boundary. + # Allow one relative BF16 spacing as well as accumulated near-zero error. + state_rtol = 2e-4 if dtype == torch.float32 else torch.finfo(dtype).eps + reqs = torch.tensor([2, 0, slots - 1], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, width, 2 * width - 1, 2 * width - 1], device="cuda", dtype=torch.int32) + if width == 1: + cu = torch.tensor([0, 1, 2, 2], device="cuda", dtype=torch.int32) + tokens = 3 * width + q = torch.randn(tokens, h, kdim, device="cuda", dtype=activation_dtype) + k = torch.randn_like(q) + v = torch.randn(tokens, hv, vdim, device="cuda", dtype=activation_dtype) + a = torch.randn(tokens, hv * kdim if kda else hv, device="cuda", dtype=activation_dtype) + gate_shift + b = torch.randn(tokens, hv, device="cuda", dtype=activation_dtype) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv * kdim if kda else hv, device="cuda") * 0.1 + accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") + + def step(): + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=run_config) + for layer in range(2) + ] + accept_replay_window(ssm_backend, reqs, accepted, cu) + return out + + # Compile on disposable state, then capture with the HOLD slot only. + saved = state.clone() + step() + state.copy_(saved) + ssm_backend.cursors.zero_() + reqs.fill_(slots - 1) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = step() + reqs.copy_(torch.tensor([2, 0, slots - 1], device="cuda", dtype=torch.int32)) + for iteration in range(3 * capacity + 5): + counts = [iteration % width + 1, min((iteration + 1) % width + 1, max(width - 1, 1))] + accepted[2] = counts[0] - 1 + accepted[0] = counts[1] - 1 + for seq, req in enumerate([2, 0]): + query_len = width if seq == 0 else max(width - 1, 1) + if pending[seq] + query_len > capacity: + ssm_backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device="cuda")) + expected[:, req] = expected[:, req].to(dtype).float() + pending[seq] = 0 + pending[seq] += counts[seq] + graph.replay() + for layer in range(2): + for seq, req in enumerate([2, 0]): + start = seq * width + length = width if seq == 0 else max(width - 1, 1) + cur = expected[layer, req].clone() + for j in range(length): + t = start + j + out, cur = reference( + q[t], + k[t], + v[t], + a[t], + b[t], + alog, + bias, + cur, + kda=kda, + gate_lower_bound=gate_lower_bound, + ) + torch.testing.assert_close(outputs[layer][t].float(), out, atol=0.006, rtol=0.02) + if j + 1 == counts[seq]: + expected[layer, req].copy_(cur) + if iteration == 2 * capacity + 1: + ssm_backend.fold_accepted_prefix(reqs[:2]) + expected = expected.to(dtype).float() + pending = [0, 0] + torch.testing.assert_close(state[:, :4].float(), expected[:, :4], atol=state_atol, rtol=state_rtol) + ssm_backend.fold_accepted_prefix(reqs[:2]) + torch.testing.assert_close(state.float(), expected.to(dtype).float(), atol=state_atol, rtol=state_rtol) + + +@pytest.mark.parametrize( + "state_dtype,activation_dtype", + [(torch.float32, torch.float32), (torch.float32, torch.bfloat16), (torch.bfloat16, torch.bfloat16)], +) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) +@pytest.mark.parametrize("width", [1, 4]) +@pytest.mark.parametrize("gate_shift", [-3, -10]) +def test_fold_matches_reference_across_precisions( + state_dtype, + activation_dtype, + compute_mode, + dims, + width, + gate_shift, +): + test_replay_acceptance_flush_and_graph( + width, + dims, + 8, + state_dtype, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + compute_mode=compute_mode, + activation_dtype=activation_dtype, + gate_shift=gate_shift, + ) + + +@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 3), (8, 4), (8, 5)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_short_replay_history_matches_independent_reference(capacity, width, dtype): + test_replay_acceptance_flush_and_graph(width, (128, 128), capacity, dtype) + + +@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 4), (8, 5), (16, 4), (16, 16), (32, 4)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_precomputed_projections_match_independent_reference(capacity, width, dtype): + test_replay_acceptance_flush_and_graph( + width, + (128, 128), + capacity, + dtype, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + compute_mode="precompute", + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_production_gdn_geometry_matches_independent_reference(dtype): + test_replay_acceptance_flush_and_graph( + 4, + (128, 128), + 8, + dtype, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + compute_mode="precompute", + head_geometry=(4, 12), + ) + + +@pytest.mark.parametrize("warps", [1, 4]) +@pytest.mark.parametrize("decay", [-3, -10]) +@pytest.mark.parametrize( + "state_dtype,dtype", + [(torch.float32, torch.float32), (torch.bfloat16, torch.bfloat16), (torch.float32, torch.bfloat16)], +) +@pytest.mark.parametrize("bv", [8, 32]) +@pytest.mark.parametrize("width", [1, 4]) +def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, dtype, width, bv): + torch.manual_seed(35) + state = (torch.randn((1, 3, 4, 128, 128), device="cuda") * 0.1).to(state_dtype) + config = {"BV": bv, "num_warps": warps, "num_stages": 3} + ssm_backends = [ + ReplaySSMExecutor( + state.clone(), + 16 if width == 1 else 8, + width, + dtype, + num_key_heads=2, + compute_mode=mode, + ) + for mode in ("inline", "precompute") + ] + reqs = torch.tensor([1, 0], dtype=torch.int32, device="cuda") + tokens = 2 if width == 1 else 7 + cu = torch.tensor([0, 1, 2] if width == 1 else [0, 4, 7], dtype=torch.int32, device="cuda") + q = torch.randn((tokens, 2, 128), device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn((tokens, 4, 128), device="cuda", dtype=dtype) + a = torch.randn((tokens, 4), device="cuda", dtype=dtype) + decay + b = torch.randn_like(a) + alog = torch.zeros(4, device="cuda") + accepted = torch.zeros(3, dtype=torch.int32, device="cuda") + for step in range(64): + for tensor in (q, k, v, b): + tensor.normal_() + accepted[0] = step % max(1, width - 1) + accepted[1] = step % width + outputs = [] + step_lengths = (cu[1:] - cu[:-1]).tolist() + if any( + ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + for req, length in zip(reqs.tolist(), step_lengths) + ): + overflow = [ + req + for req, length in zip(reqs.tolist(), step_lengths) + if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + ] + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(overflow_reqs) + for ssm_backend in ssm_backends: + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu, run_config=config)) + accept_replay_window(ssm_backend, reqs, accepted, cu) + if dtype == torch.float32: + torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) + else: + torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=0.01) + error = torch.linalg.vector_norm(outputs[0].float() - outputs[1].float()) + scale = torch.linalg.vector_norm(outputs[0].float()) + assert error < scale * (torch.finfo(dtype).eps / 2) + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(reqs) + torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) + + +def test_bf16_commit_rounds_before_the_next_token(): + # 前缀写满后先压回 SSM，下一 token 从长度 0 重写。两次 commit 必须得到同一块 BF16 状态。 + state = torch.full((1, 2, 1, 32, 64), 0.123, device="cuda", dtype=torch.bfloat16) + ssm_backends = [ReplaySSMExecutor(state.clone(), 16, 1, torch.float32) for _ in range(2)] + reqs = torch.tensor([0], device="cuda", dtype=torch.int32) + for ssm_backend in ssm_backends: + ssm_backend.raw_keys.zero_() + ssm_backend.raw_values.zero_() + ssm_backend.gates.fill_(-0.01) + ssm_backend.betas.zero_() + ssm_backend.cursors[0] = 16 + ssm_backend.fold_accepted_prefix(reqs) + q = torch.ones((1, 1, 32), device="cuda") + k = torch.zeros_like(q) + v = torch.zeros((1, 1, 64), device="cuda") + a = torch.full((1, 1), -3.0, device="cuda") + b = torch.zeros_like(a) + alog = torch.zeros(1, device="cuda") + bias = torch.zeros_like(alog) + outputs = [] + cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + assert ssm_backend.cursors[0].item() == 0 + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)) + torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) + torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) + + +def test_large_bf16_state_precompute_matches_inline(): + # 输出投影用 tf32x3。重建缓存的 k/d 在 bf16 激活下是 fp16，输出允许一个 fp16 单位。 + # 折回 SSM 仍从原始 k/v 在 fp32 中计算，两条路径的 checkpoint 必须一致。 + heads = 16 + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + reqs = torch.tensor([0], dtype=torch.int32, device="cuda") + alog = torch.zeros(heads, device="cuda") + bias = torch.zeros_like(alog) + + def check(capacity, width, steps): + torch.manual_seed(7) + state = (torch.randn((1, 2, heads, 128, 128), device="cuda") * 10).bfloat16() + backends = [ + ReplaySSMExecutor( + state.clone(), + capacity, + width, + torch.bfloat16, + num_key_heads=heads, + compute_mode=mode, + ) + for mode in ("inline", "precompute") + ] + cu = torch.tensor([0, width], dtype=torch.int32, device="cuda") + accepted = torch.tensor([width - 1], dtype=torch.int32, device="cuda") + for step in range(steps): + torch.manual_seed(100 + step) + q = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) + a = torch.randn((width, heads), device="cuda", dtype=torch.bfloat16) - 1 + b = torch.randn_like(a) + outputs = [ + replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu, run_config=config) for backend in backends + ] + # 重建用的 k/d 在 bf16 激活下是 fp16。两条路径在存入前的舍入可以差一个 fp16 单位。 + torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) + for backend in backends: + accept_replay_window(backend, reqs, accepted, cu) + overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] + if overflow.numel(): + backend.fold_accepted_prefix(overflow) + for backend in backends: + backend.fold_accepted_prefix(reqs) + torch.testing.assert_close(backends[0].ssm_state, backends[1].ssm_state, atol=0, rtol=0) + + check(4, 1, 16) + check(16, 2, 4) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capacity", [4, 8]) +def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): + # 放不下时，GDN 和 KDA 都在 replay kernel 内折回。cursor 等两层都算完才改。 + torch.manual_seed(11 + capacity + int(kda)) + hv, kh, dim = 2, 2, 32 + state = torch.randn(2, 2, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + backend = ReplaySSMExecutor(state, capacity, 1, dtype, num_key_heads=kh, kda=kda) + twin = ReplaySSMExecutor(state.clone(), capacity, 1, dtype, num_key_heads=kh, kda=kda) + assert backend.gates.shape[-1] == (dim if kda else capacity) + reqs = torch.tensor([0], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(hv * dim if kda else hv, device="cuda") + for step in range(capacity + 1): + q = torch.randn(1, kh, dim, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(1, hv, dim, device="cuda", dtype=dtype) + a = (torch.randn(1, hv * dim if kda else hv, device="cuda", dtype=dtype) - 1).contiguous() + b = torch.randn(1, hv, device="cuda", dtype=dtype) + if step == capacity: + fold_prefix(twin, reqs, config) + twin.cursors.index_fill_(0, reqs.long(), 0) + assert int(backend.cursors[0]) == capacity + outs, twin_outs = [], [] + for layer in range(2): + outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config)) + twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config)) + if step == capacity: + assert int(backend.cursors[0]) == capacity + for actual, expected in zip(outs, twin_outs): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + accept_replay_window(backend, reqs) + accept_replay_window(twin, reqs) + assert int(backend.cursors[0]) == 1 + assert int(twin.cursors[0]) == 1 + # bf16 两条路径都先舍入到同一状态。fp32 的独立 kernel 与层内 kernel 归约顺序可差几个 ulp。 + state_atol = 0 if dtype == torch.bfloat16 else 1e-6 + torch.testing.assert_close(backend.ssm_state, twin.ssm_state, atol=state_atol, rtol=0) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("bv", [32, 128]) +@pytest.mark.parametrize("width", [1, 3]) +def test_overflow_preserves_shared_keys_across_heads_and_graph(kda, compute_mode, bv, width): + torch.manual_seed(123) + batch, kh, hv, dim, capacity = 64, 4, 16, 128, 4 + dtype = torch.bfloat16 + initial = torch.randn(2, batch + 1, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 + config = {"BV": bv, "num_warps": 4 if bv == 128 else 1, "num_stages": 1} + backend, reference_backend = [ + ReplaySSMExecutor( + initial.clone(), + capacity, + width, + dtype, + num_key_heads=kh, + kda=kda, + compute_mode=compute_mode, + ) + for _ in range(2) + ] + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) + reqs[-2:] = batch + windows = (torch.arange(reqs.numel(), device="cuda", dtype=torch.int32) % width) + 1 + windows[0] = 0 + cu = torch.cat((windows.new_zeros(1), windows.cumsum(0).int())) + tokens = int(cu[-1]) + q = torch.randn(tokens, kh, dim, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(tokens, hv, dim, device="cuda", dtype=dtype) + a = torch.full((tokens, hv * dim if kda else hv), -2.0, device="cuda", dtype=dtype) + b = torch.randn(tokens, hv, device="cuda", dtype=dtype) + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(hv * dim if kda else hv, device="cuda") + + def step(owner): + outputs = [ + replay_layer(owner, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config) for layer in range(2) + ] + accept_replay_window(owner, reqs, cu_seqlens=cu) + return outputs + + step(backend) + backend.ssm_state.copy_(initial) + backend.cursors.zero_() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = step(backend) + backend.ssm_state.copy_(initial) + backend.cursors.zero_() + for _ in range(2 * capacity + 1): + for tensor in (q, k, v, b): + tensor.normal_() + lengths = reference_backend.cursors[reqs[:-2].long()] + overflow = reqs[:-2][(windows[:-2] > 0) & (lengths + windows[:-2] > capacity)] + if overflow.numel(): + fold_prefix(state=reference_backend, reqs=overflow, run_config=config) + reference_backend.cursors.index_fill_(0, overflow.long(), 0) + expected = step(reference_backend) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(backend.ssm_state, reference_backend.ssm_state, rtol=0, atol=0) + for left, right in zip(actual, expected): + torch.testing.assert_close(left, right, rtol=torch.finfo(dtype).eps, atol=2e-5) + assert torch.equal(backend.cursors, reference_backend.cursors) + torch.testing.assert_close(backend.ssm_state[:, 0], initial[:, 0], rtol=0, atol=0) + torch.testing.assert_close(backend.ssm_state[:, batch], initial[:, batch], rtol=0, atol=0) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): + state = torch.randn(1, 3, 2, 32, 32, device="cuda") + backend = ReplaySSMExecutor(state, 8, 4, torch.float32, num_key_heads=2, kda=kda) + backend.cursors[0] = 3 + reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, 0, 0], device="cuda", dtype=torch.int32) + q = torch.randn(1, 2, 32, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, 2, 32, device="cuda") + a = torch.zeros(1, 2 * 32 if kda else 2, device="cuda") + b = torch.zeros(1, 2, device="cuda") + bias = torch.zeros(2 * 32 if kda else 2, device="cuda") + replay_layer( + backend, + 0, + q, + k, + v, + a, + b, + bias, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + accepted = torch.tensor([-1, -1, -1], device="cuda", dtype=torch.int32) + accept_replay_window(backend, reqs, accepted, cu) + assert backend.cursors.tolist() == [3, 0, 0] + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_mode, capture_metadata): + torch.manual_seed(31) + state = (torch.randn(2, 4, 4, 128, 128, device="cuda") * 0.01).to(dtype) + ssm_backend = ReplaySSMExecutor( + state, + 8, + 4, + dtype, + num_key_heads=2, + kda=kda, + compute_mode=compute_mode, + ) + expected = state.float().clone() + pending = [0, 0] + reqs = torch.full((3,), 3, dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 4, 8, 8], dtype=torch.int32, device="cuda") + accepted = torch.zeros(4, dtype=torch.int32, device="cuda") + q = torch.randn(8, 2, 128, dtype=dtype, device="cuda") + k = torch.randn_like(q) + v = torch.randn(8, 4, 128, dtype=dtype, device="cuda") + a = torch.full((8, 4 * 128 if kda else 4), -3.0, dtype=dtype, device="cuda") + b = torch.randn(8, 4, dtype=dtype, device="cuda") + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + + def forward(): + return [ + replay_layer( + ssm_backend, + layer, + q, + k, + v, + a, + b, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + for layer in range(2) + ] + + def step(): + positions = ssm_backend.cursors.index_select(0, reqs.long()) + out = forward() + accept_replay_window(ssm_backend, reqs, accepted, cu) + return positions, out + + step() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + positions = torch.zeros_like(reqs) + with torch.cuda.graph(graph): + if capture_metadata: + positions, outputs = step() + else: + outputs = forward() + lengths = [(3, 4), (2, 1), (3, 2), (1, 4), (4, 3), (0, 2), (2, 0), (4, 4)] + for iteration in range(24): + ids = [0, 1] if iteration < 3 or iteration % 2 == 0 else [1, 0] + widths = lengths[iteration % len(lengths)] + counts = ( + list(widths) if iteration < 3 else [min(width, 1 + (iteration + i) % 4) for i, width in enumerate(widths)] + ) + reqs.copy_(reqs.new_tensor(ids + [3])) + cu.copy_(cu.new_tensor([0, widths[0], sum(widths), 8])) + accepted.fill_(-1) + for req, count in zip(ids, counts): + accepted[req] = count - 1 + for tensor in (q, k, v, b): + tensor.normal_() + before = state.clone() + for req, width in zip(ids, widths): + if pending[req] + width > 8: + ssm_backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device="cuda")) + expected[:, req] = expected[:, req].to(dtype).float() + pending[req] = 0 + old_cursors = [pending[req] for req in ids] + [0] + if not capture_metadata: + positions.copy_(ssm_backend.cursors.index_select(0, reqs.long())) + graph.replay() + if not capture_metadata: + accept_replay_window(ssm_backend, reqs, accepted, cu) + assert positions.tolist() == old_cursors + for seq, (req, width, count) in enumerate(zip(ids, widths, counts)): + for layer in range(2): + current = expected[layer, req].clone() + for j in range(width): + t = (0 if seq == 0 else widths[0]) + j + out, current = reference(q[t], k[t], v[t], a[t], b[t], alog, bias, current, kda=kda) + torch.testing.assert_close(outputs[layer][t].float(), out, atol=0.006, rtol=0.02) + if j + 1 == count: + expected[layer, req].copy_(current) + pending[req] += count + assert ssm_backend.cursors.tolist() == pending + [0, 0] + for output in outputs: + assert torch.count_nonzero(output[sum(widths) :]).item() == 0 + torch.testing.assert_close(state[:, 2:], before[:, 2:], atol=0, rtol=0) + if iteration == 2: + assert pending[0] == 8 + torch.testing.assert_close(state[:, 0], before[:, 0], atol=0, rtol=0) + ssm_backend.fold_accepted_prefix(reqs[:2]) + torch.testing.assert_close( + state.float(), + expected.to(dtype).float(), + atol=2e-5 if dtype == torch.float32 else 0.002, + rtol=2e-4 if dtype == torch.float32 else torch.finfo(dtype).eps, + ) + + +def test_snapshot_without_history_copies_active_state(): + state = torch.randn((2, 3, 2, 32, 64), device="cuda", dtype=torch.bfloat16) + ssm_backend = ReplaySSMExecutor(state, 4, 3) + snapshot = ssm_snapshot(ssm_backend, 1) + torch.testing.assert_close(snapshot, state[:, 1], rtol=0, atol=0) + assert ssm_backend.cursors[1].item() == 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("run_config", get_configs(), ids=lambda c: f"bv{c['BV']}-w{c['num_warps']}-s{c['num_stages']}") +def test_autotune_candidates_match_recurrence(dtype, run_config): + test_replay_acceptance_flush_and_graph(3, (32, 64), 8, dtype, run_config) + + +@pytest.mark.parametrize("save_big", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_args, save_big, dtype): + from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig + from unit_tests.common.test_linear_state import make_state + + torch.manual_seed(42) + state_rtol = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps + config = ConvSSMStateConfig( + logical_layer_ids=(0, 1), + global_k_heads=2, + global_v_heads=2, + key_dim=32, + value_dim=64, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=dtype, + ) + state, _, layout = make_state(mode="replay", config=config) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + state.mtp_state_index[1] = 1 + state.conv_state.normal_() + state.ssm_state.zero_() + state.raw_keys.normal_() + state.raw_values.normal_() + state.gates.uniform_(-0.1, -0.01) + state.betas.uniform_(0.1, 0.9) + active_before = state.ssm_state[:, 1].clone() + conv_expected = state.conv_state[:, 1, :, :3].clone() + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + kind = "big" if save_big else "small" + pool = state.big_pages_pool if save_big else state.small_pages_pool + for count in (2, 1): + state.cursors[1] = count + expected = active_before.float().clone() + for i in range(count): + key = state.raw_keys[:, 1, :, i].float() + key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() + expected *= state.gates[:, 1, :, i].exp()[..., None, None] + delta = state.betas[:, 1, :, i, None] * ( + state.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) + ) + expected += key[..., None] * delta[..., None, :] + materialized = ssm_snapshot(state, 1) + torch.testing.assert_close(materialized, expected.to(dtype), rtol=state_rtol, atol=2e-5) + state.save_bounded_att_runtime_to_pages([1], [8], kind, [0]) + state.write_to_pd_page(1, block, 0, 8, layout) + torch.cuda.synchronize() + torch.testing.assert_close(pool.ssm_state[0], active_before.cpu(), rtol=0, atol=0) + torch.testing.assert_close(pool.conv_state[0], conv_expected.cpu(), rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, 1], active_before, rtol=0, atol=0) + assert state.cursors[1].item() == count + for restore in ( + lambda: state.restore_bounded_att_runtime_from_pages([2], [8], kind, [0]), + lambda: state.read_from_pd_page(block, 2, 0, 8, layout), + ): + state.cursors[2] = 3 + state.mtp_state_index[2] = 2 + restore() + assert state.cursors[2].item() == state.mtp_state_index[2].item() == 0 + torch.testing.assert_close(state.ssm_state[:, 2], active_before, rtol=0, atol=0) + torch.testing.assert_close(state.conv_state[:, 2, :, :3], conv_expected, rtol=0, atol=0) + state.reset_bounded_att_state([2]) + assert not state.ssm_state[:, 2].any() + assert not state.conv_state[:, 2].any() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (8, 4), (16, 5), (32, 4), (64, 4)]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("gate_lower_bound,gate_shift", [(-5.0, 0), (-2.5, -6)]) +@pytest.mark.parametrize("activation_dtype", [torch.float32, torch.bfloat16]) +def test_kda_replay_reference_fold_and_graph( + dtype, capacity, width, compute_mode, gate_lower_bound, gate_shift, activation_dtype +): + test_replay_acceptance_flush_and_graph( + width, + (128, 128), + capacity, + dtype, + run_config={"BV": 32, "num_warps": 2, "num_stages": 1}, + compute_mode=compute_mode, + gate_shift=gate_shift, + activation_dtype=activation_dtype, + head_geometry=(2, 2), + kda=True, + gate_lower_bound=gate_lower_bound, + ) + + +@pytest.mark.parametrize("config", get_configs()) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_kda_tuning_candidates_match_reference(config, compute_mode): + test_replay_acceptance_flush_and_graph( + 3, + (128, 128), + 4, + torch.bfloat16, + run_config=config, + compute_mode=compute_mode, + head_geometry=(2, 2), + kda=True, + gate_shift=-3, + ) + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize( + "bucket,live,folding,width,separate", + [ + (12, 12, 12, 4, True), + (13, 13, 13, 4, False), + (255, 255, 127, 4, True), + (256, 256, 128, 4, False), + (257, 257, 128, 4, False), + (256, 256, 128, 3, True), + (256, 12, 12, 4, True), + (256, 13, 4, 4, True), + ], +) +def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, separate, kda): + torch.manual_seed(72) + initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 + backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda) + reqs = torch.cat( + ( + torch.arange(live, device="cuda", dtype=torch.int32), + torch.full((bucket - live,), live, device="cuda", dtype=torch.int32), + ) + ) + indices = torch.zeros(live + 1, device="cuda", dtype=torch.int32) + expected = initial[0, :live].clone() + alog = torch.zeros(1, device="cuda") + bias = torch.zeros(32 if kda else 1, device="cuda") + for step, query_width in enumerate((4, 4, width)): + if step == 2: + backend.cursors.index_fill_(0, reqs[folding:live].long(), 0) + expected[folding:] = initial[0, folding:live] + cu = (torch.arange(bucket + 1, device="cuda", dtype=torch.int32) * query_width).clamp_max(live * query_width) + q = torch.randn(live, query_width, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((live, query_width, 32 if kda else 1), -2.0, device="cuda") + b = torch.randn(live, query_width, 1, device="cuda") + if (backend.cursors[:live] + query_width > backend.capacity).any(): + backend.fold_accepted_prefix(reqs[:live]) + actual = replay_layer( + backend, + 0, + q.reshape(-1, 1, 32), + k.reshape(-1, 1, 32), + v.reshape(-1, 1, 32), + a.flatten(0, 1), + b.flatten(0, 1), + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ).reshape(live, query_width, 1, 32) + for token in range(query_width): + output, expected = reference( + q[:, token], + k[:, token], + v[:, token], + a[:, token], + b[:, token], + alog, + bias.expand(live, 1, 32) if kda else bias, + expected, + kda=kda, + ) + torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) + lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) + advance_cursor(backend, reqs[:live], cu[: live + 1]) + commit_ssm_prefix(indices, reqs[:live], lengths, lengths, backend.cursors) + backend.fold_accepted_prefix(reqs[:live]) + torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) + assert torch.equal(backend.ssm_state[0, live], initial[0, live]) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): + # 最后一层按整窗全接受推进 cursor。放不下时 cursor 变成窗口长度，而不是旧 cursor 加窗口。 + torch.manual_seed(19 + int(kda)) + hv, kh, dim = 2, 2, 32 + dtype = torch.bfloat16 + state = torch.randn(1, 2, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 + backend = ReplaySSMExecutor( + state, + 4, + 2, + dtype, + num_key_heads=kh, + kda=kda, + compute_mode=compute_mode, + ) + reqs = torch.tensor([0], device="cuda", dtype=torch.int32) + indices = torch.zeros(2, device="cuda", dtype=torch.int32) + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(hv * dim if kda else hv, device="cuda") + cu = torch.tensor([0, 2], device="cuda", dtype=torch.int32) + for step in range(3): + q = torch.randn(2, kh, dim, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(2, hv, dim, device="cuda", dtype=dtype) + a = torch.full((2, hv * dim if kda else hv), -2.0, device="cuda", dtype=dtype) + b = torch.randn(2, hv, device="cuda", dtype=dtype) + replay_layer( + backend, + 0, + q, + k, + v, + a, + b, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + lengths = torch.tensor([2], device="cuda", dtype=torch.int32) + advance_cursor(backend, reqs, cu) + commit_ssm_prefix(indices, reqs, lengths, lengths, backend.cursors) + if step < 2: + assert int(backend.cursors[0]) == 2 * (step + 1) + else: + assert int(backend.cursors[0]) == 2 + assert torch.isfinite(backend.ssm_state).all() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py new file mode 100644 index 0000000000..02797b61b7 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -0,0 +1,530 @@ +import pytest +import torch + +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, + attach_replay_history as ReplaySSMExecutor, + replay_layer, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from types import SimpleNamespace + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import _advance_cursor + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay + + kernels = {"_decode_replay": _decode_replay, "_advance_cursor": _advance_cursor} + hashes = {name: set() for name in kernels} + launches = [] + for name, kernel in kernels.items(): + + def record(*args, run=kernel.run, name=name, **kwargs): + compiled = run(*args, **kwargs) + hashes[name].add(compiled.hash) + launches.append(name) + return compiled + + monkeypatch.setattr(kernel, "run", record) + + torch.manual_seed(137) + initial = torch.randn(1, 513, 1, 32, 32, device="cuda") * 0.01 + eager, captured = [ReplaySSMExecutor(initial.clone(), 8, 1, torch.float32, kda=kda) for _ in range(2)] + reqs = torch.arange(512, device="cuda", dtype=torch.int32) + q = torch.randn(512, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((512, 32 if kda else 1), -3.0, device="cuda") + beta = torch.randn(512, 1, device="cuda") + alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") + + def cycle(backend, ids): + count = ids.numel() + cu = torch.arange(count + 1, dtype=torch.int32, device=ids.device) + out = replay_layer( + backend, + 0, + q[:count], + k[:count], + v[:count], + a[:count], + beta[:count], + alog, + bias, + ids, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + accept_replay_window(backend, ids, cu_seqlens=cu) + return out + + graphs = {} + for bucket in (32, 512): + ids = reqs[:bucket].clone() + cycle(captured, ids) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = cycle(captured, ids) + graphs[bucket] = graph, ids, out + captured.ssm_state.copy_(initial) + captured.cursors.zero_() + counts = (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511) + for count in counts: + graph, ids, output = graphs[32 if count <= 32 else 512] + ids.fill_(captured.hold_req_id) + ids[:count].copy_(reqs[:count]) + for _ in range(8): + live = reqs[:count] + if (eager.cursors[live.long()] + 1 > eager.capacity).any(): + eager.fold_accepted_prefix(live) + captured.fold_accepted_prefix(ids[:count]) + expected = cycle(eager, reqs[:count]) + calls = len(launches) + graph.replay() + assert len(launches) == calls + torch.testing.assert_close(output[:count], expected, atol=3e-5, rtol=3e-5) + assert torch.equal(captured.cursors, eager.cursors) + torch.testing.assert_close(captured.ssm_state, eager.ssm_state, atol=3e-5, rtol=3e-5) + assert len(hashes["_decode_replay"]) == 1 + assert len(hashes["_advance_cursor"]) == 1 + assert captured.cursors[-1] == 0 + assert torch.equal(captured.ssm_state[:, -1], initial[:, -1]) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +@pytest.mark.parametrize("varlen", [False, True]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varlen, kda): + torch.manual_seed(827) + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + initial = torch.randn(3, 5, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + ssm_backends = [ + cls( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + kda=kda, + compute_mode="precompute", + ) + for cls in [ReplaySSMExecutor, ReplaySSMExecutor] + ] + manager = SimpleNamespace() + manager._ssm_executor = ssm_backends[1] + manager.mtp_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) + rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) + mtp = torch.zeros_like(rows) + flags = torch.zeros_like(rows) + starts = torch.zeros(3, device="cuda", dtype=torch.int32) + accepted = torch.full((5,), -1, device="cuda", dtype=torch.int32) + q = torch.randn(12, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(12, 4, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((12, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + + def inputs(iteration): + ids = [0, 1, 3] if iteration % 2 == 0 else [3, 0, 1] + lengths = [1 + (iteration + j) % 4 for j in range(3)] if varlen else [4] * 3 + counts = [1 + (iteration // 2 + j) % length for j, length in enumerate(lengths)] + offsets, req_rows, indexes, mask = [], [], [], [] + table = [-1] * 5 + for req, length, count in zip(ids, lengths, counts): + offsets.append(len(req_rows)) + req_rows.extend([req] * length) + indexes.extend(range(length)) + mask.extend([int(j < count) for j in range(length)]) + table[req] = count - 1 + logical = len(req_rows) + req_rows.extend([4] * (12 - logical)) + indexes.extend([0] * (12 - logical)) + mask.extend([0] * (12 - logical)) + rows.copy_(rows.new_tensor(req_rows)) + mtp.copy_(mtp.new_tensor(indexes)) + flags.copy_(flags.new_tensor(mask)) + starts.copy_(starts.new_tensor(offsets)) + accepted.copy_(accepted.new_tensor(table)) + return ids, lengths, counts, logical + + def metadata(ssm_backend): + cu, reqs, _counts = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.mtp_state_index, 4) + return cu, reqs, ssm_backend.cursors.index_select(0, reqs.long()) + + def forward_accept(index, cu, reqs): + ssm_backend = ssm_backends[index] + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu, run_config=config) + for layer in range(3) + ] + if index == 0: + accept_replay_window(ssm_backend, reqs, accepted, cu) + else: + count = starts.numel() + commit_reqs = rows[starts.long()] + window = (cu[1 : count + 1] - cu[:count]).contiguous() + advance_cursor(ssm_backend, commit_reqs, cu[: count + 1]) + commit_ssm_prefix( + manager.mtp_state_index, + commit_reqs, + accepted[commit_reqs.long()] + 1, + window, + ssm_backend.cursors, + ) + return out + + inputs(0) + graphs, captured, outputs = [], [], [] + for index, ssm_backend in enumerate(ssm_backends): + cu, reqs, pos = metadata(ssm_backend) + forward_accept(index, cu, reqs) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + manager.mtp_state_index.zero_() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + if capture_metadata: + cu, reqs, pos = metadata(ssm_backend) + out = forward_accept(index, cu, reqs) + graphs.append(g) + captured.append((cu, reqs, pos)) + outputs.append(out) + for ssm_backend in ssm_backends: + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + manager.mtp_state_index.zero_() + folds = 0 + for iteration in range(32): + ids, lengths, counts, logical = inputs(iteration) + for tensor in [q, k, v, beta]: + tensor.normal_() + overflow = [ + req for req, length in zip(ids, lengths) if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + ] + if overflow: + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(overflow_reqs) + folds += len(overflow) + for index, ssm_backend in enumerate(ssm_backends): + if not capture_metadata: + for dst, src in zip(captured[index], metadata(ssm_backend)): + dst.copy_(src) + graphs[index].replay() + torch.cuda.synchronize() + for expected, actual in zip(captured[0], captured[1]): + assert torch.equal(actual, expected) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for expected, actual in zip(outputs[0], outputs[1]): + torch.testing.assert_close( + actual[:, :logical].float(), expected[:, :logical].float(), atol=0.006, rtol=0.02 + ) + torch.testing.assert_close( + ssm_backends[1].ssm_state.float(), + ssm_backends[0].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + table = manager.mtp_state_index.cpu().tolist() + for req, count in zip(ids, counts): + assert table[req] == count - 1 + assert table[2] == table[4] == 0 + assert ssm_backends[1].cursors[4].item() == 0 + assert folds > 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +@pytest.mark.parametrize("stagger", [False, True]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, stagger, kda): + torch.manual_seed(731) + initial = torch.randn(3, 3, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + ssm_backends = [ + ReplaySSMExecutor( + initial.clone(), + 8, + 1, + dtype, + num_key_heads=2, + kda=kda, + ) + for _ in range(2) + ] + for ssm_backend in ssm_backends: + ssm_backend.fold_programs = 3 + reqs = torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32) + q = torch.randn(3, 2, 128, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(3, 4, 128, device="cuda", dtype=dtype) + a = torch.full((3, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=dtype) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + + cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + + def forward(ssm_backend): + return [ + replay_layer( + ssm_backend, + layer, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for layer in range(3) + ] + + def warm(ssm_backend): + out = forward(ssm_backend) + accept_replay_window(ssm_backend, reqs) + return out + + graphs, positions, outputs = [], [], [] + for ssm_backend in ssm_backends: + warm(ssm_backend) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + if stagger: + reqs.copy_(reqs.new_tensor([0, 2, 2])) + for _ in range(3): + warm(ssm_backend) + reqs.copy_(reqs.new_tensor([0, 1, 2])) + pos = torch.zeros_like(reqs) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + if capture_metadata: + pos = ssm_backend.cursors.index_select(0, reqs.long()) + out = forward(ssm_backend) + if capture_metadata: + accept_replay_window(ssm_backend, reqs) + graphs.append(graph) + positions.append(pos) + outputs.append(out) + for iteration in range(24): + reqs.copy_(reqs.new_tensor([0, 1, 2] if iteration % 2 == 0 else [1, 0, 2])) + for tensor in [q, k, v, beta]: + tensor.normal_() + overflow = [int(req) for req in reqs.tolist() if ssm_backends[0].cursors[req] + 1 > ssm_backends[0].capacity] + if overflow: + overflow_reqs = torch.tensor(sorted(set(overflow)), dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(overflow_reqs) + old = ssm_backends[0].cursors[reqs.long()].clone() + for ssm_backend, graph, pos in zip(ssm_backends, graphs, positions): + if not capture_metadata: + pos.copy_(ssm_backend.cursors.index_select(0, reqs.long())) + graph.replay() + if not capture_metadata: + accept_replay_window(ssm_backend, reqs) + torch.cuda.synchronize() + assert torch.equal(positions[0], old) + assert torch.equal(positions[0], positions[1]) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for expected, actual in zip(outputs[0], outputs[1]): + torch.testing.assert_close(actual.float(), expected.float(), atol=0.006, rtol=0.02) + torch.testing.assert_close( + ssm_backends[1].ssm_state.float(), + ssm_backends[0].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_cross_layer_fold_with_state_offset_above_int32(kda): + if torch.cuda.mem_get_info()[0] < 6 * 1024 ** 3: + pytest.skip("large-stride regression requires 6 GiB of free device memory") + torch.manual_seed(791) + shape = (48, 3, 12, 128, 128) + strides = (50_528_256, 196608, 16384, 128, 1) + assert strides[0] < 2 ** 31 < 47 * strides[0] + dense = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.01 + sparse = torch.empty_strided(shape, strides, device="cuda", dtype=dense.dtype) + sparse.copy_(dense) + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + ssm_backends = [ + ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=4, kda=kda) for state in [dense, sparse] + ] + q = torch.randn(8, 4, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(8, 12, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((8, 12 * 128 if kda else 12), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(8, 12, device="cuda", dtype=torch.bfloat16) + alog = torch.zeros(12, device="cuda") + bias = torch.zeros(12 * 128 if kda else 12, device="cuda") + reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) + cu = reqs.new_tensor([0, 4, 8]) + + def step(ssm_backend, accepted): + overflow = [int(req) for req in reqs.tolist() if ssm_backend.cursors[req] + 4 > ssm_backend.capacity] + if overflow: + ssm_backend.fold_accepted_prefix(reqs.new_tensor(overflow)) + pos = ssm_backend.cursors.index_select(0, reqs.long()) + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu, run_config=config) + for layer in range(48) + ] + accept_replay_window(ssm_backend, reqs, accepted, cu) + return pos, out + + for count in [2, 3]: + accepted = reqs.new_tensor([count - 1, -1, -1]) + for ssm_backend in ssm_backends: + step(ssm_backend, accepted) + reqs.copy_(reqs.new_tensor([0, 1])) + expected_pos, expected_out = step(ssm_backends[0], reqs.new_tensor([0, 0, -1])) + actual_pos, actual_out = step(ssm_backends[1], reqs.new_tensor([0, 0, -1])) + torch.cuda.synchronize() + assert torch.equal(actual_pos, expected_pos) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for layer in range(48): + torch.testing.assert_close(actual_out[layer].float(), expected_out[layer].float(), atol=0.006, rtol=0.02) + torch.testing.assert_close( + sparse[layer].float(), dense[layer].float(), atol=0.002, rtol=torch.finfo(torch.bfloat16).eps + ) + + +@pytest.mark.parametrize("batch", [1, 8, 12]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): + torch.manual_seed(1087 + batch) + initial = torch.randn(2, 14, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + ssm_backends = [ + ReplaySSMExecutor(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute") for _ in range(2) + ] + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) + padded = torch.cat((reqs, reqs.new_full((13 - batch,), 13))) + cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * 4 + padded_cu = torch.cat((cu, cu.new_full((13 - batch,), batch * 4))) + q = torch.randn(batch * 4, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(batch * 4, 4, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((batch * 4, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + accepted = torch.zeros(14, device="cuda", dtype=torch.int32) + graphs, outputs = [], [] + for ssm_backend, ids, offsets in zip(ssm_backends, [reqs, padded], [cu, padded_cu]): + + def step(): + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets, run_config=config) + for layer in range(2) + ] + accept_replay_window(ssm_backend, ids, accepted, offsets) + return out + + step() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = step() + graphs.append(graph) + outputs.append(out) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + for iteration in range(32): + accepted[:batch] = (torch.arange(batch, device="cuda") + iteration) % 4 + for value in [q, k, v, beta]: + value.normal_() + overflow = [idx for idx in range(batch) if ssm_backends[0].cursors[idx] + 4 > ssm_backends[0].capacity] + if overflow: + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(overflow_reqs) + for graph in graphs: + graph.replay() + for small, padded_out in zip(outputs[0], outputs[1]): + torch.testing.assert_close(small.float(), padded_out.float(), atol=0.006, rtol=0.02) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + torch.testing.assert_close( + ssm_backends[0].ssm_state.float(), + ssm_backends[1].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_full_prefix_does_not_address_past_capacity(kda): + torch.manual_seed(4) + capacity = 4 + backend = ReplaySSMExecutor( + torch.randn(1, 3, 2, 32, 32, device="cuda") * 0.01, + capacity, + 1, + torch.bfloat16, + kda=kda, + ) + reqs = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + q = torch.randn(2, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((2, 32 if kda else 1), -3.0, device="cuda") + beta = torch.randn(2, 1, device="cuda") + alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") + backend.ssm_state.zero_() + backend.cursors.zero_() + for name in ("keys", "deltas", "gates", "raw_keys", "raw_values", "betas"): + getattr(backend, name).zero_() + backend.cursors[reqs.long()] = capacity + 2 + replay_layer( + backend, + 0, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + backend.fold_accepted_prefix(reqs) + torch.cuda.synchronize() + assert torch.equal(backend.cursors[reqs.long()], torch.zeros(2, device="cuda", dtype=torch.int32)) + assert torch.isfinite(backend.ssm_state).all() + for _ in range(12): + replay_layer( + backend, + 0, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + accept_replay_window(backend, reqs) + overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] + if overflow.numel(): + backend.fold_accepted_prefix(overflow) + torch.cuda.synchronize() + assert int(backend.cursors[reqs.long()].max()) <= capacity + assert torch.isfinite(backend.ssm_state).all() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py new file mode 100644 index 0000000000..a94411fbd5 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -0,0 +1,70 @@ +import pytest +import torch + +from lightllm.common.basemodel.attention.linear import replay_gdn_decode as tuning +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + attach_replay_history as ReplaySSMExecutor, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("width", [1, 5]) +def test_rebuild_uses_one_layer_scratch(width): + state = torch.zeros((4, 3, 2, 32, 32), dtype=torch.bfloat16, device="cuda") + ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) + tokens = 3 * width + q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") + gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") + a_log = torch.zeros(2, device="cuda") + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width + reqs = torch.arange(3, device="cuda", dtype=torch.int32) + _, kwargs = tuning.rebuild_inputs(ssm_backend, 0, q, q, q, gates, gates, a_log, a_log, reqs, cu) + scratch = kwargs["state"] + assert scratch is not ssm_backend + assert scratch.ssm_state.shape[0] == 1 + assert scratch.capacity == ssm_backend.capacity + assert torch.count_nonzero(ssm_backend.ssm_state) == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("width", [1, 5]) +def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, width): + from lightllm.common.triton_utils import autotuner as autotuner_module + from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + + state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") + ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) + for tensor in vars(ssm_backend).values(): + if isinstance(tensor, torch.Tensor): + tensor.zero_() + tokens = 3 * width + q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") + gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") + a_log = torch.zeros(2, device="cuda") + bias = torch.zeros(2, device="cuda") + cu = torch.arange(4, device="cuda", dtype=torch.int32) * width + reqs = torch.arange(3, device="cuda", dtype=torch.int32) + before = {name: x.clone() for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} + + def oom(*args, **kwargs): + raise torch.cuda.OutOfMemoryError("scratch") + + tuner = tuning.replay + monkeypatch.setattr(tuner, "rebuild_input_func", oom) + monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", {}) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: 1) + monkeypatch.setattr( + tuner, "_bench", lambda *args, **kwargs: pytest.fail("Failed rebuild must not reach benchmarks") + ) + with pytest.raises(torch.cuda.OutOfMemoryError, match="scratch"): + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + tuning.replay(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu) + assert not list(tmp_path.iterdir()) + config = {"BV": 32, "num_warps": 1, "num_stages": 3} + for name, expected in before.items(): + torch.testing.assert_close(getattr(ssm_backend, name), expected, atol=0, rtol=0) + + tuning.replay(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) + torch.cuda.synchronize() diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py new file mode 100644 index 0000000000..ba349372ca --- /dev/null +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -0,0 +1,145 @@ +import multiprocessing +import traceback +from multiprocessing.reduction import ForkingPickler + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import GDNState, KDAState +from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor +from unit_tests.common.test_linear_state import make_state + + +class _SerializationFailure: + def __reduce_ex__(self, protocol): + raise RuntimeError("injected serialization failure") + + +def _assert_received_manager_roundtrip(payload, state_name, mode, expected): + resources = ForkingPickler.loads(payload) + state = getattr(resources, state_name) + assert state.big_pages_pool is state.small_pages_pool is None + assert resources.full_att_allocator is None + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state_name] + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(1, block, 0, 8, layout) + state.read_from_pd_page(block, 2, 0, 8, layout) + width = state.config.conv_width + ssm_index = 2 if mode == "replay" else 2 * (state.mtp_step + 1) + actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, ssm_index]) + for tensor, values in zip(actual, expected): + torch.testing.assert_close(tensor, torch.tensor(values, dtype=tensor.dtype, device="cuda"), rtol=0, atol=0) + assert state.mtp_state_index[2] == 0 + assert state.mtp_state_index[1] == 1 + if mode == "replay": + assert state.cursors[2] == 0 + assert state.cursors[1] == 1 + torch.cuda.synchronize() + + +def _receive_pd_managers(connection): + torch.cuda.set_device(0) + while True: + message = connection.recv() + if message is None: + return + try: + _assert_received_manager_roundtrip(*message) + connection.send(None) + except Exception: + connection.send(traceback.format_exc()) + + +@pytest.fixture(scope="module") +def pd_receiver(): + context = multiprocessing.get_context("spawn") + sender, receiver = context.Pipe() + process = context.Process(target=_receive_pd_managers, args=(receiver,)) + process.start() + receiver.close() + try: + yield sender, process + finally: + if process.is_alive(): + sender.send(None) + process.join(timeout=30) + if process.is_alive(): + process.terminate() + process.join(timeout=5) + sender.close() + assert process.exitcode == 0 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): + state, resources, _ = make_state(cls=cls, mode=mode, device="meta") + state.big_pages_pool, state.small_pages_pool = object(), object() + before = vars(state).copy() + exported = resources.get_pd_operator().create_ipc_copy() + copy = getattr(exported, state.state_name) + assert copy.big_pages_pool is copy.small_pages_pool is None + assert copy.conv_state is state.conv_state + assert copy.ssm_state is state.ssm_state + if mode == "replay": + assert copy.cursors is state.cursors + assert copy.keys is state.keys + assert copy.mtp_state_index is state.mtp_state_index + assert all(vars(state)[name] is value for name, value in before.items()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC and pinned memory required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("fail_serialization", [False, True]) +def test_shared_attention_state_roundtrip_preserves_worker_resources( + runtime_start_args, monkeypatch, pd_receiver, cls, mode, fail_serialization +): + state, resources, _ = make_state(cls=cls, mode=mode) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + state.conv_state.copy_(torch.arange(5, dtype=torch.bfloat16, device="cuda")) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + state.mtp_state_index[1], state.mtp_state_index[2] = 1, 2 + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors[1], state.cursors[2] = 1, 3 + width = state.config.conv_width + if mode == "replay": + expected = [ + state.conv_state[:, 1, :, :width].cpu().tolist(), + state.ssm_state[:, 1].cpu().tolist(), + ] + else: + expected = [ + state.conv_state[:, 1, :, :width].cpu().tolist(), + state.ssm_state[:, state.mtp_step + 1].cpu().tolist(), + ] + before = vars(state).copy() + source_slot = 1 if mode == "replay" else state.mtp_step + 1 + source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, source_slot].clone() + snapshots = [] + for pool in (state.big_pages_pool, state.small_pages_pool): + for tensor in (pool.conv_state, pool.ssm_state): + tensor.fill_(3) + snapshots.append((tensor, tensor.data_ptr(), tensor.clone())) + monkeypatch.setitem(ForkingPickler._extra_reducers, torch.Tensor, reduce_tensor) + exported = resources.get_pd_operator().create_ipc_copy() + if fail_serialization: + with pytest.raises(RuntimeError, match="injected serialization failure"): + ForkingPickler.dumps((_SerializationFailure(), exported)) + else: + sender, receiver = pd_receiver + sender.send((ForkingPickler.dumps(exported).tobytes(), state.state_name, mode, expected)) + assert sender.poll(30), f"PD receiver did not respond; exit code: {receiver.exitcode}" + error = sender.recv() + assert error is None, error + assert all(vars(state)[name] is value for name, value in before.items()) + torch.testing.assert_close(state.conv_state[:, 1], source_conv, rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, source_slot], source_ssm, rtol=0, atol=0) + for tensor, pointer, contents in snapshots: + assert tensor.data_ptr() == pointer + torch.testing.assert_close(tensor, contents, rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py new file mode 100644 index 0000000000..e538002a17 --- /dev/null +++ b/unit_tests/common/test_linear_state.py @@ -0,0 +1,334 @@ +from dataclasses import replace +from importlib import import_module + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState + + +def make_state(cls=GDNState, mode="native", step=2, dtype=torch.float32, device="cuda", config=None): + config = config or ConvSSMStateConfig( + logical_layer_ids=(0, 1), + global_k_heads=1, + global_v_heads=1 if cls is KDAState else 2, + key_dim=32, + value_dim=32, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=dtype, + ) + if mode == "replay": + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + + cls = ReplayKDAState if cls is KDAState else ReplayGDNState + state = cls(config, 4, mtp_step=step, device=device, ssm_replay_config={"cache_len": 4}) + else: + state = cls(config, 4, mtp_step=step, device=device) + resources = AttentionStateManager(request_capacity=4) + resources.register_state("main", state) + resources.finish_state_registration() + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state.state_name] + return state, resources, layout + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode,step", [("native", 0), ("native", 2), ("replay", 0), ("replay", 2)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( + runtime_start_args, cls, mode, step, dtype +): + state, _, layout = make_state(cls, mode, step, dtype) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + src, dst, offset = 1, 2, min(1, step) + state.conv_state.copy_(torch.arange(3 + step, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + if state.mtp_enabled: + state.mtp_state_index[src] = offset + if mode == "native": + expected_conv = state.conv_state[:, src, ..., :3].clone() + expected_ssm = state.ssm_state[:, src * (state.mtp_step + 1)].clone() + else: + # prefill 页只保存基准窗口和 canonical SSM，不按 MTP 偏移取窗口，也不折回 history。 + expected_conv = state.conv_state[:, src, ..., :3].clone() + expected_ssm = state.ssm_state[:, src].clone() + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors[src] = 1 + before = state.ssm_state.clone() + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(src, block, 0, 8, layout) + for kind in ("small", "big"): + state.save_bounded_att_runtime_to_pages([src], [8], kind, [0]) + torch.cuda.synchronize() + for pool in (state.small_pages_pool, state.big_pages_pool): + torch.testing.assert_close(pool.conv_state[0], expected_conv.cpu(), rtol=0, atol=0) + torch.testing.assert_close(pool.ssm_state[0], expected_ssm.cpu(), rtol=2e-5, atol=2e-5) + assert torch.equal(state.ssm_state, before) + if mode == "replay": + assert state.cursors[src].item() == 1 + + for restore in ( + lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), + lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "big", [0]), + lambda: state.read_from_pd_page(block, dst, 0, 8, layout), + lambda: state.copy_runtime_state_from_dp( + src_req_indices=[src], dst_req_indices=[dst], remote_dp_bounded_state=state + ), + ): + if state.mtp_enabled: + state.mtp_state_index[dst] = step + if mode == "replay": + state.cursors[dst] = 3 + restore() + actual_conv = state.conv_state[:, dst, ..., :3] + if mode == "replay": + actual_ssm = state.ssm_state[:, dst] + else: + actual_ssm = state.ssm_state[:, dst * (state.mtp_step + 1)] + torch.testing.assert_close(actual_conv, expected_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_ssm, expected_ssm, rtol=2e-5, atol=2e-5) + if state.mtp_enabled: + assert state.mtp_state_index[dst].item() == 0 + if mode == "replay": + assert state.cursors[dst].item() == 0 + assert state.cursors[src].item() == 1 + state.reset_bounded_att_state([dst]) + assert not state.conv_state[:, dst].any() + if mode == "replay": + assert not state.ssm_state[:, dst].any() + else: + start = dst * (state.mtp_step + 1) + assert not state.ssm_state[:, start : start + state.mtp_step + 1].any() + + +def copy_on_cpu(*, copies): + for source, destination in copies: + destination.copy_(source) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("destination", ["cpu", "cuda"]) +def test_checkpoint_batch_export(runtime_start_args, cls, mode, dtype, destination): + state, _, _ = make_state(cls, mode, 2, dtype) + state.initialize_page_pools(3, device=destination) + state.conv_state.copy_(torch.arange(5, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + state.mtp_state_index.copy_(torch.tensor([0, 2, 1, 0], device="cuda")) + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) + requests = [2, 0, 1] + width = state.config.conv_width + if mode == "replay": + expected = [(state.conv_state[:, req, :, :width].clone(), state.ssm_state[:, req].clone()) for req in requests] + else: + expected = [ + ( + state.conv_state[:, req, :, :width].clone(), + state.ssm_state[:, req * (state.mtp_step + 1)].clone(), + ) + for req in requests + ] + before = state.ssm_state.clone() + history = getattr(state, "runtime", state) + cursors = history.cursors.clone() if mode == "replay" else None + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + state.save_bounded_att_runtime_to_pages(requests, [8] * 3, "big", [1, 2, 0]) + stream.synchronize() + pool = state.big_pages_pool + for slot, (conv, ssm) in zip([1, 2, 0], expected): + torch.testing.assert_close(pool.conv_state[slot].cpu(), conv.cpu(), atol=0, rtol=0) + torch.testing.assert_close(pool.ssm_state[slot].cpu(), ssm.cpu(), atol=0, rtol=0) + assert torch.equal(state.ssm_state, before) + if cursors is not None: + assert torch.equal(state.cursors, cursors) + target, _, _ = make_state(cls, mode, 2, dtype) + if target.mtp_enabled: + target.mtp_state_index.fill_(2) + target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) + for req, (conv, ssm) in enumerate(expected): + actual_conv = target.conv_state[:, req, :, : target.config.conv_width] + if mode == "replay": + actual_ssm = target.ssm_state[:, req] + else: + actual_ssm = target.ssm_state[:, req * (target.mtp_step + 1)] + torch.testing.assert_close(actual_conv, conv, atol=0, rtol=0) + torch.testing.assert_close(actual_ssm, ssm, atol=0, rtol=0) + assert torch.equal(state.ssm_state, before) + if cursors is not None: + assert torch.equal(state.cursors, cursors) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_native_checkpoint_scratch_is_bounded(runtime_start_args): + config = ConvSSMStateConfig( + logical_layer_ids=(0, 1, 2, 3), + global_k_heads=1, + global_v_heads=2, + key_dim=128, + value_dim=128, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + state = GDNState(config, 65, mtp_step=2, device="cuda") + state.initialize_page_pools(64) + requests = list(range(64)) + state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) + torch.cuda.synchronize() + assert torch.cuda.max_memory_allocated() - before < 1024 * 1024 + + +def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(monkeypatch): + # Exercise layout and accepted-state selection independently of the CUDA copy kernel. + monkeypatch.setattr( + import_module("lightllm.common.attention_state.bounded_att.gdn.gdn_state"), + "copy_bounded_att_tensors", + copy_on_cpu, + ) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=2, + key_dim=2, + value_dim=3, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + tp_world_size=4, + ) + states, layouts, snapshots = [], [], [] + src, dst = 1, 2 + for rank in range(4): + state, _, layout = make_state(device="cpu", config=replace(config, tp_rank=rank)) + state.conv_state.fill_(-1) + state.ssm_state.fill_(-2) + state.mtp_state_index.fill_(1) + state.conv_state[:, src, :2, :3].fill_(10 + rank) + state.conv_state[:, src, 2:4, :3].fill_(20 + rank) + state.conv_state[:, src, 4:, :3].fill_(30 + rank) + state.ssm_state[:, src * (state.mtp_step + 1)].fill_(40 + rank) + states.append(state) + layouts.append(layout) + snapshots.append((state.conv_state.clone(), state.ssm_state.clone())) + size = layouts[0].byte_end + block = torch.full((1, size + 16), 0xA5, dtype=torch.uint8) + for state, layout in zip(states, layouts): + state.write_to_pd_page(src, block, 0, 8, layout) + expected_conv = torch.cat( + [torch.full((1, channels, 3), value) for channels, value in ((2, 10), (2, 20), (3, 30), (3, 32))], + dim=1, + ).float() + expected_ssm = torch.stack((torch.full((1, 2, 3), 40), torch.full((1, 2, 3), 42)), dim=1).float() + conv_bytes = expected_conv.numel() * 4 + torch.testing.assert_close(block[0, :conv_bytes].view(torch.float32).view_as(expected_conv), expected_conv) + torch.testing.assert_close(block[0, conv_bytes:size].view(torch.float32).view_as(expected_ssm), expected_ssm) + assert torch.all(block[0, size:] == 0xA5) + exported = block.clone() + other_requests = torch.tensor([0, 1, 3]) + for rank, (state, layout, (conv_before, ssm_before)) in enumerate(zip(states, layouts, snapshots)): + torch.testing.assert_close(state.conv_state, conv_before) + torch.testing.assert_close(state.ssm_state, ssm_before) + state.read_from_pd_page(block, dst, 0, 8, layout) + expected_local_conv = torch.cat( + (expected_conv[:, :4], expected_conv[:, 4 + rank // 2 * 3 : 7 + rank // 2 * 3]), dim=1 + ) + torch.testing.assert_close(state.conv_state[:, dst, ..., :3], expected_local_conv) + torch.testing.assert_close(state.ssm_state[:, dst * 3], expected_ssm[:, rank // 2 : rank // 2 + 1]) + assert state.mtp_state_index.tolist() == [1, 1, 0, 1] + torch.testing.assert_close(state.conv_state[:, other_requests], conv_before[:, other_requests]) + other_slots = (other_requests[:, None] * 3 + torch.arange(3)).flatten() + torch.testing.assert_close(state.ssm_state[:, other_slots], ssm_before[:, other_slots]) + assert torch.equal(block, exported) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("step", [0, 2]) +@pytest.mark.parametrize("restore_kind", ["checkpoint", "pd", "dp"]) +def test_replay_restored_slot_decodes_like_clean_slot(runtime_start_args, cls, step, restore_kind): + from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import replay_layer + + torch.manual_seed(1690) + state, _, layout = make_state(cls=cls, mode="replay", step=step) + state.initialize_page_pools(big_page_capacity=1) + state.ssm_state.normal_(mean=0, std=0.02) + state.conv_state.normal_(mean=0, std=0.1) + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.save_bounded_att_runtime_to_pages(req_indices=[0], state_token_lens=[8], page_kind="big", page_slot_ids=[0]) + state.write_to_pd_page(req_idx=0, pd_page_buffers=block, pd_page_index=0, state_token_len=8, layout=layout) + hv, kd, vd = state.ssm_state.shape[-3:] + kh = state.num_key_heads + width = step + 1 + q = torch.randn(width, kh, kd, device="cuda", dtype=torch.bfloat16).repeat(2, 1, 1) + k = torch.randn_like(q[:width]).repeat(2, 1, 1) + v = torch.randn(width, hv, vd, device="cuda", dtype=torch.bfloat16).repeat(2, 1, 1) + a = torch.full((2 * width, hv * kd if state.kda else hv), -3.0, device="cuda") + b = torch.zeros(2 * width, hv, device="cuda") + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(a.shape[-1], device="cuda") + reqs = torch.tensor([1, 2], dtype=torch.int32, device="cuda") + cu = torch.arange(3, dtype=torch.int32, device="cuda") * width + + def decode(): + outputs = [] + for layer in state.config.logical_layer_ids: + outputs.append( + replay_layer( + state=state, + layer_idx=layer, + q=q, + k=k, + v=v, + a=a, + b=b, + a_log=alog, + bias=bias, + reqs=reqs, + cu_seqlens=cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + ) + state.finish_decode_token(logical_layer_id=layer, reqs=reqs, cu_seqlens=cu) + return outputs + + decode() + assert state.cursors[1].item() == width + state.reset_bounded_att_state(req_indices=[2]) + state.restore_bounded_att_runtime_from_pages( + req_indices=[2], state_token_lens=[8], page_kind="big", page_slot_ids=[0] + ) + if restore_kind == "checkpoint": + state.restore_bounded_att_runtime_from_pages( + req_indices=[1], state_token_lens=[8], page_kind="big", page_slot_ids=[0] + ) + elif restore_kind == "pd": + state.read_from_pd_page(pd_page_buffers=block, req_idx=1, pd_page_index=0, state_token_len=8, layout=layout) + else: + state.copy_runtime_state_from_dp(src_req_indices=[0], dst_req_indices=[1], remote_dp_bounded_state=state) + assert state.cursors[1].item() == 0 + for _ in range(3): + for output in decode(): + torch.testing.assert_close(output[:width], output[width:], rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, 1], state.ssm_state[:, 2], rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py new file mode 100644 index 0000000000..fbfc4a43a9 --- /dev/null +++ b/unit_tests/common/test_linear_state_contract.py @@ -0,0 +1,241 @@ +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import KDAState +from unit_tests.common.test_linear_state import make_state +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState +from lightllm.models.glm5_next.config import Glm5NextResourceConfig +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.fixture +def linear_config(): + return ConvSSMStateConfig( + logical_layer_ids=(0, 1, 3, 4, 6, 7), + global_k_heads=1, + global_v_heads=2, + key_dim=32, + value_dim=32, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + + +@pytest.mark.parametrize("layer,row", [(0, 0), (1, 1), (3, 2), (4, 3), (6, 4), (7, 5)]) +def test_layer_buffers_follow_packed_linear_layers(linear_config, layer, row): + owner = GDNState(linear_config, request_capacity=1, device="cpu") + for i in range(6): + owner.conv_state[i].fill_(i) + owner.ssm_state[i].fill_(i + 10) + actual_conv, actual_ssm = owner.get_layer_view(layer) + assert (actual_conv == row).all() + assert (actual_ssm == row + 10).all() + + +@pytest.mark.parametrize("layer", [-1, 2, 5, 8]) +def test_linear_mapping_rejects_full_attention_and_out_of_range_layers(linear_config, layer): + owner = GDNState(linear_config, request_capacity=1, device="cpu") + with pytest.raises(KeyError): + owner.get_layer_view(layer) + + +def test_model_configuration_retains_nondefault_kda_gate_bound(): + 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], + "gate_lower_bound": -2.5, + }, + }, + StartArgs(tp=4, data_type="bfloat16"), + ) + assert config.gate_lower_bound == -2.5 + + +def test_kda_backend_receives_gate_bound_at_construction(linear_config): + owner = ReplayKDAState(linear_config, 3, mtp_step=2, gate_lower_bound=-2.5, device="cpu") + assert owner.gate_lower_bound == -2.5 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mtp_step,expected", [(0, "inline"), (2, "precompute")]) +def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): + replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState + owner = replay_cls(linear_config, request_capacity=2, mtp_step=mtp_step, device="cpu") + history = getattr(owner, "runtime", owner) + assert history.compute_mode == expected + assert history.verify_width == mtp_step + 1 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_mtp0_decode_keeps_explicit_precompute(linear_config, cls): + replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState + owner = replay_cls( + linear_config, + request_capacity=2, + mtp_step=0, + device="cpu", + ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, + ) + assert owner.verify_width == 1 + assert owner.compute_mode == "precompute" + assert owner.capacity == 8 + + +class _CaptureKernel: + def __init__(self): + self.args = None + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.args = args + self.kwargs = kwargs + + return launch + + +@pytest.mark.parametrize( + "module_name,cls,bv,precompute", + [ + ("replay_gdn_decode", ReplayGDNState, 32, True), + ("replay_gdn_decode", ReplayGDNState, 8, False), + ("replay_kda_decode", ReplayKDAState, 32, True), + ("replay_kda_decode", ReplayKDAState, 8, False), + ], +) +def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, module_name, cls, bv, precompute): + import importlib + + decode = importlib.import_module(f"lightllm.common.basemodel.attention.linear.{module_name}") + captured = _CaptureKernel() + monkeypatch.setattr(decode, "_decode_replay", captured) + owner = cls( + linear_config, + request_capacity=2, + mtp_step=0, + device="cpu", + ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, + ) + heads = owner.num_key_heads + q = torch.zeros(1, heads, linear_config.key_dim) + v = torch.zeros(1, owner.ssm_state.shape[-3], linear_config.value_dim) + gate = v.shape[-3] * linear_config.key_dim if cls is ReplayKDAState else v.shape[-3] + run_config = {"BV": bv, "num_warps": 1, "num_stages": 1} + if cls is ReplayKDAState: + decode.replay( + state_ssm=owner.ssm_state[0], + state_keys=owner.keys[0], + state_deltas=owner.deltas[0], + state_gates=owner.gates[0], + state_raw_keys=owner.raw_keys[0], + state_raw_values=owner.raw_values[0], + state_betas=owner.betas[0], + state_cursors=owner.cursors, + state_tile_sync_mark=owner.tile_sync_mark, + q=q, + k=q, + v=v, + a=torch.zeros(1, gate), + b=torch.zeros(1, v.shape[-3]), + a_log=torch.zeros(v.shape[-3]), + bias=torch.zeros(gate), + reqs=torch.zeros(1, dtype=torch.int32), + cu_seqlens=torch.arange(2, dtype=torch.int32), + state_verify_width=owner.verify_width, + state_compute_mode=owner.compute_mode, + state_gate_lower_bound=owner.gate_lower_bound, + run_config=run_config, + ) + else: + decode.replay( + owner, + 0, + q, + q, + v, + torch.zeros(1, gate), + torch.zeros(1, v.shape[-3]), + torch.zeros(v.shape[-3]), + torch.zeros(gate), + torch.zeros(1, dtype=torch.int32), + torch.arange(2, dtype=torch.int32), + run_config=run_config, + ) + assert captured.kwargs["PRECOMPUTE_PROJECTIONS"] is precompute + assert captured.kwargs["KDA"] is (cls is ReplayKDAState) + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_native_commit_selects_mtp_slot(cls): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + owner, _, _ = make_state(cls, step=2) + accepted = torch.tensor([3], dtype=torch.int32, device="cuda") + owner.commit_bounded_att_state(torch.tensor([1], device="cuda"), accepted, accepted) + assert owner.mtp_state_index.tolist() == [0, 2, 0, 0] + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_replay_fused_acceptance_keeps_cursors_inside_mode(monkeypatch, cls): + import sys + + owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") + history = getattr(owner, "runtime", owner) + calls = [] + monkeypatch.setattr(sys.modules[type(history).__module__], "commit_ssm_prefix", lambda *args: calls.append(args)) + widths = torch.tensor([1, 2, 3]) + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), widths) + assert len(calls) == 1 + assert calls[0][3] is widths + assert calls[0][4] is history.cursors + assert calls[0][0] is owner.mtp_state_index + + +def _commit_prefix_cpu(indices, reqs, lengths, req_verify_width, replay_cursors) -> None: + """CPU 上按请求行核对索引和 cursor。正式提交算子只跑 CUDA。""" + + indices[reqs] = (lengths - 1).to(indices.dtype) + hold = replay_cursors.numel() - 1 + for row in range(reqs.numel()): + req = int(reqs[row]) + if req == hold: + continue + rejected = int(req_verify_width[row]) - int(lengths[row]) + replay_cursors[req] = int(replay_cursors[req]) - rejected + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_acceptance_keeps_request_order_with_strided_inputs(device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + owner, _, _ = make_state(mode="replay", step=3, device=device) + owner.mtp_state_index.fill_(2) + history = getattr(owner, "runtime", owner) + history.cursors.zero_() + reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] + lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] + if device == "cuda": + flat_lengths = lengths.contiguous() + cu = torch.zeros(flat_lengths.numel() + 1, dtype=torch.int32, device=device) + cu[1:] = torch.cumsum(flat_lengths, 0) + advance_cursor(history, reqs.contiguous(), cu) + owner.commit_bounded_att_state(reqs, lengths, lengths) + else: + # 正式提交只接受 CUDA。CPU 用例只核对 stride 下的请求顺序。 + history.cursors[0] = 3 + history.cursors[2] = 1 + _commit_prefix_cpu(owner.mtp_state_index, reqs, lengths, lengths, history.cursors) + assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] + assert history.cursors.tolist() == [3, 0, 1, 0] diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index 1da9c1e286..622c8087ad 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -22,6 +22,19 @@ def single_rank_environment(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv(key, value) +@pytest.fixture +def runtime_start_args(monkeypatch: pytest.MonkeyPatch): + import json + from dataclasses import asdict + 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()))) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + @pytest.fixture def kv_allocator_environment(monkeypatch: pytest.MonkeyPatch): """Provide a service identity and reclaim the allocator's real shared counters.""" diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 530b95436a..6da740353d 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -48,11 +48,13 @@ def test_nope_cache_config_uses_native_mla_width(): @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("mode", ["native", "replay"]) def test_hybrid_checkpoint_restore_and_packed_kv_copy( monkeypatch, small_page, tp_world_size, mtp_step, + mode, ): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") @@ -106,6 +108,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( config=config, request_capacity=4, mtp_step=mtp_step, + ssm_state_mode=mode, full_state=main_full, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), draft_full_state=draft_full, @@ -143,6 +146,15 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( linear.ssm_state[:, 0].normal_() conv = linear.conv_state[:, 0, ..., :3].clone() ssm = linear.ssm_state[:, 0].clone() + if mode == "replay": + linear.raw_keys.zero_() + linear.raw_values.zero_() + linear.gates.fill_(-0.5) + linear.betas.zero_() + linear.cursors[0] = 1 + linear.cursors[2] = 3 + if mtp_step: + linear.mtp_state_index[2] = mtp_step resources.save_bounded_att_runtime_to_pages( req_indices=[0], state_token_lens=[8], @@ -162,12 +174,20 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( ) torch.cuda.synchronize() assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) - assert torch.equal(linear.ssm_state[:, 2 * (mtp_step + 1)], ssm) + restored_slot = 2 if mode == "replay" else 2 * (linear.mtp_step + 1) + assert torch.equal(linear.ssm_state[:, restored_slot], ssm) + if mode == "replay": + assert linear.cursors[2] == 0 + if mtp_step: + assert linear.mtp_state_index[2] == 0 # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 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() + if mode == "replay": + assert not linear.ssm_state[:, 2].any() + else: + assert not linear.ssm_state[:, 2 * (linear.mtp_step + 1) : 3 * (linear.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]]) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 8a1d70e3ac..30c4fc22c4 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -5,7 +5,10 @@ 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.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import ( + fused_recurrent_kda, +) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -23,7 +26,11 @@ def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_ste states = torch.randn(6 * width, heads, dim, dim, device="cuda", dtype=state_dtype) * 0.1 accepted = torch.tensor([width, 1, 1], device="cuda", dtype=torch.int32) for lengths in ([width, 1, 0], [2, width, 0], [1, 1, 0]): - cu = torch.tensor([0, lengths[0], sum(lengths), sum(lengths)], device="cuda", dtype=torch.int32) + cu = torch.tensor( + [0, lengths[0], sum(lengths), sum(lengths)], + device="cuda", + dtype=torch.int32, + ) tokens = sum(lengths) q, k, v, gate = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(4)] beta = torch.randn(1, tokens, heads, device="cuda", dtype=torch.bfloat16) @@ -100,7 +107,12 @@ def prefill(values, destination, tail): prefill(raw[:prefix], packed, ring) # Alternate expanded verification, rewound single-token drafting, and re-verification. - for start, count in [(prefix, width), (prefix + 1, 1), (prefix + 2, 1), (prefix + 1, width)]: + for start, count in [ + (prefix, width), + (prefix + 1, 1), + (prefix + 2, 1), + (prefix + 1, width), + ]: raw[start : start + count] = torch.randn(count, 256, device="cuda", dtype=torch.bfloat16) lengths = torch.arange(start + 1, start + count + 1, device="cuda", dtype=torch.int32) compress_pools( @@ -117,7 +129,11 @@ def prefill(values, destination, tail): b_mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), ) expected = torch.zeros_like(packed) - prefill(raw[: start + count], expected, torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16)) + prefill( + raw[: start + count], + expected, + torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16), + ) closing = torch.arange(start, start + count, device="cuda") closing = closing[(closing + 1) % 4 == 0] assert torch.equal(packed[closing], expected[closing]) @@ -125,17 +141,26 @@ def prefill(values, destination, tail): @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize("role", ["main", "draft0"]) -def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): +@pytest.mark.parametrize("mode,mtp_step", [("native", 2), ("replay", 2), ("native", 0), ("replay", 0)]) +def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mode, mtp_step): from types import SimpleNamespace from lightllm.common.basemodel.attention.base_att import AttControl from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update + from lightllm.common.basemodel.attention.linear.replay_kda import ( + ReplayKDADecodeAttState, + ) + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( + ReplayKDAState, + ) + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_update, + ) torch.manual_seed(53) - heads, dim, width = 2, 128, 3 + heads, dim, width = 2, 128, mtp_step + 1 hidden = heads * dim - lengths = [2, 1, 3] if dynamic else [3, 3, 3] + lengths = [2, 1, 3] if dynamic and mtp_step else [width] * 3 requests = [2, 0, 1] request_rows = [req for req, length in zip(requests, lengths) for _ in range(length)] offsets = [offset for length in lengths for offset in range(length)] @@ -143,10 +168,10 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): request_rows += [3, 3] offsets += [0, 0] tokens = len(request_rows) - conv = torch.randn(4, 3 * hidden, 5, device="cuda", dtype=torch.bfloat16) + conv = torch.randn(4, 3 * hidden, 3 + mtp_step, device="cuda", dtype=torch.bfloat16) ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") * 0.01 old_conv, expected_ssm = conv.clone(), ssm.clone() - accepted_offsets = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device="cuda") + accepted_offsets = torch.tensor([1, 2, 0, 0] if mtp_step else [0, 0, 0, 0], dtype=torch.int32, device="cuda") mixed = torch.randn(tokens, 3 * hidden, device="cuda", dtype=torch.bfloat16) gate = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) beta = torch.randn(tokens, heads, device="cuda", dtype=torch.bfloat16) @@ -184,7 +209,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): expected_ssm[req * width + token] = state[0] row += 1 backend = SimpleNamespace( - mtp_step=2, + mtp_step=mtp_step, tp_num_heads=heads, head_dim=dim, tp_hidden_size=hidden, @@ -203,11 +228,20 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): conv_dtype=conv.dtype, ssm_dtype=ssm.dtype, ) + selected = None 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) + state_cls = ReplayKDAState if mode == "replay" else KDAState + linear = state_cls(config, request_capacity=4, mtp_step=mtp_step, device="cuda") + if mtp_step: + linear.mtp_state_index.copy_(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) + linear.conv_state[0].copy_(conv) + if mode == "native": + linear.ssm_state[0].copy_(ssm) + else: + indexes = torch.arange(4, device="cuda") * width + accepted_offsets + linear.ssm_state[0].copy_(ssm[indexes]) + selected = linear resources.register_state(model_role, linear) resources.finish_state_registration() manager = SimpleNamespace(HOLD_REQUEST_ID=3) @@ -223,7 +257,8 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): linear_A_log=SimpleNamespace(weight=a), linear_dt_bias=SimpleNamespace(weight=bias), ) - state = KDADecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + decode_cls = ReplayKDADecodeAttState if mode == "replay" else KDADecodeAttState + state = decode_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) state.init_state() actual = state.decode_att( None, @@ -232,11 +267,42 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): AttControl( linear_att_decode=True, linear_att_decode_dict=dict( - layer_weight=weight, layer_num=0, mixed_qkv=mixed, raw_gate=gate, raw_beta=beta + layer_weight=weight, + layer_num=0, + mixed_qkv=mixed, + raw_gate=gate, + raw_beta=beta, ), ), ) - torch.testing.assert_close(actual.reshape(tokens, heads, dim)[:row], torch.cat(expected), atol=2e-3, rtol=1e-2) - torch.testing.assert_close(ssm, expected_ssm, atol=2e-3, rtol=1e-2) - if dynamic: - assert torch.equal(conv[3], old_conv[3]) + torch.testing.assert_close( + actual.reshape(tokens, heads, dim)[:row], + torch.cat(expected), + atol=2e-3, + rtol=1e-2, + ) + if mode == "native": + checked_slots = len(requests) if dynamic and not mtp_step else expected_ssm.shape[0] + torch.testing.assert_close( + selected.ssm_state[0, :checked_slots], + expected_ssm[:checked_slots], + atol=2e-3, + rtol=1e-2, + ) + else: + for req, length in zip(requests, lengths): + saved = selected.ssm_state + clone = saved.clone() + selected.ssm_state = clone + try: + fold_prefix(selected, torch.tensor([req], dtype=torch.int32, device="cuda")) + finally: + selected.ssm_state = saved + torch.testing.assert_close( + clone[0, req], + expected_ssm[req * width + length - 1], + atol=2e-3, + rtol=1e-2, + ) + if dynamic and mtp_step: + assert torch.equal(selected.conv_state[0, 3], old_conv[3]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index e72e4d6bee..4d15d6bf1c 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -46,7 +46,7 @@ def make_mems(monkeypatch, request): get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - def create(tp, target_mtp_step=None): + def create(tp, target_mtp_step=None, mode="native"): local_mtp_step = mtp_step if target_mtp_step is None else target_mtp_step config = Glm5NextResourceConfig( tp_world_size=tp, @@ -87,6 +87,7 @@ def create(tp, target_mtp_step=None): config=config, request_capacity=4, mtp_step=local_mtp_step, + ssm_state_mode=mode, tp_rank=rank, full_state=main, full_allocator=FullAttKVAllocator(full.capacity, full.page_size), @@ -112,16 +113,21 @@ def create(tp, target_mtp_step=None): @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) @pytest.mark.parametrize("remainder", range(4)) @pytest.mark.parametrize("decode_mtp_step", [0, 1, 5]) +@pytest.mark.parametrize("source_mode", ["native", "replay"]) +@pytest.mark.parametrize("dest_mode", ["native", "replay"]) def test_pd_kv_and_runtime_state_roundtrip( make_mems, prefill_tp, decode_tp, remainder, decode_mtp_step, + source_mode, + dest_mode, ): - source, dest = make_mems(prefill_tp), make_mems(decode_tp, target_mtp_step=decode_mtp_step) + source = make_mems(prefill_tp, mode=source_mode) + dest = make_mems(decode_tp, target_mtp_step=decode_mtp_step, mode=dest_mode) + dst_ssm_stride = 1 if dest_mode == "replay" else dest[0].linear.mtp_step + 1 mtp_size = get_env_start_args().mtp_step + 1 - 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 @@ -159,7 +165,15 @@ def test_pd_kv_and_runtime_state_roundtrip( 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]) + src_slot = src_req if source_mode == "replay" else src_req * (mem.linear.mtp_step + 1) + mem.linear.ssm_state[:, src_slot].copy_(ssm[:, heads]) + if source_mode == "replay": + backend = mem.linear + backend.raw_keys.zero_() + backend.raw_values.zero_() + backend.gates.fill_(-0.5) + backend.betas.zero_() + backend.cursors[src_req] = 1 offset = 0 for state in mem.tails: end = offset + len(state.config.logical_layer_ids) @@ -173,6 +187,10 @@ def test_pd_kv_and_runtime_state_roundtrip( state.storage.zero_() mem.linear.conv_state.normal_() mem.linear.ssm_state.normal_() + if dest_mode == "replay": + mem.linear.cursors[dst_req] = 3 + if mem.linear.mtp_enabled: + mem.linear.mtp_state_index[dst_req] = decode_mtp_step for state in mem.tails: state.buffer.normal_() original_tails.append(torch.cat([state.buffer[:, dst_req] for state in mem.tails]).clone()) @@ -181,7 +199,7 @@ def test_pd_kv_and_runtime_state_roundtrip( x[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( (mem.linear.conv_state, 1), - (mem.linear.ssm_state, dst_mtp_size), + (mem.linear.ssm_state, dst_ssm_stride), *((state.buffer, 1) for state in mem.tails), ) ] @@ -235,18 +253,20 @@ def test_pd_kv_and_runtime_state_roundtrip( 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]) + assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_ssm_stride], 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: + if dest_mode == "replay": + assert mem.linear.cursors[dst_req] == 0 + if mem.linear.mtp_enabled: assert mem.linear.mtp_state_index[dst_req] == 0 for (state, stride), expected in zip( ( (mem.linear.conv_state, 1), - (mem.linear.ssm_state, dst_mtp_size), + (mem.linear.ssm_state, dst_ssm_stride), *((state.buffer, 1) for state in mem.tails), ), untouched[rank], diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py new file mode 100644 index 0000000000..c25576a19a --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -0,0 +1,151 @@ +from types import SimpleNamespace + +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.attention.base_att import AttControl +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("mtp_step", [0, 2]) +@pytest.mark.parametrize( + "start,end,count,logical_layer,storage_layer", + [(4, 28, 1, 4, 3), (4, 28, 3, 28, 21), (1, 6, 2, 13, 10)], + ids=["no-loop", "aligned-loop", "unaligned-loop"], +) +def test_gdn_decode_keeps_logical_layer_state( + monkeypatch, device, mtp_step, start, end, count, logical_layer, storage_layer +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + torch.manual_seed(1585) + settings = dict(start=start, end=end, eval_count=count, use_gate=False) + config = dict( + num_hidden_layers=32, + full_attention_interval=4, + layer_types=["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(32)], + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + hidden_act="silu", + ) + indices = get_loop_layer_indices(settings, 32) + kinds = get_loop_layer_types(settings, config) + layers = tuple(i for i, kind in enumerate(kinds) if kind == "linear_attention") + args = StartArgs(mtp_step=mtp_step, linear_att_ssm_data_type="float32") + monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.basemodel.attention.base_att.get_env_start_args", lambda: args) + backend = object.__new__(TritonLinearAttBackend) + backend.__init__( + SimpleNamespace( + config=config, + tp_world_size_=1, + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: mtp_step), + ) + ) + dtype = torch.bfloat16 if device == "cuda" else torch.float32 + layout = ConvSSMStateConfig(layers, 1, 2, 32, 32, 3, dtype, torch.float32) + resource = GDNState(layout, request_capacity=3, mtp_step=mtp_step, device=device) + resources = AttentionStateManager(request_capacity=3) + resources.register_state("main", resource) + resources.finish_state_registration() + conv_view, ssm_view = resource.get_layer_view(logical_layer_id=logical_layer) + assert conv_view.data_ptr() == resource.conv_state[storage_layer].data_ptr() + assert ssm_view.data_ptr() == resource.ssm_state[storage_layer].data_ptr() + physical_weights = [ + SimpleNamespace( + layer_num_=i, + linear_conv1d=SimpleNamespace( + mm_param=SimpleNamespace(weight=torch.randn(layout.conv_channels, 4, device=device, dtype=dtype) * 0.1), + bias=None, + ), + linear_A_log=SimpleNamespace(weight=torch.zeros(2, device=device)), + linear_dt_bias=SimpleNamespace(weight=torch.zeros(2, device=device)), + ) + for i in range(32) + ] + execution_weights = [physical_weights[i] for i in indices] + weight = execution_weights[logical_layer] + assert weight is physical_weights[indices[logical_layer]] + + if device == "cpu": + + def conv_pack(mixed, z, a, b, conv, *args): + conv[1].add_(1) + return (*backend._rearrange_mixed_qkv(mixed, decode=True), z, a, b) + + def conv_mtp(mixed, conv, *args, **kwargs): + conv[1].add_(1) + return mixed + + def recurrent(*, initial_state, v, **kwargs): + initial_state[mtp_step + 1 : 2 * (mtp_step + 1)].add_(1) + return v, None + + monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp + ) + monkeypatch.setattr( + "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent + ) + monkeypatch.setattr( + "lightllm.common.basemodel.attention.linear.gdn.mtp_fused_recurrent_gated_delta_rule", + recurrent, + ) + + infer = InferStateInfo() + infer.att_state_manager, infer.att_state_role = resources, "main" + infer.b_req_idx = torch.tensor([1], device=device, dtype=torch.int32) + control = dict(layer_num=logical_layer, layer_weight=weight) + mixed_dim = backend.tp_qkvz_dim + backend.tp_ba_dim + if device == "cuda": + infer.b1_cu_q_seq_len = torch.tensor([0, 5], device=device, dtype=torch.int32) + infer.b_ready_cache_len = torch.zeros(1, device=device, dtype=torch.int32) + prefill = backend.create_att_prefill_state(infer) + prefill.init_state() + prefill.prefill_att( + None, + None, + None, + AttControl( + linear_att_prefill=True, + linear_att_prefill_dict=dict( + **control, mixed_qkvzba=torch.randn(5, mixed_dim, device=device, dtype=dtype) + ), + ), + ) + + infer.batch_size = mtp_step + 1 + infer.b_req_idx = infer.b_req_idx.repeat(infer.batch_size) + decode = backend.create_att_decode_state(infer) + untouched = [i for i in range(len(layers)) if i != storage_layer] + for _ in range(2): + before_conv, before_ssm = resource.conv_state.clone(), resource.ssm_state.clone() + decode.init_state() + output, _ = decode.decode_att( + None, + None, + None, + AttControl( + linear_att_decode=True, + linear_att_decode_dict=dict( + **control, mixed_qkvzba=torch.randn(infer.batch_size, mixed_dim, device=device, dtype=dtype) + ), + ), + ) + assert torch.isfinite(output).all() + assert not torch.equal(resource.conv_state[storage_layer], before_conv[storage_layer]) + assert not torch.equal(resource.ssm_state[storage_layer], before_ssm[storage_layer]) + assert torch.equal(resource.conv_state[untouched], before_conv[untouched]) + assert torch.equal(resource.ssm_state[untouched], before_ssm[untouched]) + assert weight.layer_num_ == indices[logical_layer] diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py index 230e23b0e5..1c391c09cd 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -102,3 +102,37 @@ def test_whole_image_blocks_fit_batch_at_startup( else: with pytest.raises(AssertionError, match="batch_max_tokens >= max_image_token_count"): _launch_subprocesses(args) + + +@pytest.mark.parametrize("mode,step", [("native", 0), ("replay", 0), ("replay", 4)]) +def test_s3_configuration_is_prepared_before_ssm_validation(launch_args, monkeypatch, tmp_path, mode, step): + from pathlib import Path + from lightllm.utils.petrel_helper import PetrelHelper + + monkeypatch.chdir(tmp_path) + model_dir = "s3://test-bucket/qwen" + config = {"model_type": "qwen3_5_text"} + downloads = [] + + def list_files(path, extension): + assert path == model_dir + return ["config.json"] if extension == "json" else [] + + def download_json(remote, local): + downloads.append(remote) + Path(local).write_text(json.dumps(config)) + + monkeypatch.setattr(PetrelHelper, "list", list_files) + monkeypatch.setattr(PetrelHelper, "download_json", download_json) + args = launch_args("qwen3_5_text", ssm_state_mode=mode, mtp_step=step, mtp_mode="eagle_with_att" if step else None) + args.model_dir = model_dir + assert not Path(model_dir, "config.json").exists() + with pytest.raises(RuntimeError, match="startup validation passed"): + _launch_subprocesses(args) + assert downloads == [model_dir + "/config.json"] + if mode == "replay": + expected_compute_mode = "precompute" if step > 0 else "inline" + assert args.ssm_replay_config == { + "cache_len": 16, + "compute_mode": expected_compute_mode, + } diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py new file mode 100644 index 0000000000..e60e71028d --- /dev/null +++ b/unit_tests/server/test_ssm_state_mode.py @@ -0,0 +1,165 @@ +import json + +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.api_start import _validate_and_normalize_ssm_config +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize( + "mode,dtype,step,model,capacity,error", + [ + ("native", "bfloat16", 0, "llama", 16, None), + ("native", "bfloat16", 32, "llama", 16, None), + ("replay", "float32", 0, "sensenova_loop", 16, None), + ("replay", "bfloat16", 2, "sensenova_loop", 8, None), + ("replay", "float32", 0, "glm5_next", 8, None), + ("replay", "bfloat16", 2, "glm5_next_text", 8, None), + ("replay", "float32", 0, "qwen3_next", 16, None), + ("replay", "float32", 15, "qwen3_5", 16, None), + ("replay", "bfloat16", 2, "qwen3_5", 16, None), + ("replay", "bfloat16", 2, "qwen3_5", 4, None), + ("replay", "bfloat16", 0, "qwen3_5", 16, None), + ("replay", "float32", 4, "qwen3_5", 4, "capacity"), + ("replay", "float32", 2, "llama", 16, "require a GDN"), + ], +) +def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, error): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = make_argument_parser().parse_args( + [ + "--ssm_state_mode", + mode, + "--linear_att_ssm_data_type", + dtype, + "--mtp_step", + str(step), + "--ssm_replay_config", + json.dumps({"cache_len": capacity} if mode == "replay" else {}), + ] + ) + if error: + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + else: + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("step", [15, 16]) +def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs(ssm_state_mode=mode, mtp_step=step) + if step <= 15: + _validate_and_normalize_ssm_config(args) + else: + with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("cache_len", [32, 64]) +def test_replay_large_cache_does_not_bypass_mtp_width_limit(monkeypatch, model, compute_mode, cache_len): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs( + ssm_state_mode="replay", + ssm_replay_config={"cache_len": cache_len, "compute_mode": compute_mode}, + linear_att_ssm_data_type="bfloat16", + mtp_step=16, + ) + with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): + _validate_and_normalize_ssm_config(args) + + +def test_ssm_mode_defaults_and_removed_flag(): + parser = make_argument_parser() + assert parser.parse_args([]).ssm_state_mode == StartArgs().ssm_state_mode == "native" + assert parser.parse_args([]).ssm_replay_config == StartArgs().ssm_replay_config == {} + with pytest.raises(SystemExit): + parser.parse_args(["--enable_replayssm"]) + with pytest.raises(SystemExit): + parser.parse_args(["--ssm_state_mode", "auto"]) + + +@pytest.mark.parametrize( + "overrides,error", + [ + ({"ssm_replay_config": {"cache_len": 2}}, "power of two and at least 4"), + ({"ssm_replay_config": {"cache_len": 6}}, "power of two and at least 4"), + ({"mtp_step": -1}, "mtp_step >= 0"), + ({"linear_att_ssm_data_type": "float16"}, "FP32 or BF16"), + ], +) +def test_programmatic_replay_config_rejected_before_model_loading(monkeypatch, overrides, error): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = StartArgs(ssm_state_mode="replay") + for key, value in overrides.items(): + setattr(args, key, value) + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + + +def test_replay_compute_mode_requires_replay(monkeypatch): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + parser = make_argument_parser() + args = parser.parse_args(["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}']) + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize( + "config,error", + [ + ([], "JSON object"), + (None, "JSON object"), + ({"projection_mode": "inline"}, "Unknown"), + ({"cache_len": True}, "capacity"), + ({"cache_len": 16.0}, "capacity"), + ({"cache_len": "16"}, "capacity"), + ({"cache_len": 128}, "capacity"), + ({"compute_mode": "auto"}, "compute_mode"), + ({"fold_method": "raw"}, "Unknown"), + ], +) +def test_invalid_replay_config(config, error): + args = StartArgs(ssm_state_mode="replay", ssm_replay_config=config) + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("mode", ["native"]) +def test_replay_config_rejects_other_modes(mode): + args = StartArgs(ssm_state_mode=mode, ssm_replay_config={"cache_len": 16}) + with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): + _validate_and_normalize_ssm_config(args) + + +def test_partial_replay_config_defaults_and_serialization(monkeypatch): + from dataclasses import asdict + + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = make_argument_parser().parse_args( + ["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}'] + ) + _validate_and_normalize_ssm_config(args) + assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "precompute"} + original = StartArgs(ssm_state_mode="replay", ssm_replay_config=args.ssm_replay_config) + restored = StartArgs(**json.loads(json.dumps(asdict(original)))) + _validate_and_normalize_ssm_config(restored) + assert restored.ssm_replay_config == args.ssm_replay_config + assert StartArgs().ssm_replay_config == {} + + +@pytest.mark.parametrize( + "option,value", + [ + ("--replayssm_cache_len", "16"), + ("--replayssm_projection_mode", "inline"), + ("--ssm_replay_config", "{invalid}"), + ], +) +def test_removed_flags_and_malformed_json(option, value): + with pytest.raises(SystemExit): + make_argument_parser().parse_args([option, value])