diff --git a/lightllm/common/attention_state/bounded_att/base_state.py b/lightllm/common/attention_state/bounded_att/base_state.py index 6a96e76dc9..c852a5c7e9 100644 --- a/lightllm/common/attention_state/bounded_att/base_state.py +++ b/lightllm/common/attention_state/bounded_att/base_state.py @@ -7,3 +7,7 @@ class BoundedAttentionState(CoreInterface, CpuCacheInterface, DpCopyInterface, PdInterface): """单种 bounded-attention 方案的数据操作，索引和槽位由调用方分配。""" + + def validate_bounded_att_commit(self, context): + + pass diff --git a/lightllm/common/attention_state/bounded_att/conv_ssm_state.py b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py new file mode 100644 index 0000000000..47f6e15515 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/conv_ssm_state.py @@ -0,0 +1,556 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING, List, NamedTuple +from math import prod + +import torch + +from lightllm.utils.log_utils import init_logger + +from .ssm import create_ssm_state +from .ssm.forward import SSMForward + +from ..utils import deepcopy_result, dtype_name +from ..types import AxisMapping, StateTransferLayout, TensorField, TransferRegion +from .base_state import BoundedAttentionState +from .conv_ssm_types import ConvSSMStateConfig +from .transfer_utils import bounded_att_buffer_copies +from .copy_kernel import copy_bounded_att_tensors + +if TYPE_CHECKING: + from lightllm.common.attention_state.types import PageKind + + +logger = init_logger(__name__) + + +class ConvSSMTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +class ConvSSMState(BoundedAttentionState): + + algorithm: str + pages_pool_class: type + + def __init__( + self, + config: ConvSSMStateConfig, + 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, + ): + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self._logical_to_storage_layer = {layer: index for index, layer in enumerate(config.logical_layer_ids)} + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.conv_state = torch.zeros( + (len(config.logical_layer_ids), request_capacity, config.conv_channels, config.conv_width + mtp_step), + dtype=config.conv_dtype, + device=device, + ) + self._ssm = create_ssm_state( + config, + request_capacity, + mtp_step, + device, + ssm_state_mode, + ssm_replay_config, + kda=self.algorithm == "kda", + gate_lower_bound=gate_lower_bound, + ) + self.big_pages_pool = None + self.small_pages_pool = None + + @property + def ssm_state(self): + return self._ssm.state + + @property + def ssm_slots_per_req(self): + return self._ssm.slots_per_request + + def _reset_tracking(self, req_idx): + self._ssm.clear_history(req_idx) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].zero_() + + def reset(self, req_idx): + self._reset_tracking(req_idx) + self.conv_state[:, req_idx].zero_() + start = req_idx * self.ssm_slots_per_req + self.ssm_state[:, start : start + self.ssm_slots_per_req].zero_() + + def prepare_prefill(self, reqs): + self._ssm.merge_accepted_updates(reqs) + return reqs, reqs * self.ssm_slots_per_req + + def prepare_decode(self, reqs, cu_seqlens=None): + speculative = cu_seqlens is not None + accepted = None + if speculative: + if self.mtp_state_index is None: + raise ValueError("speculative decode requires MTP state slots") + accepted = self.mtp_state_index[reqs] + 1 + accepted = torch.where(reqs == self.request_capacity - 1, 1, accepted) + self._ssm.prepare_decode(reqs, cu_seqlens) + if not speculative: + self._ssm.accept_single_token(reqs) + return SSMForward(self, reqs, self._ssm.decode_indices(reqs, speculative), cu_seqlens, accepted) + + def get_accepted_state(self, req_idx): + + width = self.config.conv_width + accepted = self.mtp_state_index[req_idx] if self.mtp_state_index is not None else None + conv = self.conv_state[:, req_idx, ..., :width] + if accepted is not None: + offsets = torch.arange(width, device=conv.device) + accepted + conv = self.conv_state[:, req_idx].index_select(-1, offsets) + return ConvSSMTensors(conv, self._ssm.get_accepted_state(req_idx, accepted)) + + def prepare_restore(self, req_idx): + + self._reset_tracking(req_idx) + return ConvSSMTensors( + self.conv_state[:, req_idx, ..., : self.config.conv_width], + self.ssm_state[:, req_idx * self.ssm_slots_per_req], + ) + + def restore(self, req_idx, conv, ssm): + target = self.prepare_restore(req_idx) + target.conv.copy_(conv, non_blocking=True) + target.ssm.copy_(ssm, non_blocking=True) + + def get_storage_layer_index(self, logical_layer_id): + return self._logical_to_storage_layer[logical_layer_id] + + def get_state_type(self) -> str: + 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 = ( + self.pages_pool_class( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + self.pages_pool_class( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" + return logical_layer_id in self._logical_to_storage_layer + + def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" + storage_layer_index = self._logical_to_storage_layer[logical_layer_id] + return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] + + def reset_bounded_att_state(self, req_indices: List[int]) -> None: + """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.reset(req_idx) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + accepted = self.get_accepted_state(req_idx) + conv_copies.append((accepted.conv, pool.conv_state[slot_id])) + ssm_copies.append((accepted.ssm, pool.ssm_state[slot_id])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: List[int], + state_token_lens: List[int], + page_kind: PageKind, + page_slot_ids: List[int], + ) -> None: + """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" + assert isinstance(req_indices, list) and isinstance( + page_slot_ids, list + ), "request and slot indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + target = self.prepare_restore(req_idx) + conv_copies.append((pool.conv_state[slot_id], target.conv)) + ssm_copies.append((pool.ssm_state[slot_id], target.ssm)) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def validate_bounded_att_commit(self, context): + if self.mtp_state_index is not None: + self._ssm.validate_commit(context, self.state_name) + + def commit_bounded_att_state(self, req_indices, accepted_steps, *, context=None): + if context is not None: + context.validate(req_indices.numel()) + if self.mtp_state_index is not None: + self._ssm.commit(req_indices, accepted_steps, self.mtp_state_index, context, self.state_name) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 + + region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 + 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 + """ + return self._get_transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: List[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: List[int], + big_page_slot_ids: List[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 + pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] + # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 + # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp( + self, + src_page_ids: List[int], + dst_page_ids: List[int], + remote_dp_bounded_state: ConvSSMState, + ) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 + # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def copy_runtime_state_from_dp( + self, + src_req_indices: List[int], + dst_req_indices: List[int], + remote_dp_bounded_state: ConvSSMState, + ) -> None: + """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + conv_copies, ssm_copies = [], [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + source = remote_dp_bounded_state.get_accepted_state(src_req) + target = self.prepare_restore(dst_req) + conv_copies.append((source.conv, target.conv)) + ssm_copies.append((source.ssm, target.ssm)) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._get_transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """写出 state_token_len 处的已接受状态，不修改源 Replay 历史。""" + 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._fields(*self.get_accepted_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: + """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + fields = self._fields(*self.prepare_restore(req_idx)) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、私有辅助实现 + # ======================================================================== + + def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + conv_region = TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping( + axis="window", + local_start=0, + external_start=0, + length=config.conv_width, + ), + ), + write_enabled=write_enabled, + ) + regions.append(conv_region) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=head_start, + length=config.num_v_heads, + ), + AxisMapping( + axis="key_dim", + local_start=0, + external_start=0, + length=config.key_dim, + ), + AxisMapping( + axis="value_dim", + local_start=0, + external_start=0, + length=config.value_dim, + ), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + @staticmethod + def _fields(conv: torch.Tensor, ssm: torch.Tensor): + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } diff --git a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py index 3fe6f2317e..fff372e93e 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,6 +13,10 @@ 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 请求运行态，大小页池由后续容量分配流程初始化。""" return GDNState( @@ -20,4 +24,7 @@ def create_gdn_state( request_capacity=request_capacity, mtp_step=mtp_step, device=device, + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, ) 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..59c64b8ea0 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/gdn_state.py @@ -1,522 +1,8 @@ -from __future__ import annotations - -from functools import lru_cache -from typing import TYPE_CHECKING, List -from math import prod - -import torch - -from lightllm.utils.log_utils import init_logger - -from ...utils import deepcopy_result, dtype_name -from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion -from ..base_state import BoundedAttentionState -from ..conv_ssm_types import ConvSSMStateConfig +from ..conv_ssm_state import ConvSSMState from .gdn_pages_pool import GDNPagesPool -from ..transfer_utils import bounded_att_buffer_copies -from ..copy_kernel import copy_bounded_att_tensors - -if TYPE_CHECKING: - from lightllm.common.attention_state.types import PageKind - - -logger = init_logger(__name__) - - -class GDNState(BoundedAttentionState): - """Gated DeltaNet 的 Conv 历史、SSM 和 MTP 游标，支持不同的 K/V head 数及维度。""" - - big_pages_pool: GDNPagesPool | None - small_pages_pool: GDNPagesPool | None - - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - - def __init__( - self, - config: ConvSSMStateConfig, - request_capacity: int, - mtp_step: int = 0, - device: str | torch.device = "cuda", - ) -> None: - """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" - assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" - self.config = config - self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 - self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 - self.mtp_step = mtp_step - self._logical_to_storage_layer = { - logical_layer_id: storage_layer_index - for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) - } - storage_layer_count = len(config.logical_layer_ids) - self.conv_state = torch.zeros( - (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), - dtype=config.conv_dtype, - device=device, - ) - self.ssm_state = torch.zeros( - ( - storage_layer_count, - request_capacity * (mtp_step + 1), - config.num_v_heads, - config.key_dim, - config.value_dim, - ), - dtype=config.ssm_dtype, - device=device, - ) - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None - self.big_pages_pool = None - self.small_pages_pool = None - - def get_state_type(self) -> str: - """返回 GDN 算法及 Conv/SSM 存储精度。""" - return f"gdn_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" - - def initialize_page_pools( - self, - big_page_capacity: int, - small_page_capacity: int = 0, - device: str | torch.device | int = "cpu", - pin_memory: bool = True, - ) -> None: - """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" - assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" - assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" - self.big_pages_pool = ( - GDNPagesPool( - size=big_page_capacity, - config=self.config, - device=device, - pin_memory=pin_memory, - ) - if big_page_capacity - else None - ) - self.small_pages_pool = ( - GDNPagesPool( - size=small_page_capacity, - config=self.config, - device=device, - pin_memory=pin_memory, - ) - if small_page_capacity - else None - ) - - # ======================================================================== - # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 - # ======================================================================== - - def contains_logical_layer(self, logical_layer_id: int) -> bool: - """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" - return logical_layer_id in self._logical_to_storage_layer - - def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: - """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] - return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] - - def reset_bounded_att_state(self, req_indices: List[int]) -> None: - """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" - assert isinstance(req_indices, list), "request indices must be a list" - for req_idx in req_indices: - self.conv_state[:, req_idx].zero_() - start = req_idx * (self.mtp_step + 1) - self.ssm_state[:, start : start + self.mtp_step + 1].zero_() - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].fill_(0) - - def save_bounded_att_runtime_to_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance( - page_slot_ids, list - ), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = self._pool(page_kind) - # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 - # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) - ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - def restore_bounded_att_runtime_from_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance( - page_slot_ids, list - ), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = self._pool(page_kind) - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) - ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for req_idx in req_indices: - self.mtp_state_index[req_idx].fill_(0) - - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - ) -> None: - """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" - assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" - if self.mtp_state_index is not None: - self.mtp_state_index[req_indices] = accepted_steps - 1 - - # ======================================================================== - # 三、CPU Cache 数据交换 - # ======================================================================== - - @deepcopy_result - @lru_cache(maxsize=None) - def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 - - region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 - 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 - """ - return self._get_transfer_regions() - - def write_state_pages_to_cpu( - self, - page_kind: PageKind, - page_slot_ids: List[int], - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - layout: StateTransferLayout, - cpu_page_ready_flags: list[bool] | None = None, - ) -> None: - """将指定页池槽位的快照写入 CPU 页。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - pool = self._pool(page_kind) - assert isinstance(page_slot_ids, list), "slot indices must be a list" - assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" - if cpu_page_ready_flags is not None: - assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" - page_fields, page_indices = [], [] - for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): - # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 - if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: - continue - page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) - page_indices.append(page) - for copies in bounded_att_buffer_copies( - page_fields=page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=page_indices, - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - def read_state_pages_from_cpu( - self, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - big_page_slot_ids: List[int], - layout: StateTransferLayout, - ) -> None: - """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - pool = self.big_pages_pool - assert isinstance(big_page_slot_ids, list), "slot indices must be a list" - assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" - # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 - pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] - # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 - # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 - for copies in bounded_att_buffer_copies( - page_fields=pool_page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=cpu_page_indices, - write=False, - ): - copy_bounded_att_tensors(copies=copies) - - # ======================================================================== - # 四、DP 迁移：大页快照与基准运行态 - # ======================================================================== - - def copy_big_pages_from_dp( - self, - src_page_ids: List[int], - dst_page_ids: List[int], - remote_dp_bounded_state: GDNState, - ) -> None: - """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" - assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" - assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" - assert ( - type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - src_pool = remote_dp_bounded_state.big_pages_pool - dst_pool = self.big_pages_pool - assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" - # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 - # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 - conv_copies, ssm_copies = [], [] - for src_page, dst_page in zip(src_page_ids, dst_page_ids): - conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) - ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - def copy_runtime_state_from_dp( - self, - src_req_indices: List[int], - dst_req_indices: List[int], - remote_dp_bounded_state: GDNState, - ) -> None: - """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" - assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" - assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" - assert ( - type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； - # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 - # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 - conv_copies, ssm_copies = [], [] - for src_req, dst_req in zip(src_req_indices, dst_req_indices): - conv_copies.append( - ( - remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], - self.conv_state[:, dst_req, :, : self.config.conv_width], - ) - ) - ssm_copies.append( - ( - remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], - self.ssm_state[:, dst_req * (self.mtp_step + 1)], - ) - ) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for dst_req in dst_req_indices: - self.mtp_state_index[dst_req].fill_(0) - - # ======================================================================== - # 五、PD 分离传输 - # ======================================================================== - - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" - return self._get_transfer_regions() - - def write_to_pd_page( - self, - req_idx: int, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._runtime_fields(req_idx) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - req_idx: int, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._runtime_fields(req_idx) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=False, - ): - copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].fill_(0) - - # ======================================================================== - # 六、私有辅助实现 - # ======================================================================== - - def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: - """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" - # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 - config = self.config - storage_layer_count = len(config.logical_layer_ids) - conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim - conv_shape = (storage_layer_count, conv_channels, config.conv_width) - conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize - regions = [] - logger.debug( - "%s building transfer regions: TP=%d/%d, Conv shape=%s", - self.state_name, - config.tp_rank, - config.tp_world_size, - conv_shape, - ) - - # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 - local_channel_start = 0 - global_channel_start = 0 - for name, local_head_count, global_head_count, head_dim in ( - ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), - ("K", config.num_k_heads, config.global_k_heads, config.key_dim), - ("V", config.num_v_heads, config.global_v_heads, config.value_dim), - ): - # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 - replica_count = config.tp_world_size * local_head_count // global_head_count - head_start = config.tp_rank // replica_count * local_head_count - write_enabled = config.tp_rank % replica_count == 0 - local_channel_count = local_head_count * head_dim - target_channel_start = global_channel_start + head_start * head_dim - conv_region = TransferRegion( - field="conv", - byte_start=0, - byte_end=conv_byte_end, - dtype=config.conv_dtype, - axes=("layer", "channel", "window"), - shape=conv_shape, - axis_mappings=( - AxisMapping( - axis="channel", - local_start=local_channel_start, - external_start=target_channel_start, - length=local_channel_count, - ), - AxisMapping( - axis="window", - local_start=0, - external_start=0, - length=config.conv_width, - ), - ), - write_enabled=write_enabled, - ) - regions.append(conv_region) - logger.debug( - "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " - "channels=%d, replicas=%d, write_enabled=%s", - self.state_name, - name, - local_channel_start, - target_channel_start, - local_channel_count, - replica_count, - write_enabled, - ) - local_channel_start += local_channel_count - global_channel_start += global_head_count * head_dim - - # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 - replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads - head_start = config.tp_rank // replica_count * config.num_v_heads - write_enabled = config.tp_rank % replica_count == 0 - ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) - element_bytes = config.ssm_dtype.itemsize - ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes - ssm_region = TransferRegion( - field="ssm", - byte_start=ssm_byte_start, - byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, - dtype=config.ssm_dtype, - axes=("layer", "head", "key_dim", "value_dim"), - shape=ssm_shape, - axis_mappings=( - AxisMapping( - axis="head", - local_start=0, - external_start=head_start, - length=config.num_v_heads, - ), - AxisMapping( - axis="key_dim", - local_start=0, - external_start=0, - length=config.key_dim, - ), - AxisMapping( - axis="value_dim", - local_start=0, - external_start=0, - length=config.value_dim, - ), - ), - write_enabled=write_enabled, - ) - regions.append(ssm_region) - logger.debug( - "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", - self.state_name, - head_start, - config.num_v_heads, - write_enabled, - ssm_region.byte_start, - ssm_region.byte_end, - ) - return tuple(regions) - def _pool(self, page_kind: PageKind): - """根据大小页类型取得已启用的页池，不访问 allocator。""" - assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" - pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool - assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" - return pool - def _runtime_fields(self, req: int): - """返回请求基准位置的 Conv/SSM 视图，供恢复和传输使用。""" - # TODO：当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0。 - # 若以后支持该场景，需要按已接受游标选取 Conv 窗口及 SSM 槽位。 - return self._fields( - self.conv_state[:, req, :, : self.config.conv_width], self.ssm_state[:, req * (self.mtp_step + 1)] - ) +class GDNState(ConvSSMState): - @staticmethod - def _fields(conv: torch.Tensor, ssm: torch.Tensor): - return { - "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), - "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), - } + algorithm = "gdn" + pages_pool_class = GDNPagesPool diff --git a/lightllm/common/attention_state/bounded_att/interface.py b/lightllm/common/attention_state/bounded_att/interface.py index fe1fceb3bd..fa5d1e01ee 100644 --- a/lightllm/common/attention_state/bounded_att/interface.py +++ b/lightllm/common/attention_state/bounded_att/interface.py @@ -102,8 +102,9 @@ def commit_bounded_att_state( self, req_indices: torch.Tensor, accepted_steps: torch.Tensor, + *, + context=None, ) -> None: - """两个输入均为 CUDA Tensor；int32 接受步数包含基准输入步，候选区与游标由实现管理。""" ... @@ -173,10 +174,7 @@ def copy_runtime_state_from_dp( dst_req_indices: List[int], remote_dp_bounded_state: CoreInterface, ) -> None: - """复制其他 DP 同一 TP 分片的 prefill 基准运行态并重置目标游标，不修改页池。 - 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 - """ ... @@ -205,10 +203,7 @@ def write_to_pd_page( state_token_len: int, layout: StateTransferLayout, ) -> None: - """将指定槽位在 state_token_len 绝对长度处的 prefill 基准运行态写入 PD 页。 - 调用方保证源状态尚未执行 MTP 提交；KDA/GDN 的源游标必须为 0。 - """ ... @abstractmethod 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..180c5dcc7b 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,6 +13,10 @@ 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 请求运行态，大小页池由后续容量分配流程初始化。""" return KDAState( @@ -20,4 +24,7 @@ def create_kda_state( request_capacity=request_capacity, mtp_step=mtp_step, device=device, + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, ) 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..903da44b68 100644 --- a/lightllm/common/attention_state/bounded_att/kda/kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/kda_state.py @@ -1,522 +1,8 @@ -from __future__ import annotations - -from functools import lru_cache -from typing import TYPE_CHECKING, List -from math import prod - -import torch - -from lightllm.utils.log_utils import init_logger - -from ...utils import deepcopy_result, dtype_name -from ...types import AxisMapping, StateTransferLayout, TensorField, TransferRegion -from ..base_state import BoundedAttentionState -from ..conv_ssm_types import ConvSSMStateConfig +from ..conv_ssm_state import ConvSSMState from .kda_pages_pool import KDAPagesPool -from ..transfer_utils import bounded_att_buffer_copies -from ..copy_kernel import copy_bounded_att_tensors - -if TYPE_CHECKING: - from lightllm.common.attention_state.types import PageKind - - -logger = init_logger(__name__) - - -class KDAState(BoundedAttentionState): - """KDA 的 Conv 历史、方形 SSM 和 MTP 游标，要求 Q/K/V head 数及维度相同。""" - - big_pages_pool: KDAPagesPool | None - small_pages_pool: KDAPagesPool | None - - # ======================================================================== - # 一、资源定义与初始化 - # ======================================================================== - - def __init__( - self, - config: ConvSSMStateConfig, - request_capacity: int, - mtp_step: int = 0, - device: str | torch.device = "cuda", - ) -> None: - """按请求容量创建 Conv/SSM 运行态和 MTP 游标；大小页池延后初始化。""" - assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" - self.config = config - self.state_name: str | None = None # 由 AttentionStateManager.register_state() 绑定时设置。 - self.request_capacity = request_capacity # 包含预留的 HOLD 请求槽位。 - self.mtp_step = mtp_step - self._logical_to_storage_layer = { - logical_layer_id: storage_layer_index - for storage_layer_index, logical_layer_id in enumerate(config.logical_layer_ids) - } - storage_layer_count = len(config.logical_layer_ids) - self.conv_state = torch.zeros( - (storage_layer_count, request_capacity, config.conv_channels, config.conv_width + mtp_step), - dtype=config.conv_dtype, - device=device, - ) - self.ssm_state = torch.zeros( - ( - storage_layer_count, - request_capacity * (mtp_step + 1), - config.num_v_heads, - config.key_dim, - config.value_dim, - ), - dtype=config.ssm_dtype, - device=device, - ) - self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None - self.big_pages_pool = None - self.small_pages_pool = None - - def get_state_type(self) -> str: - """返回 KDA 算法及 Conv/SSM 存储精度。""" - return f"kda_conv_{dtype_name(dtype=self.config.conv_dtype)}" f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" - - def initialize_page_pools( - self, - big_page_capacity: int, - small_page_capacity: int = 0, - device: str | torch.device | int = "cpu", - pin_memory: bool = True, - ) -> None: - """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" - assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" - assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" - self.big_pages_pool = ( - KDAPagesPool( - size=big_page_capacity, - config=self.config, - device=device, - pin_memory=pin_memory, - ) - if big_page_capacity - else None - ) - self.small_pages_pool = ( - KDAPagesPool( - size=small_page_capacity, - config=self.config, - device=device, - pin_memory=pin_memory, - ) - if small_page_capacity - else None - ) - - # ======================================================================== - # 二、Bounded attention：计算视图、运行态重置、大小页与 MTP 提交 - # ======================================================================== - - def contains_logical_layer(self, logical_layer_id: int) -> bool: - """判断指定逻辑层是否属于当前 State 的层映射，共享层同样返回 True。""" - return logical_layer_id in self._logical_to_storage_layer - - def get_layer_view(self, logical_layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: - """按逻辑层号映射到物理层，返回 Conv/SSM Tensor 引用；MTP 游标和步数直接从 State 读取。""" - storage_layer_index = self._logical_to_storage_layer[logical_layer_id] - return self.conv_state[storage_layer_index], self.ssm_state[storage_layer_index] - - def reset_bounded_att_state(self, req_indices: List[int]) -> None: - """清零指定请求的 Conv 历史、全部 SSM 候选及 MTP 游标，保留页池快照。""" - assert isinstance(req_indices, list), "request indices must be a list" - for req_idx in req_indices: - self.conv_state[:, req_idx].zero_() - start = req_idx * (self.mtp_step + 1) - self.ssm_state[:, start : start + self.mtp_step + 1].zero_() - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].fill_(0) - - def save_bounded_att_runtime_to_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """保存各请求指定绝对 token 长度处的基准状态；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance( - page_slot_ids, list - ), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = self._pool(page_kind) - # TODO：大小页仅在 prefill 阶段保存，mtp_state_index 始终为 0，可直接复制基准状态。 - # 若以后支持 decode/MTP 后保存，需要按已接受游标选择 Conv 窗口和 SSM 槽位。 - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - conv_copies.append((self.conv_state[:, req_idx, :, : self.config.conv_width], pool.conv_state[slot_id])) - ssm_copies.append((self.ssm_state[:, req_idx * (self.mtp_step + 1)], pool.ssm_state[slot_id])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - def restore_bounded_att_runtime_from_pages( - self, - req_indices: List[int], - state_token_lens: List[int], - page_kind: PageKind, - page_slot_ids: List[int], - ) -> None: - """恢复各请求指定绝对 token 长度处的基准状态并归零游标；Conv/SSM 不依赖长度定位。""" - assert isinstance(req_indices, list) and isinstance( - page_slot_ids, list - ), "request and slot indices must be lists" - assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" - pool = self._pool(page_kind) - conv_copies, ssm_copies = [], [] - for req_idx, slot_id in zip(req_indices, page_slot_ids): - conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, : self.config.conv_width])) - ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx * (self.mtp_step + 1)])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for req_idx in req_indices: - self.mtp_state_index[req_idx].fill_(0) - - def commit_bounded_att_state( - self, - req_indices: torch.Tensor, - accepted_steps: torch.Tensor, - ) -> None: - """用 CUDA 请求索引和 int32 接受步数选择 SSM 候选位置；步数包含基准输入步。""" - assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" - if self.mtp_state_index is not None: - self.mtp_state_index[req_indices] = accepted_steps - 1 - - # ======================================================================== - # 三、CPU Cache 数据交换 - # ======================================================================== - - @deepcopy_result - @lru_cache(maxsize=None) - def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回已接受 Conv/SSM 快照的字段区间与切片，供 CPU 页布局和容量计算使用。 - - region.shape 描述 CPU 页中的全局字段；axis_mappings 将本地页池的 TP 分片映射到该字段。 - 起止位置相对本 State，由管理器平移到 CPU 联合页；不包含 MTP 候选区或游标。 - """ - return self._get_transfer_regions() - - def write_state_pages_to_cpu( - self, - page_kind: PageKind, - page_slot_ids: List[int], - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - layout: StateTransferLayout, - cpu_page_ready_flags: list[bool] | None = None, - ) -> None: - """将指定页池槽位的快照写入 CPU 页。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - pool = self._pool(page_kind) - assert isinstance(page_slot_ids, list), "slot indices must be a list" - assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" - if cpu_page_ready_flags is not None: - assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" - page_fields, page_indices = [], [] - for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): - # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 - if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: - continue - page_fields.append(self._fields(pool.conv_state[slot], pool.ssm_state[slot])) - page_indices.append(page) - for copies in bounded_att_buffer_copies( - page_fields=page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=page_indices, - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - def read_state_pages_from_cpu( - self, - cpu_page_buffers: torch.Tensor, - cpu_page_indices: List[int], - big_page_slot_ids: List[int], - layout: StateTransferLayout, - ) -> None: - """读入指定大页槽位；请求运行态另行通过 restore_bounded_att_runtime_from_pages 恢复。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" - pool = self.big_pages_pool - assert isinstance(big_page_slot_ids, list), "slot indices must be a list" - assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" - assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" - # 本地目标：选定大页槽位的 Conv/SSM 快照，保留本 TP rank 的分片形状。 - pool_page_fields = [self._fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids] - # 外部来源：layout.regions 描述 CPU 全局字段和 TP 范围；write=False 将其读入本地页池。 - # 窗口保持快照顺序；请求 ring 的位置由后续 restore_bounded_att_runtime_from_pages 按长度恢复。 - for copies in bounded_att_buffer_copies( - page_fields=pool_page_fields, - page_buffers=cpu_page_buffers, - layout=layout, - page_indices=cpu_page_indices, - write=False, - ): - copy_bounded_att_tensors(copies=copies) - - # ======================================================================== - # 四、DP 迁移：大页快照与基准运行态 - # ======================================================================== - - def copy_big_pages_from_dp( - self, - src_page_ids: List[int], - dst_page_ids: List[int], - remote_dp_bounded_state: KDAState, - ) -> None: - """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" - assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" - assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" - assert ( - type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - src_pool = remote_dp_bounded_state.big_pages_pool - dst_pool = self.big_pages_pool - assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" - # 大页保存的是基准快照，与请求槽位及 MTP 候选区无关；页的引用和使用期限由调用方管理。 - # 页池均须在本进程注册；GPU 在当前流直接访问两端 CPU 存储，完成后才能归还页面。 - conv_copies, ssm_copies = [], [] - for src_page, dst_page in zip(src_page_ids, dst_page_ids): - conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) - ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - - def copy_runtime_state_from_dp( - self, - src_req_indices: List[int], - dst_req_indices: List[int], - remote_dp_bounded_state: KDAState, - ) -> None: - """复制同一 TP 分片的 prefill 基准运行态并重置目标游标，不写入页池。""" - assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" - assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" - assert ( - type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config - ), "DP copy requires the same algorithm and equivalent TP state configurations" - # TODO：当前 DP fetch 只复制 prefill 基准状态，源 MTP 游标为 0； - # 若以后支持 decode/MTP 后迁移，需按已接受游标选择源 Conv 窗口及 SSM 槽位。 - # 跨设备源存储必须已通过 peer/IPC 映射，批量复制不改变双方页池。 - conv_copies, ssm_copies = [], [] - for src_req, dst_req in zip(src_req_indices, dst_req_indices): - conv_copies.append( - ( - remote_dp_bounded_state.conv_state[:, src_req, :, : self.config.conv_width], - self.conv_state[:, dst_req, :, : self.config.conv_width], - ) - ) - ssm_copies.append( - ( - remote_dp_bounded_state.ssm_state[:, src_req * (remote_dp_bounded_state.mtp_step + 1)], - self.ssm_state[:, dst_req * (self.mtp_step + 1)], - ) - ) - copy_bounded_att_tensors(copies=conv_copies) - copy_bounded_att_tensors(copies=ssm_copies) - if self.mtp_state_index is not None: - for dst_req in dst_req_indices: - self.mtp_state_index[dst_req].fill_(0) - - # ======================================================================== - # 五、PD 分离传输 - # ======================================================================== - - @deepcopy_result - @lru_cache(maxsize=None) - def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: - """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" - return self._get_transfer_regions() - - def write_to_pd_page( - self, - req_idx: int, - pd_page_buffers: torch.Tensor, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """写出 state_token_len 处的 prefill 基准状态；Conv/SSM 不依赖长度，源 MTP 游标须为 0。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._runtime_fields(req_idx) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=True, - ): - copy_bounded_att_tensors(copies=copies) - - def read_from_pd_page( - self, - pd_page_buffers: torch.Tensor, - req_idx: int, - pd_page_index: int, - state_token_len: int, - layout: StateTransferLayout, - ) -> None: - """恢复 state_token_len 处的基准状态并规范化控制信息；Conv/SSM 不依赖长度定位。""" - assert layout.state_name == self.state_name, "transfer layout does not describe this resource" - assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" - fields = self._runtime_fields(req_idx) - for copies in bounded_att_buffer_copies( - page_fields=[fields], - page_buffers=pd_page_buffers, - layout=layout, - page_indices=[pd_page_index], - write=False, - ): - copy_bounded_att_tensors(copies=copies) - if self.mtp_state_index is not None: - self.mtp_state_index[req_idx].fill_(0) - - # ======================================================================== - # 六、私有辅助实现 - # ======================================================================== - - def _get_transfer_regions(self) -> tuple[TransferRegion, ...]: - """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" - # 阶段一：确定全局 Conv 形状；只传有效历史窗口，不包含 MTP 候选区和游标。 - config = self.config - storage_layer_count = len(config.logical_layer_ids) - conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim - conv_shape = (storage_layer_count, conv_channels, config.conv_width) - conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize - regions = [] - logger.debug( - "%s building transfer regions: TP=%d/%d, Conv shape=%s", - self.state_name, - config.tp_rank, - config.tp_world_size, - conv_shape, - ) - - # 阶段二：Conv 按 Q/K/V 拼接；三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 - local_channel_start = 0 - global_channel_start = 0 - for name, local_head_count, global_head_count, head_dim in ( - ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), - ("K", config.num_k_heads, config.global_k_heads, config.key_dim), - ("V", config.num_v_heads, config.global_v_heads, config.value_dim), - ): - # 配置已保证完整分片或副本。每组连续 rank 保存相同 heads，仅组内首 rank 写出。 - replica_count = config.tp_world_size * local_head_count // global_head_count - head_start = config.tp_rank // replica_count * local_head_count - write_enabled = config.tp_rank % replica_count == 0 - local_channel_count = local_head_count * head_dim - target_channel_start = global_channel_start + head_start * head_dim - conv_region = TransferRegion( - field="conv", - byte_start=0, - byte_end=conv_byte_end, - dtype=config.conv_dtype, - axes=("layer", "channel", "window"), - shape=conv_shape, - axis_mappings=( - AxisMapping( - axis="channel", - local_start=local_channel_start, - external_start=target_channel_start, - length=local_channel_count, - ), - AxisMapping( - axis="window", - local_start=0, - external_start=0, - length=config.conv_width, - ), - ), - write_enabled=write_enabled, - ) - regions.append(conv_region) - logger.debug( - "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " - "channels=%d, replicas=%d, write_enabled=%s", - self.state_name, - name, - local_channel_start, - target_channel_start, - local_channel_count, - replica_count, - write_enabled, - ) - local_channel_start += local_channel_count - global_channel_start += global_head_count * head_dim - - # 阶段三：SSM 沿 V head 切分；接在 Conv 之后，按 SSM dtype 对齐字节起点。 - replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads - head_start = config.tp_rank // replica_count * config.num_v_heads - write_enabled = config.tp_rank % replica_count == 0 - ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) - element_bytes = config.ssm_dtype.itemsize - ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes - ssm_region = TransferRegion( - field="ssm", - byte_start=ssm_byte_start, - byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, - dtype=config.ssm_dtype, - axes=("layer", "head", "key_dim", "value_dim"), - shape=ssm_shape, - axis_mappings=( - AxisMapping( - axis="head", - local_start=0, - external_start=head_start, - length=config.num_v_heads, - ), - AxisMapping( - axis="key_dim", - local_start=0, - external_start=0, - length=config.key_dim, - ), - AxisMapping( - axis="value_dim", - local_start=0, - external_start=0, - length=config.value_dim, - ), - ), - write_enabled=write_enabled, - ) - regions.append(ssm_region) - logger.debug( - "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", - self.state_name, - head_start, - config.num_v_heads, - write_enabled, - ssm_region.byte_start, - ssm_region.byte_end, - ) - return tuple(regions) - def _pool(self, page_kind: PageKind): - """根据大小页类型取得已启用的页池，不访问 allocator。""" - assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" - pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool - assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" - return pool - def _runtime_fields(self, req: int): - """返回请求基准位置的 Conv/SSM 视图，供恢复和传输使用。""" - # TODO：当前 DP/PD 不会发生在 MTP 提交后，源状态游标为 0。 - # 若以后支持该场景，需要按已接受游标选取 Conv 窗口及 SSM 槽位。 - return self._fields( - self.conv_state[:, req, :, : self.config.conv_width], self.ssm_state[:, req * (self.mtp_step + 1)] - ) +class KDAState(ConvSSMState): - @staticmethod - def _fields(conv: torch.Tensor, ssm: torch.Tensor): - return { - "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), - "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), - } + algorithm = "kda" + pages_pool_class = KDAPagesPool 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..5447ac3740 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,9 @@ 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, *, context=None + ) -> None: """ring 按绝对 token 位置寻址，无需提交额外游标。""" pass diff --git a/lightllm/common/attention_state/bounded_att/ssm/__init__.py b/lightllm/common/attention_state/bounded_att/ssm/__init__.py new file mode 100644 index 0000000000..a2d69e470c --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/__init__.py @@ -0,0 +1,34 @@ +import torch + + +def create_ssm_state(config, request_capacity, mtp_step, device, mode, replay_config, *, kda, gate_lower_bound): + if mode not in ("native", "compact", "replay"): + raise ValueError(f"unknown SSM state mode: {mode}") + slots = mtp_step + 1 if mode == "native" else 1 + state = torch.zeros( + (len(config.logical_layer_ids), request_capacity * slots, config.num_v_heads, config.key_dim, config.value_dim), + dtype=config.ssm_dtype, + device=device, + ) + kwargs = dict(kda=kda, gate_lower_bound=gate_lower_bound) + if mode == "native": + from .native import NativeSSMState + + return NativeSSMState(state, mtp_step + 1, **kwargs) + kwargs.update(num_key_heads=config.num_k_heads) + if mode == "compact": + from .compact import CompactSSMState + + return CompactSSMState(state, mtp_step + 1, config.conv_dtype, **kwargs) + from .replay import ReplaySSMState + + replay_config = replay_config or {} + return ReplaySSMState( + state, + replay_config.get("cache_len", 16), + mtp_step + 1, + config.conv_dtype, + compute_mode=replay_config.get("compute_mode", "inline"), + fold_method=replay_config.get("fold_method", "raw"), + **kwargs, + ) diff --git a/lightllm/common/attention_state/bounded_att/ssm/base.py b/lightllm/common/attention_state/bounded_att/ssm/base.py new file mode 100644 index 0000000000..ceb719edf0 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/base.py @@ -0,0 +1,66 @@ +from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update + + +def update_accepted_indices(reqs, accepted_steps, indices, context, *, replay_cursors=None): + if context is None: + indices[reqs] = accepted_steps - 1 + else: + linear_att_mtp_state_index_update( + req_to_mtp_state_index=indices, + b_req_mtp_start_loc=context.b_req_mtp_start_loc, + b_req_idx=context.b_req_idx, + b_mtp_index=context.b_mtp_index, + accepted_index=context.accepted_index, + verify_width=context.verify_width, + replay_cursors=replay_cursors, + ) + + +class SSMState: + slots_per_request = 1 + + def clear_history(self, req): + pass + + def prepare_decode(self, reqs, cu_seqlens=None): + pass + + def accept_single_token(self, reqs): + pass + + def merge_accepted_updates(self, reqs): + pass + + def decode_indices(self, reqs, speculative): + return reqs + + def get_accepted_state(self, req_idx, accepted_index=None): + return self.state[:, req_idx] + + def create_commit(self, run_config): + return None + + def validate_commit(self, context, state_name): + pass + + def commit(self, reqs, accepted_steps, indices, context, state_name): + update_accepted_indices(reqs, accepted_steps, indices, context) + + def forward( + self, + layer, + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens=None, + run_config=None, + *, + state_indices=None, + accepted_token_counts=None, + ): + raise NotImplementedError diff --git a/lightllm/common/attention_state/bounded_att/ssm/compact.py b/lightllm/common/attention_state/bounded_att/ssm/compact.py new file mode 100644 index 0000000000..7d570f93cd --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/compact.py @@ -0,0 +1,177 @@ +import torch +import triton + +from lightllm.common.ssm_types import SSMRunConfig +from lightllm.common.basemodel.triton_kernel.linear_att.compact_ssm import _compact +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config +from .base import SSMState, update_accepted_indices + + +class CompactSSMState(SSMState): + def __init__( + self, + state, + verify_width, + activation_dtype, + *, + num_key_heads=None, + kda=False, + gate_lower_bound=-5.0, + run_config=None, + ): + self.state = state + self.kda = kda + self.gate_lower_bound = gate_lower_bound + self.verify_width = verify_width + # Verify and accepted-state reconstruction must share the same layout. + self._config_is_fixed = run_config is not None + self.batch_configs = {} + self.run_config = dict(run_config or {"BV": 8, "num_warps": 1}) + assert self.run_config["BV"] in (8, 16, 32, 64, 128) + assert self.run_config["num_warps"] in (1, 2, 4, 8) + layers, slots, hv, k, v = state.shape + self.num_key_heads = hv if num_key_heads is None else num_key_heads + assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 + self.hold_req_id = slots - 1 + shape = (layers, slots, hv, verify_width) + self.keys = torch.empty( + (layers, slots, self.num_key_heads, verify_width, k), device=state.device, dtype=activation_dtype + ) + self.values = torch.empty((*shape, v), device=state.device, dtype=activation_dtype) + self.decays = torch.empty((*shape, k) if kda else shape, device=state.device, dtype=torch.float32) + self.betas = torch.empty(shape, device=state.device, dtype=torch.float32) + + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: + config = resolve_run_config(self, "compact", q, k, v, a, b, a_log, bias, cu_seqlens) + return SSMRunConfig(**config) + + def forward( + self, + linear_layer_idx, + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens=None, + run_config=None, + *, + state_indices=None, + accepted_token_counts=None, + ): + assert cu_seqlens is not None, "compact replay is only used for speculative verify" + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) + _, slots, hv, kd, vd = self.state.shape + assert self.num_key_heads in (q.shape[-2], hv) + out = torch.empty_like(v) + # Match native reduction layouts: changing BV can alter BF16 rounding. + bv = config.BV + _compact[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + self.state[linear_layer_idx], + self.keys[linear_layer_idx], + self.values[linear_layer_idx], + self.decays[linear_layer_idx], + self.betas[linear_layer_idx], + reqs, + cu_seqlens, + None, + out, + q.stride(1), + k.stride(1), + v.stride(1), + a.stride(0), + b.stride(0), + slots, + q.shape[-2], + hv, + self.num_key_heads, + kd, + vd, + self.verify_width, + self.hold_req_id, + triton.next_power_of_2(kd), + bv, + False, + self.kda, + self.gate_lower_bound, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + return out, config + + def accept_updates(self, reqs, accepted_state_indices, run_config=None): + """Merge this round's prefix; accepted_state_indices contains per-request last accepted indexes.""" + assert run_config is not None or not self.batch_configs, "compact acceptance requires the verify configuration" + layers, slots, hv, kd, vd = self.state.shape + config = run_config or self.run_config + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) + bv = config.BV + _compact[(triton.cdiv(vd, bv), reqs.numel(), layers * hv)]( + None, + None, + None, + None, + None, + None, + None, + self.state, + self.keys, + self.values, + self.decays, + self.betas, + reqs, + None, + accepted_state_indices, + None, + 0, + 0, + 0, + 0, + 0, + slots, + hv, + hv, + self.num_key_heads, + kd, + vd, + self.verify_width, + self.hold_req_id, + triton.next_power_of_2(kd), + bv, + True, + self.kda, + self.gate_lower_bound, + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + + def create_commit(self, run_config): + if not isinstance(run_config, SSMRunConfig): + raise ValueError("compact commit requires the configuration of its forward") + return run_config + + def validate_commit(self, context, state_name): + if context is None or not context.batches: + raise ValueError("compact acceptance requires forward commit batches") + for _, config in context.partitions(state_name): + if not isinstance(config, SSMRunConfig): + raise ValueError(f"missing compact forward configuration for {state_name}") + + def commit(self, reqs, accepted_steps, indices, context, state_name): + self.validate_commit(context, state_name) + update_accepted_indices(reqs, accepted_steps, indices, context) + for partition, config in context.partitions(state_name): + self.accept_updates(reqs[partition], indices, run_config=config) diff --git a/lightllm/common/attention_state/bounded_att/ssm/forward.py b/lightllm/common/attention_state/bounded_att/ssm/forward.py new file mode 100644 index 0000000000..b827149a36 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/forward.py @@ -0,0 +1,64 @@ +from dataclasses import dataclass +from typing import Optional, TYPE_CHECKING + +import torch + +from lightllm.common.attention_state.commit import StateCommit +from lightllm.common.ssm_types import SSMRunConfig + +if TYPE_CHECKING: + from ..conv_ssm_state import ConvSSMState + + +@dataclass +class SSMForward: + resource: "ConvSSMState" + req_indices: torch.Tensor + state_indices: torch.Tensor + cu_seqlens: Optional[torch.Tensor] + accepted_token_counts: Optional[torch.Tensor] + run_config: Optional[SSMRunConfig] = None + + def run_layer(self, logical_layer, q, k, v, a, b, a_log, bias): + output, self.run_config = self.resource._ssm.forward( + self.resource.get_storage_layer_index(logical_layer), + q, + k, + v, + a, + b, + a_log, + bias, + self.req_indices, + self.cu_seqlens, + run_config=self.run_config, + state_indices=self.state_indices, + accepted_token_counts=self.accepted_token_counts, + ) + return output + + def finish(self): + payload = self.resource._ssm.create_commit(self.run_config) + if payload is None: + return None + if self.resource.state_name is None: + raise ValueError("forward commits require a registered resource") + return StateCommit(self.resource.state_name, payload) + + def copy_inputs(self, source: "SSMForward"): + if self.resource is not source.resource: + raise ValueError("cannot copy a forward for a different resource") + copies = [] + for name in ("req_indices", "state_indices", "cu_seqlens", "accepted_token_counts"): + target, value = getattr(self, name), getattr(source, name) + if (target is None) != (value is None): + raise ValueError(f"captured SSM layout changed: {name}") + if target is not None: + if target.shape != value.shape or target.dtype != value.dtype or target.device != value.device: + raise ValueError(f"captured SSM tensor layout changed: {name}") + copies.append((target, value)) + copied = set() + for target, value in copies: + if target is not None and target.data_ptr() != value.data_ptr() and target.data_ptr() not in copied: + target.copy_(value, non_blocking=True) + copied.add(target.data_ptr()) diff --git a/lightllm/common/attention_state/bounded_att/ssm/native.py b/lightllm/common/attention_state/bounded_att/ssm/native.py new file mode 100644 index 0000000000..b090a3131d --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/native.py @@ -0,0 +1,96 @@ +import torch + +from .base import SSMState + + +class NativeSSMState(SSMState): + def __init__(self, state, verify_width, *, kda=False, gate_lower_bound=-5.0): + self.state = state + self.slots_per_request = verify_width + self.kda = kda + self.gate_lower_bound = gate_lower_bound + + def decode_indices(self, reqs, speculative): + if not speculative: + return reqs if self.slots_per_request == 1 else reqs * self.slots_per_request + offsets = torch.arange(self.slots_per_request, device=reqs.device, dtype=reqs.dtype) + return reqs[:, None] * self.slots_per_request + offsets[None, :] + + def get_accepted_state(self, req_idx, accepted_index=None): + if self.slots_per_request == 1: + return self.state[:, req_idx] + index = req_idx * self.slots_per_request + accepted_index + return self.state.index_select(1, index.long().view(1))[:, 0] + + def forward( + self, + layer, + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens=None, + run_config=None, + *, + state_indices=None, + accepted_token_counts=None, + ): + state = self.state[layer] + if self.kda: + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda + + output, _ = fused_recurrent_kda( + q=q, + k=k, + v=v, + raw_gate=a.view(*q.shape[:2], -1), + raw_beta=b.view(*q.shape[:2], -1), + a_log=a_log, + gate_bias=bias, + initial_state=state, + lower_bound=self.gate_lower_bound, + inplace_final_state=True, + ssm_state_indices=state_indices, + cu_seqlens=cu_seqlens, + num_accepted_tokens=accepted_token_counts, + ) + elif cu_seqlens is not None: + from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, + ) + + output, _ = mtp_fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=state, + cu_seqlens=cu_seqlens.to(torch.long), + ssm_state_indices=state_indices, + ssm_state_write_indices=state_indices, + num_accepted_tokens=accepted_token_counts, + A_log=a_log, + dt_bias=bias, + a_raw=a, + b_raw=b, + ) + else: + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import fused_recurrent_gated_delta_rule + + output, _ = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=state, + inplace_final_state=True, + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=bias, + a_raw=a, + b_raw=b, + ) + return output, None diff --git a/lightllm/common/attention_state/bounded_att/ssm/replay.py b/lightllm/common/attention_state/bounded_att/ssm/replay.py new file mode 100644 index 0000000000..ce3f62f69f --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/ssm/replay.py @@ -0,0 +1,300 @@ +import torch +import triton + +from lightllm.common.ssm_types import SSMRunConfig +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( + _accept_updates, + _prepare_fold, + _fold_history, + _materialize, + _replay, +) +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config +from .base import SSMState, update_accepted_indices + + +class ReplaySSMState(SSMState): + """Request-owned scratch; canonical CPU/PD checkpoints never include it.""" + + def __init__( + self, + state, + capacity, + verify_width, + activation_dtype=torch.bfloat16, + *, + num_key_heads=None, + kda=False, + gate_lower_bound=-5.0, + compute_mode="inline", + run_config=None, + fold_method="raw", + ): + assert state.dtype in (torch.float32, torch.bfloat16), "ReplaySSM 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", "parallel") + assert fold_method in ("raw", "derived") + assert fold_method == "raw" or compute_mode == "parallel" + if compute_mode == "parallel": + assert not kda and verify_width > 1, "Parallel ReplaySSM requires GDN MTP" + assert state.shape[-2:] == (128, 128), "Parallel ReplaySSM requires 128x128 heads" + assert activation_dtype == torch.bfloat16, "Parallel ReplaySSM requires BF16 activations" + assert state.dtype == torch.bfloat16, "Parallel ReplaySSM currently supports BF16 state" + self.fold_method = fold_method + self.state = state + self.kda = kda + self.gate_lower_bound = gate_lower_bound + self.fold_programs = ( + 8 * torch.cuda.get_device_properties(state.device).multi_processor_count if state.is_cuda else 0 + ) + self.capacity = capacity + self.verify_width = verify_width + self.compute_mode = compute_mode + self.run_config = run_config + self._config_is_fixed = run_config is not None + self.batch_configs = {} + layers, slots, hv, k, v = state.shape + # Omission preserves the redundant layout for existing callers and A/B. + self.num_key_heads = hv if num_key_heads is None else num_key_heads + assert self.num_key_heads > 0 and hv % self.num_key_heads == 0 + self.hold_req_id = slots - 1 + self.cursors = torch.zeros(slots, dtype=torch.int32, device=state.device) + self.forward_cursors = torch.zeros_like(self.cursors) + # Alternate halves on fold so CTAs cannot overwrite history another + # V tile is still reading. Cursor packs phase (2*L) and count (0..L). + # Derived corrections can exceed the activation/state range even when + # the recurrence stays finite. Only BF16 has enough exponent range for + # the short-history narrowing optimization; FP16 must retain FP32. + history_dtype = ( + activation_dtype + if activation_dtype == torch.bfloat16 + and verify_width > 1 + and compute_mode == "precompute" + and capacity <= 8 + else torch.float32 + ) + key_shape = (layers, slots, self.num_key_heads, 2 * capacity, k) + self.keys = torch.empty(key_shape, dtype=history_dtype, device=state.device) + self.deltas = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=history_dtype, device=state.device) + self.gates = torch.empty( + (layers, slots, hv, 2 * capacity, k) if kda else (layers, slots, hv, 2 * capacity), + dtype=torch.float32, + device=state.device, + ) + self.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=state.device) + self.raw_values = torch.empty((layers, slots, hv, 2 * capacity, v), dtype=activation_dtype, device=state.device) + self.betas = torch.empty((layers, slots, hv, 2 * capacity), dtype=torch.float32, device=state.device) + + def clear_history(self, req): + """Invalidate history for a reused/restored request; leave its SSM state intact.""" + if type(req) is int: + # Scalar assignment copies a CPU tensor into the CUDA slot. + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs, cu_seqlens=None): + """Snapshot history for all layers and prepare the base for accepted updates.""" + positions = torch.empty_like(reqs) + layers, _, hv, kd, vd = self.state.shape + active = torch.empty((reqs.numel() + 1,), dtype=torch.int32, device=reqs.device) + _prepare_fold[(1,)]( + reqs, + self.cursors, + positions, + self.forward_cursors, + cu_seqlens, + active, + reqs.numel(), + self.hold_req_id, + self.capacity, + cu_seqlens is not None, + triton.next_power_of_2(reqs.numel()), + self.verify_width, + self.compute_mode == "parallel", + num_warps=4, + ) + programs = min(self.fold_programs, triton.cdiv(vd, 32) * reqs.numel() * layers * hv) + _fold_history[(programs,)]( + self.state, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + reqs, + positions, + active, + reqs.numel(), + layers, + hv, + self.num_key_heads, + kd, + vd, + self.capacity, + triton.next_power_of_2(kd), + 32, + self.state.stride(0), + self.raw_keys.stride(0), + self.raw_values.stride(0), + self.gates.stride(0), + self.betas.stride(0), + self.kda, + self.keys, + self.deltas, + self.fold_method == "derived", + BL=max(16, self.capacity), + num_warps=2, + num_stages=1, + ) + return positions + + def accept_updates(self, reqs, accepted_state_indices=None): + """Accept this round's prefix; accepted_state_indices contains per-request last accepted indexes.""" + _accept_updates[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width, + 256, + ) + + def merge_accepted_updates(self, reqs): + """Fold accepted history into SSM state and clear its history positions.""" + if reqs.numel() == 0: + return + self._materialize_accepted_state(reqs, self.state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def _materialize_accepted_state(self, reqs, output, snapshot): + layers, slots, hv, k, v = self.state.shape + config = self.run_config or {"BV": 32, "num_warps": 4} + bv = config["BV"] + _materialize[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + self.state, + output, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + self.cursors, + reqs, + slots, + hv, + self.num_key_heads, + k, + v, + self.capacity, + self.hold_req_id, + triton.next_power_of_2(k), + bv, + snapshot, + self.kda, + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), + ) + + def get_accepted_state(self, req_idx, accepted_index=None): + """Return a canonical checkpoint without changing the active state or history.""" + layers, _, hv, k, v = self.state.shape + output = torch.empty((layers, hv, k, v), dtype=self.state.dtype, device=self.state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.state.device) + self._materialize_accepted_state(reqs, output, snapshot=True) + return output + + def get_run_config(self, q, k, v, a, b, a_log, bias, cu_seqlens=None) -> SSMRunConfig: + if self.compute_mode == "parallel": + config = self.run_config or {"BV": 128, "num_warps": 2, "num_stages": 1} + else: + config = resolve_run_config(self, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + return SSMRunConfig(**config) + + def forward( + self, + linear_layer_idx, + q, + k, + v, + a, + b, + a_log, + bias, + reqs, + cu_seqlens=None, + run_config=None, + *, + state_indices=None, + accepted_token_counts=None, + ): + config = run_config or self.get_run_config(q, k, v, a, b, a_log, bias, cu_seqlens) + if not isinstance(config, SSMRunConfig): + config = SSMRunConfig(**config) + hv, kd, vd = self.state.shape[-3:] + assert self.num_key_heads in (q.shape[-2], hv) + axis = 1 if cu_seqlens is not None else 0 + out = torch.empty_like(v) + bv = config.BV + precompute_projections = self.verify_width > 1 and bv >= 16 and self.compute_mode in ("precompute", "parallel") + if self.compute_mode == "parallel": + assert precompute_projections, "Parallel ReplaySSM requires BV >= 16" + # Resolve tile dimensions in Python: bare max() in Triton JIT code + # lowers to a tensor operation, which cannot define static shapes. + bw = triton.next_power_of_2(self.verify_width) + _replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + self.state[linear_layer_idx], + self.keys[linear_layer_idx], + self.deltas[linear_layer_idx], + self.gates[linear_layer_idx], + self.raw_keys[linear_layer_idx], + self.raw_values[linear_layer_idx], + self.betas[linear_layer_idx], + reqs, + self.forward_cursors, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + self.num_key_heads, + kd, + vd, + self.capacity, + self.hold_req_id, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + precompute_projections, + bw, + reqs.numel() > 12 and self.compute_mode != "parallel", + self.kda, + self.gate_lower_bound, + self.compute_mode == "parallel", + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, self.capacity + bw)), + num_warps=config.num_warps, + num_stages=config.num_stages, + ) + return out, config + + def accept_single_token(self, reqs): + + self.accept_updates(reqs) + + def commit(self, reqs, accepted_steps, indices, context, state_name): + if context is not None and self.verify_width > 1: + update_accepted_indices(reqs, accepted_steps, indices, context, replay_cursors=self.cursors) + else: + update_accepted_indices(reqs, accepted_steps, indices, context) + self.accept_updates(reqs, indices) 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..2b937bf548 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,9 @@ 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, *, context=None + ) -> None: """按绝对位置寻址；接受后继续覆盖 ring，不维护额外 MTP 游标。""" pass diff --git a/lightllm/common/attention_state/commit.py b/lightllm/common/attention_state/commit.py new file mode 100644 index 0000000000..778e154eb5 --- /dev/null +++ b/lightllm/common/attention_state/commit.py @@ -0,0 +1,59 @@ +from dataclasses import dataclass +from typing import Tuple + +import torch + + +@dataclass(frozen=True) +class StateCommit: + state_name: str + payload: object + + +@dataclass(frozen=True) +class StateCommitBatch: + + request_count: int + commits: Tuple[StateCommit, ...] = () + + def __post_init__(self): + if type(self.request_count) is not int or self.request_count < 0: + raise ValueError("commit batch requires a nonnegative request count") + names = [commit.state_name for commit in self.commits] + if len(set(names)) != len(names): + raise ValueError("duplicate state commit in a forward") + + +@dataclass(frozen=True) +class MtpCommitContext: + + b_req_mtp_start_loc: torch.Tensor + b_req_idx: torch.Tensor + b_mtp_index: torch.Tensor + accepted_index: torch.Tensor + verify_width: int + batches: Tuple[StateCommitBatch, ...] = () + + def validate(self, request_count: int, state_names=None): + if type(self.verify_width) is not int or self.verify_width <= 0: + raise ValueError("verification width must be positive") + rows = (self.b_req_idx, self.b_mtp_index, self.accepted_index) + if any(row.ndim != 1 or row.numel() != self.b_req_idx.numel() for row in rows): + raise ValueError("verification row tensors must have the same one-dimensional layout") + if self.b_req_mtp_start_loc.numel() != request_count: + raise ValueError("verification and logical request counts differ") + if self.batches and sum(batch.request_count for batch in self.batches) != request_count: + raise ValueError("commit batches must partition all unpadded requests") + for batch in self.batches: + for commit in batch.commits: + if state_names is not None and commit.state_name not in state_names: + raise ValueError(f"commit belongs to an unknown state: {commit.state_name}") + + def partitions(self, state_name: str): + start = 0 + for batch in self.batches: + end = start + batch.request_count + if end > start: + payload = next((c.payload for c in batch.commits if c.state_name == state_name), None) + yield slice(start, end), payload + start = end diff --git a/lightllm/common/attention_state/manager.py b/lightllm/common/attention_state/manager.py index 6fb18c7ea1..835778a844 100644 --- a/lightllm/common/attention_state/manager.py +++ b/lightllm/common/attention_state/manager.py @@ -397,12 +397,19 @@ def commit_bounded_att_state( req_indices: torch.Tensor, accepted_steps: torch.Tensor, role: str, + *, + context=None, ) -> None: """传入 CUDA 请求索引和 int32 接受步数，包含基准输入步；无参与 State 时为空操作。""" assert req_indices.is_cuda and accepted_steps.is_cuda, "commit inputs must be CUDA tensors" assert role in MODEL_ROLES, f"unknown model role {role}" - for state in self.get_bounded_attention_states(role): - state.commit_bounded_att_state(req_indices, accepted_steps) + states = self.get_bounded_attention_states(role) + if context is not None: + context.validate(req_indices.numel(), {state.state_name for state in states}) + for state in states: + state.validate_bounded_att_commit(context) + for state in states: + state.commit_bounded_att_state(req_indices, accepted_steps, context=context) # ======================================================================== # 四、CPU/DP/PD 操作均由对应 Operator 提供。 diff --git a/lightllm/common/basemodel/attention/base_att.py b/lightllm/common/basemodel/attention/base_att.py index 1122df3551..5e08a1d711 100644 --- a/lightllm/common/basemodel/attention/base_att.py +++ b/lightllm/common/basemodel/attention/base_att.py @@ -10,6 +10,7 @@ from lightllm.utils.envs_utils import get_env_start_args if TYPE_CHECKING: + from lightllm.common.attention_state.commit import StateCommit from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.infer_struct import InferStateInfo @@ -161,6 +162,16 @@ def prefill_att( @dataclass class BaseDecodeAttState(BaseAttState): + def get_state_commits(self) -> tuple["StateCommit", ...]: + return () + + def clone_for_warmup(self, infer_state): + import copy + + state = copy.copy(self) + state._infer_state_weak_ref = weakref.ref(infer_state) + return state + def copy_for_decode_cuda_graph(self, new_state: "BaseDecodeAttState"): """Copy tensor metadata into the buffers owned by the captured graph.""" for attr_name, attr_value in vars(new_state).items(): diff --git a/lightllm/common/basemodel/attention/linear/flashqla.py b/lightllm/common/basemodel/attention/linear/flashqla.py index 6e57eca732..79ad4e20e6 100644 --- a/lightllm/common/basemodel/attention/linear/flashqla.py +++ b/lightllm/common/basemodel/attention/linear/flashqla.py @@ -1,7 +1,7 @@ -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend +from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttBackend -class FlashQlaLinearAttBackend(LinearAttBackend): +class FlashQlaLinearAttBackend(GDNLinearAttBackend): def get_prefill_kernel(self): from flash_qla import chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 986214568f..0e44139037 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -1,23 +1,17 @@ from lightllm.common.attention_state.bounded_att import GDNState +from .state import LinearDecodeAttState, LinearPrefillAttState, get_resource import dataclasses import weakref -import torch from abc import ABC, abstractmethod from typing import TYPE_CHECKING -from ..base_att import BaseAttBackend, BasePrefillAttState, BaseDecodeAttState, AttControl + +import torch + +from ..base_att import BaseAttBackend, AttControl from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn 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_fused_recurrent import ( - mtp_fused_recurrent_gated_delta_rule, -) -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.fla.ops import ( - fused_recurrent_gated_delta_rule, -) if TYPE_CHECKING: from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -26,7 +20,7 @@ from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerWeight -class LinearAttBackend(BaseAttBackend, ABC): +class GDNLinearAttBackend(BaseAttBackend, ABC): def __init__(self, model: "TpPartBaseModel"): super().__init__(model=model) self._init_linear_layer_metadata(network_config=model.config, tp_world_size=model.tp_world_size_) @@ -100,28 +94,16 @@ def _rearrange_mixed_qkv(self, mixed_qkv, decode=False): value = value.view(1, seq_len, self.tp_num_v_heads, self.head_v_dim) return query, key, value - def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "LinearAttPrefillAttState": - return LinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + def create_att_prefill_state(self, infer_state: "InferStateInfo") -> "GDNLinearAttPrefillAttState": + return GDNLinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) - def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDecodeAttState": - return LinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + def create_att_decode_state(self, infer_state: "InferStateInfo") -> "GDNLinearAttDecodeAttState": + return GDNLinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) @dataclasses.dataclass -class LinearAttPrefillAttState(BasePrefillAttState): - - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - backend: LinearAttBackend = self.backend - mtp_step = backend.mtp_step - # 每次 _prefill 都会在 runtime infer_state 上调用 init_state。 - # prefill cuda graph 回调必须走 new_infer_state.prefill_att_state1， - # 才能读到这里按当前 batch（含 token padding 后的 dummy request）更新的索引。 - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (mtp_step + 1) - return +class GDNLinearAttPrefillAttState(LinearPrefillAttState): + state_class = GDNState def prefill_att( self, @@ -138,22 +120,24 @@ def prefill_att( mixed_qkvzba: torch.Tensor = linear_att_dict["mixed_qkvzba"] layer_weight = linear_att_dict["layer_weight"] layer_num = linear_att_dict["layer_num"] - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=GDNState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) + resource = get_resource(self, layer_num) + conv_states, ssm_states = resource.get_layer_view(layer_num) # 在开启了mtp的时候，conv 状态的最后一维可能存在冗余的部分，需要进行切片对齐。 # prefill 模式下，使用不到这几个维度，所以需要扣除掉， if backend.mtp_step > 0: conv_states = conv_states[:, :, : -backend.mtp_step] mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) core_attn_out = self._gdn_prefill_kernel( - mixed_qkv, conv_states, ssm_states, a, b, self.infer_state, layer_weight + mixed_qkv, + conv_states, + ssm_states, + a, + b, + self.infer_state, + layer_weight, + self.state_indices[resource.state_name], ) return core_attn_out, z @@ -166,17 +150,19 @@ def _gdn_prefill_kernel( b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", + state_indices, ): g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) mixed_qkv = mixed_qkv.transpose(0, 1) - backend: LinearAttBackend = self.backend + conv_indices, ssm_indices = state_indices + backend: GDNLinearAttBackend = self.backend out_tensor = causal_conv1d_fn( mixed_qkv, layer_weight.linear_conv1d.mm_param.weight, bias=layer_weight.linear_conv1d.bias, query_start_loc=infer_state.b1_cu_q_seq_len, - cache_indices=self.b_conv_buffer_idx, + cache_indices=conv_indices, has_initial_state=infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation=backend.activation, @@ -185,7 +171,7 @@ def _gdn_prefill_kernel( # Recurrent processing query, key, value = backend._rearrange_mixed_qkv(mixed_qkv) - initial_state = ssm_states[self.b_ssm_buffer_idx] + initial_state = ssm_states[ssm_indices] # g and beta have shape (total_tokens, num_heads), need to unsqueeze to get (1, total_tokens, num_heads) core_attn_out, last_recurrent_state = backend.prefill_kernel( q=query, @@ -201,85 +187,16 @@ def _gdn_prefill_kernel( # The chunk kernel accumulates the recurrent state in float32 even when # the state cache is configured as bfloat16. Advanced indexing does # not perform an implicit dtype conversion for index_put. - ssm_states[self.b_ssm_buffer_idx] = last_recurrent_state.to(ssm_states.dtype, copy=False) + ssm_states[ssm_indices] = last_recurrent_state.to(ssm_states.dtype, copy=False) return core_attn_out @dataclasses.dataclass -class LinearAttDecodeAttState(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_num_accepted_tokens: torch.Tensor = None +class GDNLinearAttDecodeAttState(LinearDecodeAttState): + state_class = GDNState - def init_state(self): - draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) - if draft_step == 0: - self._init_normal_decode_state() - return - - # MTP 元数据按模型初始化，沿用每个角色只有一个 GDN State 的约束。 - # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 - (state,) = ( - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states( - self.infer_state.att_state_role - ) - if isinstance(state, GDNState) - ) - if self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(draft_step + 1, req_to_mtp_state_index=state.mtp_state_index) - else: - self._init_fixed_mtp_decode_state(draft_step, req_to_mtp_state_index=state.mtp_state_index) - - def _init_normal_decode_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx - - def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - ( - self.b1_mtp_cu_q_seq_len, - self.b_conv_buffer_idx, - self.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=self.infer_state.b_req_idx, - b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=req_to_mtp_state_index, - hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_fixed_mtp_decode_state(self, draft_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - mtp_size = draft_step + 1 - batch_size = self.infer_state.batch_size - assert batch_size % mtp_size == 0, ( - "GDN fixed-layout decode requires batch_size to be divisible by draft_step + 1, " - f"got batch_size={batch_size}, draft_step={draft_step}." - ) - - att_batch_size = batch_size // mtp_size - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch_size + 1, - mtp_size, - dtype=torch.int32, - device=self.infer_state.b_req_idx.device, - ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - att_batch_size = self.b_conv_buffer_idx.shape[0] - # Each request owns mtp_size consecutive recurrent-state slots. - b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) - state_offsets = torch.arange( - mtp_size, - device=self.infer_state.b_req_idx.device, - dtype=self.infer_state.b_req_idx.dtype, - ).view(1, mtp_size) - self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] + def get_draft_step(self): + return self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) def decode_att( self, @@ -296,27 +213,23 @@ def decode_att( mixed_qkvzba = linear_att_dict["mixed_qkvzba"] layer_weight = linear_att_dict["layer_weight"] layer_num = linear_att_dict["layer_num"] - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend mixed_qkv, z, b, a = backend._split_qkvzba(mixed_qkvzba) - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="gdn", - logical_layer_id=layer_num, - state_class=GDNState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) + forward = self.get_forward(layer_num) + conv_states, _ = forward.resource.get_layer_view(layer_num) draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) if draft_step > 0: core_attn_out = self._gdn_mtp_kernel( mixed_qkv, conv_states, - ssm_states, + forward, a, b, self.infer_state, layer_weight, + layer_num=layer_num, ) else: # 非 MTP 模式下，使用线性层 decode 状态。 @@ -324,11 +237,12 @@ def decode_att( mixed_qkv, z, conv_states, - ssm_states, + forward, a, b, self.infer_state, layer_weight, + layer_num=layer_num, ) return core_attn_out, z @@ -337,13 +251,15 @@ def _gdn_decode_kernel( mixed_qkv: torch.Tensor, z: torch.Tensor, conv_states: torch.Tensor, - ssm_states: torch.Tensor, + forward, a: torch.Tensor, b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", + *, + layer_num: int, ): - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend # Recurrent processing with fused gating. Decode uses a specialized # conv+pack kernel to avoid materializing the post-conv qkv tensor @@ -356,7 +272,7 @@ def _gdn_decode_kernel( conv_states, layer_weight.linear_conv1d.mm_param.weight, layer_weight.linear_conv1d.bias, - self.b_conv_buffer_idx, + forward.req_indices, backend.activation, backend.conv_kernel_dim, backend.tp_num_k_heads, @@ -364,18 +280,15 @@ def _gdn_decode_kernel( backend.tp_num_v_heads, backend.head_v_dim, ) - core_attn_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - inplace_final_state=True, - ssm_state_indices=self.b_ssm_buffer_idx, - use_qk_l2norm_in_kernel=True, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, + core_attn_out = forward.run_layer( + layer_num, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, ) return core_attn_out, z @@ -383,19 +296,21 @@ def _gdn_mtp_kernel( self, mixed_qkv: torch.Tensor, conv_states: torch.Tensor, - ssm_states: torch.Tensor, + forward, a: torch.Tensor, b: torch.Tensor, infer_state: "Qwen3NextInferStateInfo", layer_weight: "Qwen3NextTransformerLayerWeight", + *, + layer_num: int, ): from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( causal_conv1d_update as causal_conv1d_update_mtp, ) - backend: LinearAttBackend = self.backend + backend: GDNLinearAttBackend = self.backend - cu_seqlens_q = self.b1_mtp_cu_q_seq_len + cu_seqlens_q = forward.cu_seqlens mixed_qkv = causal_conv1d_update_mtp( mixed_qkv, conv_states, @@ -403,28 +318,20 @@ def _gdn_mtp_kernel( mtp_step=backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model), 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, + conv_state_indices=forward.req_indices, + num_accepted_tokens=forward.accepted_token_counts, query_start_loc=cu_seqlens_q, ) query, key, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=False) - assert self.b_ssm_buffer_idx.dim() == 2, "SSM buffer idx must be 2D [N, S+1]" - # #8b: b_num_accepted_tokens >= 1 is guaranteed upstream: init/cache restore set 1, - # and MTP decode only writes values in [1, mtp_step+1]. The old per-layer per-step - # .all() D2H sync stalled the GPU on the eager decode hot path; it is redundant here. - core_attn_out, _ = mtp_fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, - initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), - ssm_state_indices=self.b_ssm_buffer_idx, - ssm_state_write_indices=self.b_ssm_buffer_idx, - num_accepted_tokens=self.b_num_accepted_tokens, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, + core_attn_out = forward.run_layer( + layer_num, + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, ) return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/kda.py b/lightllm/common/basemodel/attention/linear/kda.py index cd2f8afb39..a985a6bc13 100644 --- a/lightllm/common/basemodel/attention/linear/kda.py +++ b/lightllm/common/basemodel/attention/linear/kda.py @@ -11,24 +11,19 @@ import torch from lightllm.common.attention_state.bounded_att import KDAState +from .state import LinearDecodeAttState, LinearPrefillAttState, get_resource 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.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda 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.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) if TYPE_CHECKING: from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -49,23 +44,18 @@ def __init__(self, model: "TpPartBaseModel"): self.mtp_step = model.args.mtp_step def create_att_prefill_state(self, infer_state: "InferStateInfo"): - return KDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + return KDALinearAttPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) def create_att_decode_state(self, infer_state: "InferStateInfo"): - return KDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + return KDALinearAttDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) def split_qkv(self, mixed_qkv: torch.Tensor): return mixed_qkv.split(self.tp_hidden_size, dim=-1) @dataclasses.dataclass -class KDAPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None - b_ssm_buffer_idx: torch.Tensor = None - - def init_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx * (self.backend.mtp_step + 1) +class KDALinearAttPrefillAttState(LinearPrefillAttState): + state_class = KDAState def prefill_att( self, @@ -84,12 +74,8 @@ def prefill_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=KDAState, - ) + state = get_resource(self, layer_num) + conv_indices, ssm_indices = self.state_indices[state.state_name] 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( @@ -97,7 +83,7 @@ def prefill_att( 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, + cache_indices=conv_indices, has_initial_state=self.infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation="silu", @@ -110,7 +96,7 @@ def prefill_att( raw_gate = raw_gate.view(1, -1, backend.tp_hidden_size) raw_beta = raw_beta.view(1, -1, backend.tp_num_heads) - initial_state = ssm_states[self.b_ssm_buffer_idx].contiguous() + initial_state = ssm_states[ssm_indices].contiguous() output, final_state = chunk_kda_with_fused_gate( q=q, k=k, @@ -126,84 +112,13 @@ def prefill_att( safe_gate=True, lower_bound=backend.lower_bound, ) - ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + ssm_states[ssm_indices] = final_state.to(ssm_states.dtype, copy=False) return output @dataclasses.dataclass -class KDADecodeAttState(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_num_accepted_tokens: torch.Tensor = None - - def init_state(self): - mtp_step = self.backend.mtp_step - if mtp_step == 0: - self._init_normal_decode_state() - return - - # MTP 元数据按模型初始化，沿用每个角色只有一个 KDA State 的约束。 - # 只在此读取游标；Backend 不保存 State，逐层计算按逻辑层号查询资源。 - (state,) = ( - state - for state in self.infer_state.att_state_manager.get_bounded_attention_states( - self.infer_state.att_state_role - ) - if isinstance(state, KDAState) - ) - if self.backend.uses_dynamic_spec_verify_layout(): - self._init_dynamic_mtp_decode_state(mtp_step + 1, req_to_mtp_state_index=state.mtp_state_index) - else: - self._init_fixed_mtp_decode_state(mtp_step, req_to_mtp_state_index=state.mtp_state_index) - - def _init_normal_decode_state(self): - self.b_conv_buffer_idx = self.infer_state.b_req_idx - self.b_ssm_buffer_idx = self.infer_state.b_req_idx - - def _init_dynamic_mtp_decode_state(self, mtp_size: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - ( - self.b1_mtp_cu_q_seq_len, - self.b_conv_buffer_idx, - self.b_num_accepted_tokens, - ) = build_dynamic_mtp_linear_att_state_params( - b_req_idx=self.infer_state.b_req_idx, - b_mtp_index=self.infer_state.b_mtp_index, - req_to_mtp_state_index=req_to_mtp_state_index, - hold_req_id=self.infer_state.req_manager.HOLD_REQUEST_ID, - ) - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_fixed_mtp_decode_state(self, mtp_step: int, *, req_to_mtp_state_index: torch.Tensor) -> None: - mtp_size = mtp_step + 1 - batch_size = self.infer_state.batch_size - assert batch_size % mtp_size == 0, ( - "KDA fixed-layout decode requires batch_size to be divisible by mtp_step + 1, " - f"got batch_size={batch_size}, mtp_step={mtp_step}." - ) - - att_batch_size = batch_size // mtp_size - self.b1_mtp_cu_q_seq_len = torch.arange( - 0, - batch_size + 1, - mtp_size, - dtype=torch.int32, - device=self.infer_state.b_req_idx.device, - ) - self.b_conv_buffer_idx = self.infer_state.b_req_idx.view(att_batch_size, mtp_size)[:, 0].contiguous() - self.b_num_accepted_tokens = req_to_mtp_state_index[self.b_conv_buffer_idx] + 1 - self._init_mtp_ssm_buffer_idx(mtp_size) - - def _init_mtp_ssm_buffer_idx(self, mtp_size: int): - att_batch_size = self.b_conv_buffer_idx.shape[0] - # Each request owns mtp_size consecutive recurrent-state slots. - b_ssm_buffer_start_idx = (self.b_conv_buffer_idx * mtp_size).view(att_batch_size, 1) - state_offsets = torch.arange( - mtp_size, - device=self.infer_state.b_req_idx.device, - dtype=self.infer_state.b_req_idx.dtype, - ).view(1, mtp_size) - self.b_ssm_buffer_idx = b_ssm_buffer_start_idx + state_offsets # [att_batch_size, mtp_size] +class KDALinearAttDecodeAttState(LinearDecodeAttState): + state_class = KDAState def decode_att( self, @@ -222,21 +137,16 @@ def decode_att( raw_beta = params["raw_beta"] backend: KDALinearAttBackend = self.backend - state = self.infer_state.att_state_manager.get_state( - model_role=self.infer_state.att_state_role, - state_type_filter="kda", - logical_layer_id=layer_num, - state_class=KDAState, - ) - conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) - conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) + forward = self.get_forward(layer_num) + conv_states, _ = forward.resource.get_layer_view(layer_num) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=forward.req_indices) 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, + num_accepted_tokens=forward.accepted_token_counts, + query_start_loc=forward.cu_seqlens, ) mixed_qkv = conv_update(mixed_qkv, conv_states, layer_weight.get_merged_kda_conv_weight(), **conv_kwargs) q, k, v = backend.split_qkv(mixed_qkv) @@ -246,19 +156,13 @@ def decode_att( v = v.view(*shape, backend.tp_num_heads, backend.head_dim) raw_gate = raw_gate.view(*shape, backend.tp_hidden_size) raw_beta = raw_beta.view(*shape, backend.tp_num_heads) - output, _ = fused_recurrent_kda( - q=q, - k=k, - v=v, - raw_gate=raw_gate, - raw_beta=raw_beta, - a_log=layer_weight.linear_A_log.weight, - gate_bias=layer_weight.linear_dt_bias.weight, - initial_state=ssm_states, - lower_bound=backend.lower_bound, - inplace_final_state=True, - ssm_state_indices=self.b_ssm_buffer_idx, - cu_seqlens=self.b1_mtp_cu_q_seq_len, - num_accepted_tokens=self.b_num_accepted_tokens, + return forward.run_layer( + layer_num, + q, + k, + v, + raw_gate.view(-1, backend.tp_hidden_size), + raw_beta.view(-1, backend.tp_num_heads), + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, ) - return output diff --git a/lightllm/common/basemodel/attention/linear/state.py b/lightllm/common/basemodel/attention/linear/state.py new file mode 100644 index 0000000000..64d2c2243c --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/state.py @@ -0,0 +1,89 @@ +import copy +from dataclasses import dataclass, field + +import torch + +from lightllm.common.basemodel.attention.base_att import BaseDecodeAttState, BasePrefillAttState +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + + +def get_resources(att_state): + return tuple( + state + for state in att_state.infer_state.att_state_manager.get_bounded_attention_states( + att_state.infer_state.att_state_role + ) + if isinstance(state, att_state.state_class) + ) + + +def get_resource(att_state, logical_layer): + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter=att_state.state_class.algorithm, + logical_layer_id=logical_layer, + state_class=att_state.state_class, + ) + + +@dataclass +class LinearPrefillAttState(BasePrefillAttState): + state_indices: dict = field(default_factory=dict) + + def init_state(self): + + self.state_indices = { + state.state_name: state.prepare_prefill(self.infer_state.b_req_idx) for state in get_resources(self) + } + + +@dataclass +class LinearDecodeAttState(BaseDecodeAttState): + forwards: dict = field(default_factory=dict) + + def get_draft_step(self): + return self.backend.mtp_step + + def init_state(self): + draft_step = self.get_draft_step() + infer_state = self.infer_state + reqs, cu_seqlens = infer_state.b_req_idx, None + if draft_step: + if self.backend.uses_dynamic_spec_verify_layout(): + cu_seqlens, reqs = build_dynamic_mtp_linear_att_state_params( + b_req_idx=reqs, + b_mtp_index=infer_state.b_mtp_index, + hold_req_id=infer_state.req_manager.HOLD_REQUEST_ID, + ) + else: + width = draft_step + 1 + if infer_state.batch_size % width: + raise ValueError("fixed-layout decode requires batch_size divisible by draft_step + 1") + cu_seqlens = torch.arange( + 0, + infer_state.batch_size + 1, + width, + dtype=torch.int32, + device=reqs.device, + ) + reqs = reqs.view(-1, width)[:, 0].contiguous() + self.forwards = {state.state_name: state.prepare_decode(reqs, cu_seqlens) for state in get_resources(self)} + + def get_forward(self, logical_layer): + return self.forwards[get_resource(self, logical_layer).state_name] + + def get_state_commits(self): + return tuple(commit for forward in self.forwards.values() if (commit := forward.finish()) is not None) + + def clone_for_warmup(self, infer_state): + state = super().clone_for_warmup(infer_state) + state.forwards = {name: copy.copy(forward) for name, forward in self.forwards.items()} + return state + + def copy_for_decode_cuda_graph(self, new_state): + if self.forwards.keys() != new_state.forwards.keys(): + raise ValueError("captured linear attention resources changed") + for name, forward in self.forwards.items(): + forward.copy_inputs(new_state.forwards[name]) diff --git a/lightllm/common/basemodel/attention/linear/triton.py b/lightllm/common/basemodel/attention/linear/triton.py index f0b6400f51..f579f33a6c 100644 --- a/lightllm/common/basemodel/attention/linear/triton.py +++ b/lightllm/common/basemodel/attention/linear/triton.py @@ -1,7 +1,7 @@ -from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend +from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttBackend from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule -class TritonLinearAttBackend(LinearAttBackend): +class TritonLinearAttBackend(GDNLinearAttBackend): def get_prefill_kernel(self): return chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 0edbdadb94..e8e112cc97 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -577,11 +577,18 @@ def _create_padded_prefill_model_input(self, model_input: ModelInput, new_handle return new_model_input def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferStateInfo) -> ModelOutput: + state_commits = tuple( + commit + for state in (infer_state.decode_att_state, infer_state.decode_att_state1) + if state is not None + for commit in state.get_state_commits() + ) output = ModelOutput( logits=post_output.logits.contiguous(), mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, logits_token_ids=post_output.logits_token_ids, + state_commits=state_commits, ) return output diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index d8f2329b24..8e4b39688e 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -1,11 +1,14 @@ import copy from dataclasses import dataclass -from typing import List, Optional +from typing import TYPE_CHECKING, List, Optional import torch from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor +if TYPE_CHECKING: + from lightllm.common.attention_state.commit import StateCommit + @dataclass class ModelInput: @@ -202,6 +205,9 @@ class ModelOutput: # target 模型会将候选值回填到完整词表 logits，因此不需要该映射。 logits_token_ids: Optional[torch.Tensor] = None + # Detached state-commit metadata for this forward/graph; preserved when unpadding. + state_commits: tuple["StateCommit", ...] = () + def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index ab0e0c0721..b85ce3b0f2 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -1,11 +1,13 @@ import os -import torch -import torch.distributed as dist import copy import bisect import math -import triton from typing import Optional + +import torch +import torch.distributed as dist +import triton + from lightllm.utils.log_utils import init_logger from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import dist_group_manager @@ -21,6 +23,16 @@ logger = init_logger(__name__) +def _copy_warmup_state(infer_state: InferStateInfo) -> InferStateInfo: + warmup_state = copy.copy(infer_state) + for name in ("decode_att_state", "decode_att_state1"): + state = getattr(infer_state, name, None) + if state is not None: + state = state.clone_for_warmup(warmup_state) + setattr(warmup_state, name, state) + return warmup_state + + class CudaGraph: # CudaGraph forward pass for the decoding stage. @@ -123,7 +135,7 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo): torch.cuda.synchronize() # 在正式捕获前调优 decode attention，退出作用域后再捕获选定的配置。 with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - decode_func(copy.copy(infer_state)) + decode_func(_copy_warmup_state(infer_state)) torch.cuda.synchronize() for param_name in set(vars(infer_state).keys()): if param_name not in pure_para_set: @@ -156,7 +168,7 @@ def _capture_decode_overlap( pure_para_set1 = set(vars(infer_state1).keys()) torch.cuda.synchronize() with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - decode_func(copy.copy(infer_state), copy.copy(infer_state1)) + decode_func(_copy_warmup_state(infer_state), _copy_warmup_state(infer_state1)) torch.cuda.synchronize() for para_name in set(vars(infer_state).keys()): if para_name not in pure_para_set: diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py new file mode 100644 index 0000000000..c16a9d7983 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/compact_ssm.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Compact speculative state for GDN and KDA. + +Verify retains raw inputs rather than full SSM snapshots. Commit re-executes the +accepted prefix with the same per-token rounding as the recurrent baseline. +""" + +import triton +import triton.language as tl + + +@triton.jit +def _compact( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + State, + Keys, + Values, + Decays, + Betas, + Reqs, + Cu, + Accepted, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + SLOTS: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + WIDTH: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + COMMIT: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, +): + iv, row, lh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + row).to(tl.int64) + linear_layer_idx, hv = lh // HV, lh % HV + if COMMIT: + start = 0 + end = tl.load(Accepted + req) + 1 if req != HOLD else 0 + else: + start, end = tl.load(Cu + row), tl.load(Cu + row + 1) + kk = tl.arange(0, BK) + vv = iv * BV + tl.arange(0, BV) + if req == HOLD or end == start: + if not COMMIT: + for t in range(start, end): + tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) + return + slot = (linear_layer_idx * SLOTS + req) * HV + hv + key_slot = (linear_layer_idx * SLOTS + 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) + for t in range(start, end): + record = slot * WIDTH + t - start + key_record = key_slot * WIDTH + t - start + if COMMIT: + k = tl.load(Keys + key_record * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(Values + record * V + vv, vv < V, 0).to(tl.float32) + if KDA: + decay = tl.load(Decays + record * K + kk, kk < K, 1)[:, None] + else: + decay = tl.load(Decays + record) + beta = tl.load(Betas + record) + else: + h = hv // (HV // H) + k = tl.load(Kp + t * SK + h * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(Vp + t * SV + hv * V + vv, vv < V, 0).to(tl.float32) + log_a = tl.load(Alog + hv).to(tl.float32) + if KDA: + x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + x += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32) + decay = tl.exp(LOWER * tl.sigmoid(tl.exp(log_a) * x))[:, None] + else: + x = tl.load(A + t * SA + hv).to(tl.float32) + tl.load(Bias + hv).to(tl.float32) + g = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + decay = tl.exp(g) + beta = tl.sigmoid(tl.load(B + t * SB + hv).to(tl.float32)) + tl.store(Values + record * V + vv, v, vv < V) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store(Keys + key_record * K + kk, k, kk < K) + tl.store(Betas + record, beta) + if KDA: + tl.store(Decays + record * K + kk, tl.sum(decay, 1), kk < K) + else: + tl.store(Decays + record, decay) + k = k / tl.sqrt(tl.sum(k * k) + 1.0e-6) + state *= decay + d = (v - tl.sum(state * k[:, None], 0)) * beta + state += k[:, None] * d[None, :] + if not COMMIT: + q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) + q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K ** -0.5) + tl.store(Out + (t * HV + hv) * V + vv, tl.sum(state * q[:, None], 0), vv < V) + state = state.to(State.dtype.element_ty).to(tl.float32) + if COMMIT: + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py index b81bf9ec16..b0f2ea6b79 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py @@ -7,17 +7,17 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( b_req_idx, b_mtp_index, - req_to_mtp_state_index, out_cu_q_seq_len, out_conv_buffer_idx, - out_num_accepted_tokens, batch_size, hold_req_id, + SEQUENCE_CAPACITY: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): offsets = tl.arange(0, BLOCK_SIZE) token_mask = offsets < batch_size - cu_mask = offsets <= batch_size + sequence_mask = offsets < SEQUENCE_CAPACITY + cu_mask = offsets <= SEQUENCE_CAPACITY req_idx = tl.load(b_req_idx + offsets, mask=token_mask, other=hold_req_id) mtp_index = tl.load(b_mtp_index + offsets, mask=token_mask, other=0) @@ -25,34 +25,24 @@ def _build_dynamic_mtp_linear_att_state_params_kernel( valid_row = token_mask & (req_idx != hold_req_id) actual_token_num = tl.sum(tl.where(valid_row, 1, 0), axis=0) - # Keep tensor shapes dependent only on the target graph batch size. The - # unused tail represents zero-length sequences, so one captured graph can - # replay arbitrary compact per-request widths at the same token batch size. + # A request contributes at most one sequence, regardless of verify width. + # Both the request-pool capacity and graph token bucket are static. tl.store(out_cu_q_seq_len + offsets, actual_token_num, mask=cu_mask) - tl.store(out_conv_buffer_idx + offsets, hold_req_id, mask=token_mask) - tl.store(out_num_accepted_tokens + offsets, 1, mask=token_mask) + tl.store(out_conv_buffer_idx + offsets, hold_req_id, mask=sequence_mask) tl.debug_barrier() # mtp_index restarts at zero on the first row of every request group. is_start = valid_row & (mtp_index == 0) sequence_index = tl.cumsum(tl.where(is_start, 1, 0), axis=0) - 1 - accepted_state_index = tl.load( - req_to_mtp_state_index + req_idx, - mask=is_start, - other=0, - ) - tl.store(out_cu_q_seq_len + sequence_index, offsets, mask=is_start) tl.store(out_conv_buffer_idx + sequence_index, req_idx, mask=is_start) - tl.store(out_num_accepted_tokens + sequence_index, accepted_state_index + 1, mask=is_start) def build_dynamic_mtp_linear_att_state_params( b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, - req_to_mtp_state_index: torch.Tensor, hold_req_id: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor]: """Convert compact MTP verify rows to variable-length GDN sequences. Rows remain request-major after dynamic verification trimming, but each request can @@ -70,42 +60,41 @@ def build_dynamic_mtp_linear_att_state_params( request 4 -> rows [3, 4), query length 1 request 9 -> rows [4, 6), query length 2 - Runtime ``H`` rows are graph padding and do not form query sequences. If - ``req_to_mtp_state_index`` stores ``{7: 2, 4: 0, 9: 1}``, the fixed-shape - outputs are:: + Runtime ``H`` rows are graph padding and do not form query sequences. + The fixed-shape outputs are:: b1_cu_q_seq_len = [0, 3, 4, 6, 6, 6, 6, 6, 6] b_conv_buffer_idx = [7, 4, 9, H, H, H, H, H] - b_num_accepted_tokens = [3, 1, 2, 1, 1, 1, 1, 1] ``b_conv_buffer_idx`` therefore changes from one request id per query row to one request id per GDN sequence. Repeated cumulative lengths describe zero-length tail sequences. Keeping every output shape dependent only on - the padded input batch allows the same CUDA Graph to replay different - per-request query lengths. + the padded input batch and request-pool capacity allows the same CUDA Graph + to replay different per-request query lengths. The request table includes + one HOLD slot; the number of real sequences cannot exceed the other slots. """ - assert b_req_idx.is_cuda and b_mtp_index.is_cuda and req_to_mtp_state_index.is_cuda + assert b_req_idx.is_cuda and b_mtp_index.is_cuda assert b_req_idx.ndim == 1 and b_req_idx.shape == b_mtp_index.shape assert b_req_idx.dtype == torch.int32 and b_mtp_index.dtype == torch.int32 batch_size = b_req_idx.shape[0] assert batch_size > 0 + sequence_capacity = min(batch_size, hold_req_id) + assert sequence_capacity > 0 - b1_cu_q_seq_len = torch.empty((batch_size + 1,), dtype=torch.int32, device=b_req_idx.device) - b_conv_buffer_idx = torch.empty_like(b_req_idx) - b_num_accepted_tokens = torch.empty_like(b_req_idx) + b1_cu_q_seq_len = torch.empty((sequence_capacity + 1,), dtype=torch.int32, device=b_req_idx.device) + b_conv_buffer_idx = torch.empty((sequence_capacity,), dtype=torch.int32, device=b_req_idx.device) _build_dynamic_mtp_linear_att_state_params_kernel[(1,)]( b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, - req_to_mtp_state_index=req_to_mtp_state_index, out_cu_q_seq_len=b1_cu_q_seq_len, out_conv_buffer_idx=b_conv_buffer_idx, - out_num_accepted_tokens=b_num_accepted_tokens, batch_size=batch_size, hold_req_id=int(hold_req_id), + SEQUENCE_CAPACITY=sequence_capacity, BLOCK_SIZE=triton.next_power_of_2(batch_size + 1), num_warps=8, num_stages=1, ) - return b1_cu_q_seq_len, b_conv_buffer_idx, b_num_accepted_tokens + return b1_cu_q_seq_len, b_conv_buffer_idx diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py new file mode 100644 index 0000000000..c50e0c680f --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -0,0 +1,582 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GDN/KDA ReplaySSM in LightLLM's [K, V] state layout. + +The checkpoint contains only folded tokens. Derived records reconstruct verify +outputs, while raw inputs replay the accepted suffix into the checkpoint. Verify +may overwrite the uncommitted suffix; only acceptance advances the cursor. +FP32 accumulation is rounded to the checkpoint dtype only on fold/materialize; +BF16 replay therefore differs from native per-token BF16 state rounding. +Algorithm reference: https://dao-lab.ai/blog/2026/replayssm/ +""" + +import triton +import triton.language as tl + + +@triton.jit +def _accept_updates( + cursors, + reqs, + accepted_state_indices, + N: tl.constexpr, + HOLD: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + r = tl.load(reqs + i, i < N, HOLD) + valid = (i < N) & (r != HOLD) + cursor = tl.load(cursors + r, valid, 0) + count = tl.load(accepted_state_indices + r, valid, 0) + 1 if WIDTH > 1 else 1 + tl.store(cursors + r, cursor + count, valid) + + +@triton.jit +def _prepare_fold( + Reqs, + Cursors, + Positions, + ForwardCursors, + Cu, + Active, + N: tl.constexpr, + HOLD: tl.constexpr, + L: tl.constexpr, + VARLEN: tl.constexpr, + BLOCK: tl.constexpr, + WIDTH: tl.constexpr, + SEPARATE_FOLD: tl.constexpr, +): + seq = tl.arange(0, BLOCK) + req = tl.load(Reqs + seq, seq < N, HOLD) + valid = (seq < N) & (req != HOLD) + cursor = tl.load(Cursors + req, valid, 0) + tl.store(Positions + seq, cursor, seq < N) + if VARLEN: + start = tl.load(Cu + seq, seq < N, 0) + end = tl.load(Cu + seq + 1, seq < N, 0) + length = end - start + else: + length = tl.full((BLOCK,), 1, tl.int32) + overflow = cursor % (2 * L) + length > L + tl.store(Cursors + req, (cursor & (2 * L)) ^ (2 * L), valid & overflow) + needs_fold = valid & (length > 0) & overflow + offset = tl.cumsum(needs_fold.to(tl.int32), axis=0) - 1 + tl.store(Active + offset, seq, needs_fold) + count = tl.sum(needs_fold.to(tl.int32), axis=0) + live = tl.sum((valid & (length > 0)).to(tl.int32), axis=0) + split = (count > 0) & ((live <= 12) | (2 * count <= live)) + if N >= 256 and L == 8 and WIDTH == 4: + short = tl.sum((valid & (length > 0) & (length < WIDTH)).to(tl.int32), axis=0) + split = split & ((live < 256) | (short > 0)) + if SEPARATE_FOLD: + split = count > 0 + tl.store(ForwardCursors + req, tl.where(overflow & split, (cursor & (2 * L)) ^ (2 * L), cursor), valid) + tl.store(Active + N, tl.where(split, count, 0)) + + +@triton.jit +def _replay_update(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): + 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 _fold_history( + State, + RawKeys, + RawValues, + Gates, + Betas, + Reqs, + Cursors, + Active, + N: tl.constexpr, + LAYERS: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SS: tl.constexpr, + RS: tl.constexpr, + VS: tl.constexpr, + GS: tl.constexpr, + BS: tl.constexpr, + KDA: tl.constexpr, + Keys, + Deltas, + DERIVED: tl.constexpr, + BL: tl.constexpr, +): + count = tl.load(Active + N) + work = tl.program_id(0) + tiles: tl.constexpr = (V + BV - 1) // BV + total = count * LAYERS * HV * tiles + while work < total: + iv = work % tiles + slot_work = work // tiles + seq = tl.load(Active + slot_work % count) + layer_hv = slot_work // count + linear_layer_idx, hv = (layer_hv // HV).to(tl.int64), layer_hv % HV + req = tl.load(Reqs + seq).to(tl.int64) + cursor = tl.load(Cursors + seq) + n = cursor % (2 * L) + base = (cursor // (2 * L)) * L + kk = tl.arange(0, BK) + vv = iv * BV + tl.arange(0, BV) + slot = req * HV + hv + key_slot = req * KH + hv // (HV // KH) + sp = State + linear_layer_idx * SS + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + if DERIVED: + tl.static_assert(not KDA) + ll = tl.arange(0, BL) + g = tl.load(Gates + linear_layer_idx * GS + slot * (2 * L) + base + ll, ll < n, 0).to(tl.float32) + prefix = tl.cumsum(g, 0) + total_g = tl.sum(g, 0) + keys = tl.load( + Keys + linear_layer_idx * RS + (key_slot * (2 * L) + base + ll[None, :]) * K + kk[:, None], + (kk[:, None] < K) & (ll[None, :] < n), + 0, + ).to(tl.float32) + delta = tl.load( + Deltas + linear_layer_idx * VS + (slot * (2 * L) + base + ll[:, None]) * V + vv[None, :], + (ll[:, None] < n) & (vv[None, :] < V), + 0, + ).to(tl.float32) + keys *= tl.exp(total_g - prefix)[None, :] + state = state * tl.exp(total_g) + tl.dot(keys, delta, input_precision="tf32x3") + else: + for j in range(n): + raw_k = tl.load( + RawKeys + linear_layer_idx * RS + (key_slot * (2 * L) + base + j) * K + kk, kk < K, 0 + ).to(tl.float32) + raw_v = tl.load(RawValues + linear_layer_idx * VS + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to( + tl.float32 + ) + if KDA: + g = tl.load(Gates + linear_layer_idx * GS + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to( + tl.float32 + ) + else: + g = tl.load(Gates + linear_layer_idx * GS + slot * (2 * L) + base + j).to(tl.float32) + beta = tl.load(Betas + linear_layer_idx * BS + slot * (2 * L) + base + j).to(tl.float32) + state = _replay_update(state, raw_k, raw_v, g, beta, KDA) + # Start the new history from the same rounded checkpoint future calls load. + state = state.to(State.dtype.element_ty).to(tl.float32) + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + work += tl.num_programs(0) + + +@triton.jit +def _replay( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + State, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Reqs, + Cursors, + 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, + VARLEN: tl.constexpr, + PRECOMPUTE_PROJECTIONS: tl.constexpr, + BW: tl.constexpr, + FUSED_FOLD: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + PARALLEL_SOLVE: tl.constexpr, + BQK: tl.constexpr, + BR: tl.constexpr, +): + """Output reconstruction with an exact raw-input recurrence at checkpoint folds.""" + iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + seq).to(tl.int64) + if VARLEN: + start, end = tl.load(Cu + seq), tl.load(Cu + seq + 1) + else: + start, end = seq, 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: + return + + kk = tl.arange(0, BK) + ll = tl.arange(0, L) + slot = req * HV + hv + key_slot = req * KH + hv // (HV // KH) + cursor = tl.load(Cursors + req) + n = cursor % (2 * L) + base = (cursor // (2 * L)) * L + sp = State + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + + if FUSED_FOLD and n + end - start > L: + for j in range(n): + raw_k = tl.load(RawKeys + (key_slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to(tl.float32) + if KDA: + g = tl.load(Gates + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + else: + g = tl.load(Gates + slot * (2 * L) + base + j).to(tl.float32) + beta = tl.load(Betas + slot * (2 * L) + base + j).to(tl.float32) + state = _replay_update(state, raw_k, raw_v, g, beta, KDA) + # Start the new history from the same rounded checkpoint future calls load. + state = state.to(State.dtype.element_ty).to(tl.float32) + tl.store(sp, state, (kk[:, None] < K) & (vv[None, :] < V)) + base, n = L - base, 0 + + if PRECOMPUTE_PROJECTIONS: + projection_dtype: tl.constexpr = ( + tl.bfloat16 if State.dtype.element_ty == tl.bfloat16 and Q.dtype.element_ty == tl.bfloat16 else tl.float32 + ) + projection_precision: tl.constexpr = "tf32" if Q.dtype.element_ty == tl.bfloat16 else "tf32x3" + tt = tl.arange(0, BQK) + tokens = start + tt % BW + valid = (tt < 2 * BW) & (tokens < end) + h = hv // (HV // H) + q_rows = tl.load( + Q + tokens[:, None] * SQ + h * K + kk[None, :], + (valid & (tt < BW))[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + k_rows = tl.load( + Kp + tokens[:, None] * SK + h * K + kk[None, :], + (valid & (tt >= BW))[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + qk = q_rows + k_rows + if PARALLEL_SOLVE: + qk /= tl.sqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] + else: + qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] + qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] + if KDA: + history = ll < n + gates = tl.load( + Gates + (slot * (2 * L) + base + 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_inputs = tl.load( + A + (start + wi[:, None]) * SA + hv * K + kk[None, :], + (start + wi[:, None] < end) & (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. + projections = tl.dot( + (qk * tl.exp(query_gates)).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 + ) + + if PRECOMPUTE_PROJECTIONS and not KDA: + # Project the whole verify window against both accepted and new keys. + # Only the small corrected-value solve below depends on token order. + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + cached_keys = tl.load( + Keys + (key_slot * (2 * L) + base) * K + jj[:, None] * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + raw_keys = tl.load( + Kp + fresh_tokens[:, 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] + all_keys = cached_keys + fresh_keys + values = tl.load( + Vp + fresh_tokens[:, 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 + fresh_tokens * 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 * (2 * L) + base + 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 * (2 * L) + base) * V + jj[:, None] * V + vv[None, :], + old[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + betas = tl.sigmoid(tl.load(B + fresh_tokens * SB + hv, fresh, 0).to(tl.float32)) + if PARALLEL_SOLVE: + projected_history = projections + tl.dot( + coefficients.to(projection_dtype), deltas.to(projection_dtype), input_precision=projection_precision + ) + wi = tl.arange(0, BQK) + active = (wi < BW) & (start + wi < end) + key_rows = tl.gather( + projected_history, tl.broadcast_to((BW + wi % BW)[:, None], projected_history.shape), 0 + ) + fresh_values = tl.gather(values, tl.broadcast_to((L + wi % BW)[:, None], (BQK, BV)), 0) + beta_rows = tl.sigmoid(tl.load(B + (start + wi) * SB + hv, active, 0).to(tl.float32)) + rhs = tl.where(active[:, None], beta_rows[:, None] * (fresh_values - key_rows), 0.0) + cross = tl.gather(coefficients, tl.broadcast_to((BW + wi % BW)[:, None], coefficients.shape), 0) + cross = tl.gather(cross, tl.broadcast_to((L + wi % BW)[None, :], (BQK, BQK)), 1) + lower = active[:, None] & active[None, :] & (wi[:, None] > wi[None, :]) + power = tl.where(lower, -beta_rows[:, None] * cross, 0.0) + inverse = (wi[:, None] == wi[None, :]).to(tl.float32) + # The strictly lower-triangular coupling is nilpotent. Repeated + # squaring builds (I - power)^-1 in log2(BW) matrix steps; capacity + # is at most 64, so six steps cover every supported verify width. + for solve in tl.static_range(6): + if (1 << solve) < BW: + inverse += tl.dot(power, inverse, input_precision="tf32x3") + power = tl.dot(power, power, input_precision="tf32x3") + solved = tl.dot( + inverse.to(projection_dtype), rhs.to(projection_dtype), input_precision=projection_precision + ) + mapped = tl.gather(solved, tl.broadcast_to(((jj - L) % BW)[:, None], deltas.shape), 0) + deltas = tl.where(fresh[:, None], mapped, deltas) + else: + 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 + tl.sum(deltas * weights[:, None], 0) + 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 + ) + tl.store( + Out + (tokens[:, None] * HV + hv) * V + vv[None, :], + outputs, + (valid & (tt < BW))[:, None] & (vv[None, :] < V), + ) + records = base + n + jj - L + tl.store( + Deltas + (slot * (2 * L) + records[:, None]) * V + vv[None, :], + deltas, + fresh[:, None] & (vv[None, :] < V), + ) + tl.store( + RawValues + (slot * (2 * L) + records[:, None]) * V + vv[None, :], + values, + fresh[:, None] & (vv[None, :] < V), + ) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store( + Keys + (key_slot * (2 * L) + records[:, None]) * K + kk[None, :], + fresh_keys, + fresh[:, None] & (kk[None, :] < K), + ) + tl.store( + RawKeys + (key_slot * (2 * L) + records[:, None]) * K + kk[None, :], + raw_keys, + fresh[:, None] & (kk[None, :] < K), + ) + tl.store(Gates + slot * (2 * L) + records, fresh_gates, fresh) + tl.store(Betas + slot * (2 * L) + records, betas, fresh) + return + + if not PRECOMPUTE_PROJECTIONS or KDA: + history = ll < n + keys = tl.load( + Keys + (key_slot * (2 * L) + base) * K + ll[:, None] * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + deltas = tl.load( + Deltas + (slot * (2 * L) + base) * V + ll[:, None] * V + vv[None, :], + history[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + if KDA: + gates = tl.load( + Gates + (slot * (2 * L) + base + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + else: + gates = tl.load(Gates + slot * (2 * L) + base + 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(Kp + t * SK + h * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(Vp + 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 = tl.sum(keys * k[None, :], axis=1) + hq = tl.sum(keys * q[None, :], axis=1) + if PRECOMPUTE_PROJECTIONS: + checkpoint_sk = tl.sum(tl.where((tt == BW + t - start)[:, None], projections, 0), 0) + checkpoint_sq = tl.sum(tl.where((tt == t - start)[:, None], projections, 0), 0) + else: + 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 + tl.sum(deltas * hk[:, None], axis=0) + sq = checkpoint_sq + tl.sum(deltas * hq[:, None], axis=0) + else: + weights = tl.where(history, tl.exp(total_g - gate_prefix), 0.0) + hk = tl.sum(keys * k[None, :], axis=1) * weights + hq = tl.sum(keys * q[None, :], axis=1) * 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) + tl.sum(deltas * hk[:, None], axis=0) + sq = checkpoint_sq * tl.exp(total_g) + tl.sum(deltas * hq[:, None], axis=0) + d = beta * (v - sk) + tl.store(Out + (t * HV + hv) * V + vv, sq + d * tl.sum(k * q), vv < V) + record = base + n + tl.store(Deltas + (slot * (2 * L) + record) * V + vv, d, vv < V) + tl.store(RawValues + (slot * (2 * L) + record) * V + vv, v, 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 * (2 * L) + record) * K + kk, k, kk < K) + tl.store(RawKeys + (key_slot * (2 * L) + record) * K + kk, raw_k, kk < K) + if KDA: + tl.store(Gates + (slot * (2 * L) + record) * K + kk, g, kk < K) + else: + tl.store(Gates + slot * (2 * L) + record, g) + tl.store(Betas + slot * (2 * L) + record, beta) + 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 + + +@triton.jit +def _materialize( + State, + Output, + 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, + SNAPSHOT: tl.constexpr, + KDA: tl.constexpr, +): + 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 + cursor = tl.load(Cursors + req) + n = cursor % (2 * L) + base = (cursor // (2 * L)) * L + if n == 0 and not SNAPSHOT: + return + linear_layer_idx, head = lh // HV, lh % HV + 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 + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + for j in range(n): + raw_k = tl.load(RawKeys + (key_slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * (2 * L) + base + j) * V + vv, vv < V, 0).to(tl.float32) + if KDA: + g = tl.load(Gates + (slot * (2 * L) + base + j) * K + kk, kk < K, 0).to(tl.float32) + else: + g = tl.load(Gates + slot * (2 * L) + base + j).to(tl.float32) + beta = tl.load(Betas + slot * (2 * L) + base + j).to(tl.float32) + raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + state *= tl.exp(g)[:, None] if KDA else tl.exp(g) + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + out_slot = linear_layer_idx * HV + head if SNAPSHOT else slot + out = Output + out_slot * K * V + kk[:, None] * V + vv[None, :] + tl.store(out, state, (kk[:, None] < K) & (vv[None, :] < V)) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py new file mode 100644 index 0000000000..052f61fb8f --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -0,0 +1,186 @@ +"""Tune one layer's SSM update cycle using the shared autotuner.""" + +import ast + +import torch +import triton + +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, autotune +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +def get_configs(): + 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 static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens): + axis = 1 if cu_seqlens is not None else 0 + _, _, hv, kd, vd = ssm_backend.state.shape + # Keep filenames short, but distinguish packed projections, TP layouts, + # state precision, recurrent rule and history capacity. + key = { + "mode": mode, + "v": 6, + "shape": str((q.shape[-2], hv, ssm_backend.num_key_heads, kd, vd, ssm_backend.verify_width)), + "dtype": str((q.dtype, ssm_backend.state.dtype)), + "strides": str((q.stride(axis), k.stride(axis), v.stride(axis), a.stride(0), b.stride(0))), + "history": ssm_backend.capacity if mode == "replay" else 0, + "gate": int(ssm_backend.kda), + "varlen": cu_seqlens is not None, + } + if ssm_backend.kda: + key["lower"] = ssm_backend.gate_lower_bound + if mode == "replay": + key["history_dtype"] = str(ssm_backend.keys.dtype) + if ssm_backend.compute_mode == "parallel": + key["projection"] = "parallel" + key["fold"] = ssm_backend.fold_method + return key + key["projection"] = ( + ("kda_mma" if ssm_backend.kda else "gdn_mma") + if ssm_backend.compute_mode == "precompute" and ssm_backend.verify_width > 1 + else "inline" + ) + return key + + +def run_key(ssm_backend, q, cu_seqlens): + tokens = q.shape[1 if cu_seqlens is not None else 0] + return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens + + +def scratch_size_bytes(ssm_backend, batch, sequences, inputs): + # State/history have a leading layer axis; cursors are shared across layers. + layers, slots = ssm_backend.state.shape[:2] + pool_bytes = sum( + x.numel() * x.element_size() // (layers if x.ndim > 1 else 1) + for x in vars(ssm_backend).values() + if isinstance(x, torch.Tensor) + ) + input_bytes = sum( + (1 + sum((size - 1) * stride for size, stride in zip(x.shape, x.stride()))) * x.element_size() + for x in inputs + if x.numel() + ) + return pool_bytes // slots * (batch + 1) + input_bytes + 32 * (batch + sequences + 1) + + +def rebuild_inputs(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None): + tokens, sequences = run_key(ssm_backend, q, cu_seqlens) + batch = min(sequences, triton.cdiv(tokens, ssm_backend.verify_width)) if cu_seqlens is not None else tokens + if q.is_cuda: + required = scratch_size_bytes(ssm_backend, batch, sequences, (q, k, v, a, b, a_log, bias)) + free, _ = torch.cuda.mem_get_info(q.device) + available = free + torch.cuda.memory_reserved(q.device) - torch.cuda.memory_allocated(q.device) + # Leave half for benchmark intermediates, CUDA Graphs and allocator fragmentation. + if required > available // 2: + raise torch.cuda.OutOfMemoryError( + f"SSM tuning scratch needs {required} bytes, exceeding its {available // 2}-byte budget" + ) + _, _, hv, kd, vd = ssm_backend.state.shape + state = torch.empty((1, batch + 1, hv, kd, vd), device=q.device, dtype=ssm_backend.state.dtype).normal_(0, 0.02) + if mode == "replay": + scratch = type(ssm_backend)( + state, + ssm_backend.capacity, + ssm_backend.verify_width, + q.dtype, + num_key_heads=ssm_backend.num_key_heads, + kda=ssm_backend.kda, + gate_lower_bound=ssm_backend.gate_lower_bound, + compute_mode=ssm_backend.compute_mode, + fold_method=ssm_backend.fold_method, + ) + else: + scratch = type(ssm_backend)( + state, + ssm_backend.verify_width, + q.dtype, + num_key_heads=ssm_backend.num_key_heads, + kda=ssm_backend.kda, + gate_lower_bound=ssm_backend.gate_lower_bound, + ) + # Benchmark callbacks must not recursively select a configuration. + scratch._config_is_fixed = True + inputs = [ + torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=x.dtype).normal_(0, 0.2) + for x in (q, k, v, a, b, a_log, bias) + ] + reqs = torch.cat( + ( + torch.arange(batch, dtype=torch.int32, device=q.device), + torch.full((sequences - batch,), batch, dtype=torch.int32, device=q.device), + ) + ) + cu = (torch.arange(sequences + 1, dtype=torch.int32, device=q.device) * ssm_backend.verify_width).clamp_max(tokens) + lengths = cu[1 : batch + 1] - cu[:batch] + accepted = torch.cat((lengths - 1, lengths.new_zeros(1))) + partial = accepted.clone() + partial[:-1] = torch.minimum(reqs[:batch] % ssm_backend.verify_width, lengths - 1) + return (), dict( + ssm_backend=scratch, + mode=mode, + q=inputs[0], + k=inputs[1], + v=inputs[2], + a=inputs[3], + b=inputs[4], + a_log=inputs[5], + bias=inputs[6], + cu_seqlens=cu if cu_seqlens is not None else None, + workload=(reqs, accepted, partial), + ) + + +@autotune( + kernel_name="ssm_update_cycle:v1", + configs_gen_func=get_configs, + static_key_func=static_key, + run_key_func=run_key, + run_key_distance_func=lambda a, b: sum(abs(x - y) for x, y in zip(ast.literal_eval(a), ast.literal_eval(b))), + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + rebuild_input_func=rebuild_inputs, + warmup_all_exist_config=False, +) +def select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None): + """Return a layout; only rebuilt disposable inputs execute the timed cycle.""" + if workload is None: + return run_config + ssm_backend.run_config = run_config + reqs, accepted, partial = workload + if mode == "replay": + ssm_backend.cursors.zero_() + # Include an entire history period and the subsequent fold. Repeating + # just the first round would benchmark an empty history forever. + tokens = q.shape[1 if cu_seqlens is not None else 0] + rounds = ssm_backend.capacity // min(ssm_backend.verify_width, tokens) + 1 + else: + rounds = 2 # Partial acceptance followed by full acceptance. + for step in range(rounds): + ssm_backend.prepare_decode(reqs, cu_seqlens) + ssm_backend.forward(0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) + ssm_backend.accept_updates(reqs, partial if mode != "replay" and step == 0 else accepted) + return run_config + + +def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens): + default = ssm_backend.run_config or {"BV": 32, "num_warps": 1} + if ssm_backend._config_is_fixed: + return default + key = (tuple(static_key(ssm_backend, mode, q, k, v, a, b, cu_seqlens).items()), run_key(ssm_backend, q, cu_seqlens)) + if key in ssm_backend.batch_configs: + return ssm_backend.batch_configs[key] + if not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + return default + config = select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) + config = dict(config or default) + # Each captured graph keeps its own layout, even after other buckets warm up. + ssm_backend.batch_configs[key] = config + logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") + return config diff --git a/lightllm/common/basemodel/triton_kernel/mtp_utils.py b/lightllm/common/basemodel/triton_kernel/mtp_utils.py index 7525a754df..c72f8191df 100644 --- a/lightllm/common/basemodel/triton_kernel/mtp_utils.py +++ b/lightllm/common/basemodel/triton_kernel/mtp_utils.py @@ -349,6 +349,8 @@ def _fwd_kernel_linear_att_mtp_state_index_update( b_mtp_index, accepted_index, req_mtp_all_num, + replay_cursors, + HOLD: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): cur_index = tl.program_id(0) @@ -369,6 +371,9 @@ def _fwd_kernel_linear_att_mtp_state_index_update( max_mtp_index = tl.max(valid_mtp_index, axis=0) tl.store(req_to_mtp_state_index + cur_req_idx, max_mtp_index) + if replay_cursors is not None: + cursor = tl.load(replay_cursors + cur_req_idx, mask=cur_req_idx != HOLD, other=0) + tl.store(replay_cursors + cur_req_idx, cursor + max_mtp_index + 1, mask=cur_req_idx != HOLD) return @@ -379,6 +384,7 @@ def linear_att_mtp_state_index_update( b_mtp_index: torch.Tensor, accepted_index: torch.Tensor, verify_width: int, + replay_cursors: torch.Tensor = None, ): """ Update req_to_mtp_state_index with the max b_mtp_index among accepted tokens per request. @@ -406,6 +412,8 @@ def linear_att_mtp_state_index_update( b_mtp_index=b_mtp_index, accepted_index=accepted_index, req_mtp_all_num=req_mtp_all_num, + replay_cursors=replay_cursors, + HOLD=req_to_mtp_state_index.numel() - 1, BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=1, diff --git a/lightllm/common/ssm_types.py b/lightllm/common/ssm_types.py new file mode 100644 index 0000000000..636f7ceffe --- /dev/null +++ b/lightllm/common/ssm_types.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class SSMRunConfig: + BV: int + num_warps: int + num_stages: int = 3 + + def __post_init__(self): + for value in (self.BV, self.num_warps, self.num_stages): + if type(value) is not int: + raise TypeError("SSM kernel configuration must be an integer") diff --git a/lightllm/common/triton_utils/autotuner.py b/lightllm/common/triton_utils/autotuner.py index b7a94d657e..c664bbe2b0 100644 --- a/lightllm/common/triton_utils/autotuner.py +++ b/lightllm/common/triton_utils/autotuner.py @@ -181,6 +181,7 @@ def __init__( self.cached_configs = {} self.fast_match_configs = collections.defaultdict(dict) self.warmuped_configs_set = set() + self._oom_keys = set() self.arg_names = [param.name for param in inspect.signature(self.fn).parameters.values()] self._argname_to_pos = {name: idx for idx, name in enumerate(self.arg_names)} self._pos_to_argname = {idx: name for idx, name in enumerate(self.arg_names)} @@ -244,6 +245,8 @@ def __call__(self, *args, **kwargs): need_tuning = (autotune_level == AutotuneLevel.FORCE_AUTOTUNE) or ( run_key not in self.cached_configs.get(static_key, {}) ) + if (static_key, run_key) in self._oom_keys: + need_tuning = False if world_size > 1: _need_tunings = [None for _ in range(world_size)] dist.all_gather_object(_need_tunings, obj=need_tuning, group=self._get_autotune_group()) @@ -390,7 +393,22 @@ def _autotune(self, args, kwargs, static_key, run_key, rank_id, world_size): # 仅为本次调优重建输入，构造开销不计入计时，最终执行仍使用调用方的原始输入。 if self.rebuild_input_func is not None: - args, kwargs = self.rebuild_input_func(*args, **kwargs) + ready = True + try: + args, kwargs = self.rebuild_input_func(*args, **kwargs) + except torch.cuda.OutOfMemoryError as error: + ready = False + logger.warning(f"Cannot allocate tuning inputs for {self.kernel_name}: {error}") + if world_size > 1: + all_ready = [None for _ in range(world_size)] + dist.all_gather_object(all_ready, obj=ready, group=self._get_autotune_group()) + ready = all(all_ready) + if not ready: + # Every rank skips the search. Keep existing/default configs without + # persisting a memory-pressure fallback as a measured tuning result. + self._oom_keys.add((static_key, run_key)) + logger.warning(f"Skipping autotune for {self.kernel_name} / {run_key}: insufficient scratch memory") + return best_config = None best_time = float("inf") diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 0e52e736a9..f19f9028d3 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -17,6 +17,7 @@ class Glm5NextResourceConfig(Qwen3NextResourceConfig): """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" # 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 @@ -45,6 +46,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/indexer.py b/lightllm/models/glm5_next/indexer.py index 5df510e5de..8bb48cc44e 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -126,12 +126,15 @@ def _get_decode_indices(self, q_fp8, weights, packed_buffer, infer_state, att_st lengths, max_pools, ) - metadata = deep_gemm.get_paged_mqa_logits_metadata(lengths, 64, deep_gemm.get_num_sms()) + # DeepGEMM's SM90 scheduler cannot handle a large batch of empty HOLD rows. + # Their initialized dummy pool is discarded by top-k's original lengths. + schedule_lengths = lengths.clamp_min(1) + metadata = deep_gemm.get_paged_mqa_logits_metadata(schedule_lengths, 64, deep_gemm.get_num_sms()) # Each MTP position has its own pool length; HOLD rows have no valid pools. # Fixed-width mtp_step=2 gives Q=[3 * num_requests, 1, heads, dim], so next_n stays 1. # All three verify positions are preserved despite the SM90 kernel's native next_n limit of 2. logits = deep_gemm.fp8_paged_mqa_logits( - q_fp8.unsqueeze(1), pages, weights, lengths, block_table, metadata, max_pools, clean_logits=False + q_fp8.unsqueeze(1), pages, weights, schedule_lengths, block_table, metadata, max_pools, clean_logits=False ) groups = torch.empty((q_fp8.shape[0], self.topk // 4), dtype=torch.int32, device=q_fp8.device) self.select_topk_indices(logits, lengths.view(-1), groups) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index 5b5235c76e..2b2977bcb9 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -79,6 +79,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()) diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 888af5cd35..0e0ed924d2 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -1,9 +1,10 @@ import os import torch -from lightllm.common.basemodel.batch_objs import PostLayerOutput import torch.functional as F import torch.distributed as dist import numpy as np + +from lightllm.common.basemodel.batch_objs import PostLayerOutput from lightllm.common.basemodel.layer_weights.base_layer_weight import BaseLayerWeight from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo @@ -102,6 +103,17 @@ def _lm_head_and_gather( ) -> PostLayerOutput: """执行原始 lm-head，并收集完整词表 logits。""" + if self.tp_world_size_ == 1 and hidden.is_cuda and hidden.dtype == torch.bfloat16: + # Produce token-major logits directly so the FP32 conversion reads + # contiguous rows instead of transposing the full vocabulary matrix. + normed = self._norm(hidden, infer_state, layer_weight) + vocab_size = layer_weight.lm_head_weight_.vocab_size + local_logits = self.alloc_tensor((token_num, vocab_size), dtype=hidden.dtype) + torch.mm(normed, layer_weight.lm_head_weight_.weight.t(), out=local_logits) + logits = self.alloc_tensor((token_num, vocab_size), dtype=torch.float32) + logits.copy_(local_logits) + return PostLayerOutput(logits=logits) + local_logits = self._project_local_logits(hidden, token_num, layer_weight, infer_state) vocab_size = layer_weight.lm_head_weight_.vocab_size diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 39f9c42919..4fe3e4e93c 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -1,3 +1,6 @@ +from typing import Tuple +from functools import partial + import torch import torch.distributed as dist @@ -9,12 +12,10 @@ from lightllm.utils.log_utils import init_logger from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor from lightllm.common.basemodel.attention.base_att import AttControl -from typing import Tuple from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ from lightllm.distributed import all_reduce from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd from lightllm.utils.envs_utils import get_env_start_args -from functools import partial logger = init_logger(__name__) @@ -316,7 +317,7 @@ def _linear_post( ) -> torch.Tensor: num_tokens = z.shape[0] core_attn_out = core_attn_out.view(-1, core_attn_out.shape[-1]) - z = z.contiguous().view(-1, z.shape[-1]) + # Gated RMSNorm accepts the strided [tokens, heads, dim] projection slice. norm_out = layer_weight.linear_norm(core_attn_out, z, self.eps_) core_attn_out = norm_out.view(num_tokens, -1) output = layer_weight.linear_out_proj.mm(core_attn_out) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index f1a3a65748..20fbeb4749 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -81,6 +81,8 @@ 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] diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 627f2c741a..fee7c6401d 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: @@ -974,6 +975,20 @@ 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", "compact", "replay"], + default="native", + help="SSM state storage: native snapshots, compact MTP records, or deferred GDN/KDA replay (FP32/BF16).", + ) + parser.add_argument( + "--ssm_replay_config", + type=json.loads, + default={}, + help='Replay-only JSON object: {"cache_len":16,"compute_mode":"inline","fold_method":"raw"}. ' + "cache_len: 4/8/16/32/64, at least the verify width; compute_mode: inline/precompute/parallel; " + "fold_method: raw/derived. Parallel requires BF16 GDN MTP with 128x128 heads; derived requires parallel.", + ) 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 3e2fb44b5b..52da4fb5e1 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -43,6 +43,65 @@ def _set_envs_and_config(args: StartArgs): mp.set_start_method("spawn", force=True) +def _validate_and_normalize_ssm_config(args: StartArgs): + assert args.ssm_state_mode in ("native", "compact", "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": "inline", "fold_method": "raw"} + 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", "parallel"), "Invalid replay compute_mode" + assert config["fold_method"] in ("raw", "derived"), "Invalid replay fold_method" + assert ( + config["fold_method"] == "raw" or config["compute_mode"] == "parallel" + ), "derived fold requires parallel ReplaySSM compute_mode" + if config["compute_mode"] == "parallel": + assert args.mtp_step > 0, "Parallel ReplaySSM requires MTP" + assert args.linear_att_ssm_data_type == "bfloat16", "Parallel ReplaySSM requires BF16 state" + 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 + + is_gdn_or_kda = get_model_type(args.model_dir) in ( + "glm5_next", + "glm5_next_text", + "qwen3_next", + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + ) + 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" + if args.ssm_state_mode == "compact": + assert args.mtp_step > 0, "ssm_state_mode=compact requires MTP (mtp_step > 0)" + else: + 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" + + def _launch_subprocesses(args: StartArgs): _validate_pd_features(args) _set_envs_and_config(args) @@ -60,6 +119,12 @@ def _launch_subprocesses(args: StartArgs): not args.disable_cudagraph or args.run_mode == "prefill" ), "--disable_cudagraph is only supported on Prefill nodes when --mtp_mode is enabled" + # Prepare remote configuration before any model-dependent defaults or validation. + if args.model_dir and "s3://" in args.model_dir: + from lightllm.utils.petrel_helper import s3_model_prepare + + s3_model_prepare(args.model_dir) + auto_set_max_req_total_len(args) auto_set_fused_shared_experts(args) set_unique_server_name(args) @@ -297,6 +362,8 @@ def _launch_subprocesses(args: StartArgs): f"but got {args.batch_max_tokens}, {args.chunked_prefill_size}" ) + _validate_and_normalize_ssm_config(args) + args.disable_chunked_mm_input = args.disable_chunked_mm_input or get_default_disable_chunked_mm_input( args.model_dir ) @@ -327,12 +394,6 @@ def _launch_subprocesses(args: StartArgs): args.cpu_cache_token_page_size % args.page_size == 0 ), "--cpu_cache_token_page_size must be divisible by --page_size" - # help to manage data stored on Ceph - if "s3://" in args.model_dir: - from lightllm.utils.petrel_helper import s3_model_prepare - - s3_model_prepare(args.model_dir) - # 如果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 diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index baaefae2e3..cb2239b93c 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -251,6 +251,8 @@ class StartArgs: weight_version: str = "default" # hybrid attention model (Qwen3Next) + ssm_state_mode: str = "native" + ssm_replay_config: dict = 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/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 056c36f596..640c049aa3 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -1,9 +1,8 @@ import enum -import torch -import numpy as np import collections import pickle - +import torch +import numpy as np from sortedcontainers import SortedDict from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Dict, Tuple, Optional, Callable, Any, Union @@ -415,7 +414,7 @@ def get_can_alloc_token_num(self): def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): """Snapshot request-level attention state at big/small-page boundaries.""" - if not self.is_hybrid_att_model: + if not self.is_hybrid_att_model or self.radix_cache is None: return # Request-state snapshot at a big-page boundary. diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 9bc8c16cf7..b18d0a6186 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -1,6 +1,9 @@ -import torch +from lightllm.common.attention_state.commit import StateCommitBatch import time from typing import List + +import torch + from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventPack @@ -296,6 +299,7 @@ def decode_mtp( b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, + state_commit_batches=(StateCommitBatch(req_num, model_output.state_commits),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index ed7e5bb167..00b298d463 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -1,6 +1,9 @@ -import torch +from lightllm.common.attention_state.commit import StateCommitBatch import time from typing import List, Tuple + +import torch + from lightllm.common.basemodel.triton_kernel.mtp_utils import gen_b_req_mtp_start_loc from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput @@ -531,6 +534,7 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): b_req_idx=model_input.b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=model_input.b_mtp_index, + state_commit_batches=(StateCommitBatch(req_num, model_output.state_commits),), ) accepted_index_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor( key="accepted_index", @@ -805,6 +809,10 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf b_req_idx=b_req_idx, b_req_mtp_start_loc=b_req_mtp_start_loc, b_mtp_index=b_mtp_index, + state_commit_batches=( + StateCommitBatch(real_request_num0, model_output0.state_commits), + StateCommitBatch(real_request_num1, model_output1.state_commits), + ), ) mtp_accept_len0 = mtp_accept_len[:real_request_num0] mtp_accept_len1 = mtp_accept_len[real_request_num0:] diff --git a/lightllm/server/router/model_infer/mtp_speculative/utils.py b/lightllm/server/router/model_infer/mtp_speculative/utils.py index a32cedd17d..3338f342f7 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/utils.py +++ b/lightllm/server/router/model_infer/mtp_speculative/utils.py @@ -5,6 +5,8 @@ import torch +from lightllm.common.attention_state.commit import MtpCommitContext + from lightllm.common.basemodel.triton_kernel.mtp_utils import ( mtp_scatter_next_token_ids, mtp_verify, @@ -22,9 +24,9 @@ def verify_mtp_tokens( b_req_idx: torch.Tensor, b_req_mtp_start_loc: torch.Tensor, b_mtp_index: torch.Tensor, + state_commit_batches=(), ) -> Tuple[torch.Tensor, torch.Tensor]: """Verify target tokens and update recurrent MTP state when required.""" - # 保留原调用接口中的步序参数；当前状态提交直接使用验证返回的接受长度。 accept_lengths, accepted_index = mtp_verify( req_to_next_token_ids=backend.model.req_manager.req_sampling_params_manager.req_to_next_token_ids, @@ -32,8 +34,19 @@ def verify_mtp_tokens( new_next_token_ids=next_token_ids, b_req_idx=b_req_idx, ) + context = MtpCommitContext( + b_req_mtp_start_loc=b_req_mtp_start_loc, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + accepted_index=accepted_index, + verify_width=backend.max_draft_step + 1, + batches=state_commit_batches, + ) 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.long()], + accept_lengths, + role="main", + context=context, ) return accept_lengths, accepted_index diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py new file mode 100644 index 0000000000..2ce999257e --- /dev/null +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -0,0 +1,150 @@ +"""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.attention_state.bounded_att.ssm.replay import ReplaySSMState +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs + + +def workload(args, batch, width, dtype, mode, acceptance): + layers, heads, dim = args.layers, args.heads, 128 + shape = (1, batch * width, heads, dim) if width > 1 else (batch, 1, heads, dim) + packed = torch.randn(batch * width, heads * dim * 3, dtype=torch.bfloat16, device="cuda") + q, k, v = [x.view(shape) for x in packed.split(heads * dim, -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") + cu = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * width 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 == "compact": + ssm_backend = CompactSSMState(state, width, q.dtype, kda=True) + elif mode != "native": + ssm_backend = ReplaySSMState(state, 8, width, q.dtype, kda=True, compute_mode=mode) + if ssm_backend is not None: + ssm_backend._config_is_fixed = True + 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 + + def cycle(): + for _ in range(args.rounds): + if ssm_backend is not None: + ssm_backend.prepare_decode(reqs, cu) + for layer in range(layers): + if ssm_backend is None: + fused_recurrent_kda( + q, + k, + v, + gate.view(*shape[:2], -1), + beta.view(*shape[:2], heads), + alog, + bias, + state[layer], + indexes, + cu_seqlens=cu, + num_accepted_tokens=counts if cu is not None else None, + ) + else: + ssm_backend.forward(layer, q, k, v, gate, beta, alog, bias, reqs, cu) + if ssm_backend is not None: + ssm_backend.accept_updates(reqs, accepted) + if isinstance(ssm_backend, ReplaySSMState): + ssm_backend._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) + ssm_backend.merge_accepted_updates(reqs) + else: + snapshot.copy_(state[:, native_index] if ssm_backend is None else state[:, 0]) + + return ssm_backend, cycle, bytes_used + + +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", "compact", "inline", "precompute"] if width > 1 else ["native", "inline"]: + ssm_backend, cycle, size = 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}]: + ssm_backend.run_config = config + cycle() + ms = triton.testing.do_bench_cudagraph(cycle, rep=50) + candidates.append(dict(config=config, ms=ms)) + ssm_backend.run_config = 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, + state_history_bytes=size, + ms_per_round=statistics.median(samples), + samples_ms=samples, + config=ssm_backend.run_config if ssm_backend 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_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index 15d29bd85d..4580453297 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", "compact", "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,10 @@ def test_hybrid_layer_mapping_and_mtp_states( model_logical_layer_ids=logical_layer_ids, device="meta", ) + assert bounded_state.ssm_slots_per_req == (3 if mode == "native" else 1) + assert bounded_state._ssm.state is bounded_state.ssm_state + if mode == "replay": + assert bounded_state._ssm.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_mtp_state_commit.py b/unit_tests/common/att_resources/test_mtp_state_commit.py index a50ff4fc71..5d4d9b8c23 100644 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -80,6 +80,7 @@ def test_commit_without_recurrent_state_and_invalid_role(): def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[int, ...]): manager = make_manager("cuda") backend = SimpleNamespace( + max_draft_step=2, model=SimpleNamespace( att_state_manager=manager, req_manager=SimpleNamespace( @@ -87,7 +88,7 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i req_to_next_token_ids=torch.tensor([[-1, 10, 11]] * 3, device="cuda") ) ), - ) + ), ) reqs, starts, tokens, expected_mask = [], [], [], [] for req, width, accepted in zip((2, 0, 1), widths, (3, 1, 2)): @@ -107,3 +108,37 @@ def test_verify_commits_main_states_for_fixed_and_dynamic_widths(widths: tuple[i assert manager.main_kda_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] assert manager.main_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 1, 2] assert manager.draft0_gdn_conv_fp32_ssm_fp32_state.mtp_state_index.tolist() == [0, 0, 0] + + +def test_verify_routes_one_generic_commit_context(monkeypatch): + from lightllm.server.router.model_infer.mtp_speculative import utils + from lightllm.common.attention_state.commit import StateCommitBatch + + reqs = torch.tensor([2, 2, 0, 0], dtype=torch.int32) + starts = torch.tensor([0, 2], dtype=torch.int32) + positions = torch.tensor([0, 1, 0, 1], dtype=torch.int32) + lengths = torch.tensor([2, 1], dtype=torch.int32) + mask = torch.tensor([1, 1, 1, 0], dtype=torch.int32) + monkeypatch.setattr(utils, "mtp_verify", lambda **kwargs: (lengths, mask)) + calls = [] + resources = SimpleNamespace(commit_bounded_att_state=lambda ids, steps, **kw: calls.append((ids, steps, kw))) + backend = SimpleNamespace( + max_draft_step=1, + model=SimpleNamespace( + att_state_manager=resources, + req_manager=SimpleNamespace(req_sampling_params_manager=SimpleNamespace(req_to_next_token_ids=None)), + ), + ) + batches = (StateCommitBatch(2),) + utils.verify_mtp_tokens(backend, torch.zeros(4), reqs, starts, positions, state_commit_batches=batches) + assert len(calls) == 1 + ids, steps, kwargs = calls[0] + assert ids.tolist() == [2, 0] + assert steps is lengths + assert kwargs["role"] == "main" + context = kwargs["context"] + assert context.batches is batches + assert context.b_req_idx is reqs + assert context.b_mtp_index is positions + assert context.accepted_index is mask + assert context.verify_width == 2 diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 8ec77c6f4d..f9b5b39f76 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -10,6 +10,7 @@ from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState from lightllm.common.basemodel.infer_struct import InferStateInfo import lightllm.common.basemodel.attention.linear.gdn as gdn +import lightllm.common.basemodel.attention.linear.state as linear_state import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils import lightllm.common.basemodel.triton_kernel.linear_att.fla.ops as fla_ops from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend @@ -38,7 +39,7 @@ def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): infer_state = InferStateInfo() infer_state.b1_cu_q_seq_len = torch.tensor([0, 1], dtype=torch.int32) infer_state.b_ready_cache_len = 0 - state = gdn.LinearAttPrefillAttState( + state = gdn.GDNLinearAttPrefillAttState( backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ) @@ -58,6 +59,7 @@ def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): torch.zeros((1, 1), dtype=cache_dtype), state.infer_state, layer_weight, + (state.b_conv_buffer_idx, state.b_ssm_buffer_idx), ) assert ssm_states.dtype == cache_dtype @@ -69,11 +71,12 @@ def test_decode_state_initializes_normal_layout(): state, infer_state = _create_decode_state(draft_step=0, dynamic_layout=False, b_req_idx=b_req_idx) state.init_state() + forward = state.get_forward(3) - assert state.b_conv_buffer_idx is b_req_idx - assert state.b_ssm_buffer_idx is b_req_idx - assert state.b1_mtp_cu_q_seq_len is None - assert state.b_num_accepted_tokens is None + assert forward.req_indices is b_req_idx + assert forward.state_indices is b_req_idx + assert forward.cu_seqlens is None + assert forward.accepted_token_counts is None @pytest.mark.parametrize("role", ["main", "draft0"]) @@ -89,12 +92,13 @@ def test_decode_state_initializes_fixed_mtp_layout(role: str): ) state.init_state() + forward = state.get_forward(3) - torch.testing.assert_close(state.b1_mtp_cu_q_seq_len, torch.tensor([0, 3, 6], dtype=torch.int32)) - torch.testing.assert_close(state.b_conv_buffer_idx, torch.tensor([2, 4], dtype=torch.int32)) - torch.testing.assert_close(state.b_num_accepted_tokens, torch.tensor([2, 3], dtype=torch.int32)) + torch.testing.assert_close(forward.cu_seqlens, torch.tensor([0, 3, 6], dtype=torch.int32)) + torch.testing.assert_close(forward.req_indices, torch.tensor([2, 4], dtype=torch.int32)) + torch.testing.assert_close(forward.accepted_token_counts, torch.tensor([2, 3], dtype=torch.int32)) torch.testing.assert_close( - state.b_ssm_buffer_idx, + forward.state_indices, torch.tensor([[6, 7, 8], [12, 13, 14]], dtype=torch.int32), ) @@ -110,9 +114,9 @@ def test_decode_state_initializes_dynamic_mtp_layout(monkeypatch, role: str): def build_params(**kwargs): build_calls.append(kwargs) - return expected_cu_q_seq_len, expected_conv_buffer_idx, expected_num_accepted_tokens + return expected_cu_q_seq_len, expected_conv_buffer_idx - monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) + monkeypatch.setattr(linear_state, "build_dynamic_mtp_linear_att_state_params", build_params) state, infer_state = _create_decode_state( draft_step=2, dynamic_layout=True, @@ -122,12 +126,13 @@ def build_params(**kwargs): ) state.init_state() + forward = state.get_forward(3) - assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len - assert state.b_conv_buffer_idx is expected_conv_buffer_idx - assert state.b_num_accepted_tokens is expected_num_accepted_tokens + assert forward.cu_seqlens is expected_cu_q_seq_len + assert forward.req_indices is expected_conv_buffer_idx + torch.testing.assert_close(forward.accepted_token_counts, expected_num_accepted_tokens) torch.testing.assert_close( - state.b_ssm_buffer_idx, + forward.state_indices, torch.tensor( [[6, 7, 8], [12, 13, 14], [15, 16, 17], [15, 16, 17], [15, 16, 17]], dtype=torch.int32, @@ -136,12 +141,11 @@ def build_params(**kwargs): assert len(build_calls) == 1 assert build_calls[0]["b_req_idx"] is state.infer_state.b_req_idx assert build_calls[0]["b_mtp_index"] is state.infer_state.b_mtp_index - assert build_calls[0]["req_to_mtp_state_index"] is req_to_mtp_state_index assert build_calls[0]["hold_req_id"] == 5 def test_linear_backend_is_abstract_and_triton_supplies_chunk_kernel(): - assert inspect.isabstract(gdn.LinearAttBackend) + assert inspect.isabstract(gdn.GDNLinearAttBackend) backend = object.__new__(TritonLinearAttBackend) assert backend.get_prefill_kernel() is fla_ops.chunk_gated_delta_rule @@ -358,6 +362,8 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st HOLD_REQUEST_ID=req_to_mtp_state_index.shape[0] - 1 if req_to_mtp_state_index is not None else -1, ) infer_state.att_state_role = role + if req_to_mtp_state_index is None: + req_to_mtp_state_index = torch.zeros(int(b_req_idx.max()) + 2, dtype=torch.int32) if req_to_mtp_state_index is not None: request_capacity = req_to_mtp_state_index.numel() resources = AttentionStateManager(request_capacity=request_capacity) @@ -379,7 +385,7 @@ def _create_decode_state(*, draft_step, dynamic_layout, b_req_idx, req_to_mtp_st resources.register_state(model_role, linear) resources.finish_state_registration() infer_state.att_state_manager = resources - state = gdn.LinearAttDecodeAttState( + state = gdn.GDNLinearAttDecodeAttState( backend=SimpleNamespace( model=model, uses_dynamic_spec_verify_layout=lambda: dynamic_layout, diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py index 52b1c1cf2d..ee76288d8d 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ b/unit_tests/common/basemodel/test_cuda_graph_autotune.py @@ -1,5 +1,6 @@ from contextlib import contextmanager from types import SimpleNamespace +import weakref import pytest import torch @@ -128,3 +129,73 @@ def decode_func(*infer_states): assert not Autotuner.is_autotune_warmup() assert capture_env.capture_count == 0 assert graph.graph == {} + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("kda", [False, True]) +def test_linear_config_is_local_to_capture_state(capture_env, overlap, kda): + from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttDecodeAttState + from lightllm.common.basemodel.attention.linear.kda import KDALinearAttDecodeAttState + from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.batch_objs import PostLayerOutput + from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState + from lightllm.common.ssm_types import SSMRunConfig + from lightllm.common.attention_state.bounded_att import GDNState, KDAState + + capture_env.args.enable_decode_microbatch_overlap = overlap + graph = CudaGraph(1, 2, 1, max_batch_size=2) + from unit_tests.common.test_linear_state import make_state + from lightllm.common.attention_state.commit import StateCommit + + manager, resources, _ = make_state(KDAState if kda else GDNState, mode="compact", device="cpu") + att_class = KDALinearAttDecodeAttState if kda else GDNLinearAttDecodeAttState + states = [] + configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] + for _ in range(2 if overlap else 1): + state = InferStateInfo() + state.input_ids = torch.ones(2, dtype=torch.int64) + state.att_state_manager = resources + state.hidden_collector = NoopHiddenCollector() + state.decode_att_state = TritonDecodeAttState(_infer_state_weak_ref=weakref.ref(state)) + state.decode_att_state1 = att_class(_infer_state_weak_ref=weakref.ref(state)) + state.decode_att_state1.forwards[manager.state_name] = manager.prepare_decode( + torch.tensor([0, 1], dtype=torch.int32), + torch.tensor([0, 3, 6], dtype=torch.int32), + ) + states.append(state) + + def decode_func(*infer_states): + outputs = [] + for index, state in enumerate(infer_states): + att = state.decode_att_state1 + assert att.infer_state is state + forward = att.get_forward(0) + assert forward.run_config is None + assert (att is states[index].decode_att_state1) == capture_env.capturing + forward.run_config = configs[index] if capture_env.capturing else SSMRunConfig(128, 4) + outputs.append(TpPartBaseModel._create_model_output(None, PostLayerOutput(torch.ones(2, 1)), state)) + return tuple(outputs) if overlap else outputs[0] + + result = graph.capture_decode(decode_func, *states) + outputs = result if overlap else (result,) + runtime_states = [cuda_graph_module._copy_warmup_state(state) for state in states] + for state in runtime_states: + state.decode_att_state1.forwards[manager.state_name] = manager.prepare_decode( + torch.tensor([1, 0], dtype=torch.int32), + torch.tensor([0, 1, 4], dtype=torch.int32), + ) + replay = graph.replay(*runtime_states) + assert replay == result + for index, (state, output) in enumerate(zip(states, outputs)): + forward = state.decode_att_state1.get_forward(0) + assert forward.run_config is configs[index] + assert forward.req_indices.tolist() == [1, 0] + assert forward.cu_seqlens.tolist() == [0, 1, 4] + forward.run_config = SSMRunConfig(64, 4) + assert output.state_commits == (StateCommit(manager.state_name, configs[index]),) + unpadded = TpPartBaseModel._create_unpad_decode_model_output(None, output, 1) + assert unpadded.logits.shape[0] == 1 + assert unpadded.state_commits is output.state_commits 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..4694b9fbe6 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -0,0 +1,189 @@ +import pytest +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from types import SimpleNamespace +from lightllm.common.attention_state.commit import MtpCommitContext +from lightllm.common.attention_state.bounded_att.ssm.native import NativeSSMState + + +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 = ReplaySSMState.__new__(ReplaySSMState) + ssm_backend.cursors = tensor([7] * (hold + 1)) + ssm_backend.hold_req_id = hold + ssm_backend.capacity = max(4, triton.next_power_of_2(width)) + ssm_backend.verify_width = width + manager = SimpleNamespace() + manager.mtp_state_index = tensor([-9] * (hold + 1)) + manager._ssm = ssm_backend + return manager, args, (ids, lengths, starts) + + +def accept_unfused(m, args): + update(m.mtp_state_index, *args) + m._ssm.accept_updates(args[1][args[0].long()], m.mtp_state_index) + + +def accept_fused(m, args): + reqs = args[1][args[0].long()] + m._ssm.commit(reqs, None, m.mtp_state_index, MtpCommitContext(*args), "main") + + +@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.cursors.fill_(7) + ref._ssm.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.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.cursors, ref._ssm.cursors) + assert m.mtp_state_index.tolist() == expected_index + assert m._ssm.cursors.tolist() == expected_cursor + checks += 1 + + m._ssm = NativeSSMState(torch.empty(0), width) + accept_fused(m, args) + assert m.mtp_state_index.tolist() == expected_index + + 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 = [ + ReplaySSMState( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + compute_mode=compute_mode, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for _ in range(2) + ] + managers = [] + for ssm_backend in ssm_backends: + manager = SimpleNamespace() + manager._ssm = 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(1, 12, 2, 32, dtype=torch.bfloat16, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, 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] + ssm_backend.prepare_decode(reqs, cu) + outputs = [ssm_backend.forward(layer, q, k, v, a, b, alog, alog, reqs, cu)[0] 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].state, ssm_backends[1].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_backends[0].get_accepted_state(3), ssm_backends[1].get_accepted_state(3)) + if iteration % 11 == 10: + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs[:3]) + assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) + if iteration == 19: + for manager in managers: + manager._ssm.clear_history(1) + manager._ssm.state[:, 1].zero_() + manager.mtp_state_index[1] = 0 + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs[:3]) + assert torch.equal(ssm_backends[0].state, ssm_backends[1].state) + assert torch.equal(ssm_backends[0].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..c4c35a78cf --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -0,0 +1,68 @@ +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(): + cu, ids = build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, hold) + forward = manager.prepare_decode(ids, cu) + return cu, forward.req_indices, forward.state_indices, forward.accepted_token_counts + + build() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + cu, ids, ssm, counts = build() + assert cu.numel() == sequences + 1 + assert ids.numel() == counts.numel() == sequences + + # 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] * (sequences + 1 - len(expected_cu)) + expected_ids = request_ids + [hold] * (sequences - len(request_ids)) + assert cu.tolist() == expected_cu + assert ids.tolist() == expected_ids + expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (sequences - len(request_ids)) + assert counts.tolist() == expected_counts + assert ssm.tolist() == [[req * 4 + offset for offset in range(4)] for req in expected_ids] 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..db1ab6d98e --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -0,0 +1,662 @@ +import pytest +import torch + +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs + +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"]) +def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state_dtype, kda): + state = torch.zeros(1, 2, 1, 128, 128, dtype=state_dtype, device="cuda") + ssm_backend = ReplaySSMState( + state, + capacity, + 2, + torch.float16, + num_key_heads=1, + kda=kda, + compute_mode="precompute", + run_config={"BV": bv, "num_warps": 1}, + ) + 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(1, 4, 1, 128, dtype=torch.float16, device="cuda") + q[0, :2, 0, 0] = 1 + k = torch.zeros_like(q) + k[0, 0, 0, 0] = -0.25 + k[0, 1, 0, 0] = 1 + v = torch.zeros_like(q) + v[0, :2, 0, 0] = 60000 + a = torch.full((4, 128 if kda else 1), -20, dtype=torch.float16, device="cuda") + b = torch.full((4, 1), 20, dtype=torch.float16, 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): + ssm_backend.prepare_decode(reqs, cu) + # 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 derived correction by reconstructing from raw inputs. + assert ssm_backend.cursors[0].item() == 2 * round_idx + assert not state.any() + actual = ssm_backend.forward(0, q, k, v, a, b, alog, bias, reqs, cu)[0] + expected = [] + for token in range(2): + output, expected_state = reference( + q[0, token], k[0, token], v[0, 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[0, :2].float(), torch.stack(expected), atol=4, rtol=1e-3) + ssm_backend.accept_updates(reqs, accepted) + # 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_backend.get_accepted_state(0)[0].float(), expected_state, atol=128, rtol=0.01) + + +@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 = ReplaySSMState( + state, + capacity, + width, + num_key_heads=h, + kda=kda, + gate_lower_bound=gate_lower_bound, + compute_mode=compute_mode, + run_config=run_config, + 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(1, tokens, h, kdim, device="cuda", dtype=activation_dtype) + k = torch.randn_like(q) + v = torch.randn(1, 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(): + ssm_backend.prepare_decode(reqs, cu) + out = [ssm_backend.forward(layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + ssm_backend.accept_updates(reqs, accepted) + 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]): + if pending[seq] + (width if seq == 0 else max(width - 1, 1)) > capacity: + 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[0, t], + k[0, t], + v[0, t], + a[t], + b[t], + alog, + bias, + cur, + kda=kda, + gate_lower_bound=gate_lower_bound, + ) + torch.testing.assert_close(outputs[layer][0, 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.merge_accepted_updates(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.merge_accepted_updates(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}, + 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} + ssm_backends = [ + ReplaySSMState( + state.clone(), + 16 if width == 1 else 8, + width, + dtype, + num_key_heads=2, + compute_mode=mode, + run_config=config, + ) + 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((1, tokens, 2, 128), device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn((1, 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 = [] + for ssm_backend in ssm_backends: + ssm_backend.prepare_decode(reqs, cu) + outputs.append(ssm_backend.forward(0, q, k, v, a, b, alog, alog, reqs, cu)[0]) + ssm_backend.accept_updates(reqs, accepted) + 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.merge_accepted_updates(reqs) + torch.testing.assert_close(ssm_backends[0].state, ssm_backends[1].state, atol=0, rtol=0) + + +def test_bf16_decode_fold_uses_rounded_checkpoint(): + # Automatic fold and explicit materialization must start from the same BF16 + # checkpoint, including the token computed inside the automatic-fold kernel. + state = torch.full((1, 2, 1, 32, 64), 0.123, device="cuda", dtype=torch.bfloat16) + ssm_backends = [ReplaySSMState(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_backends[1].merge_accepted_updates(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 = [] + for ssm_backend in ssm_backends: + ssm_backend.prepare_decode(reqs) + ssm_backend.accept_updates(reqs) + outputs.append(ssm_backend.forward(0, q, k, v, a, b, alog, bias, reqs)[0]) + torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) + torch.testing.assert_close(ssm_backends[0].state, ssm_backends[1].state, atol=0, rtol=0) + + +@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(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 = ReplaySSMState( + state, 8, 4, dtype, num_key_heads=2, compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1} + ) + expected = state.float().clone() + pending, phases = [0, 0], [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(1, 8, 2, 128, dtype=dtype, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, 8, 4, 128, dtype=dtype, device="cuda") + a = torch.full((8, 4), -3.0, dtype=dtype, device="cuda") + b = torch.randn_like(a) + alog = torch.zeros(4, device="cuda") + + def forward(): + return [ssm_backend.forward(layer, q, k, v, a, b, alog, alog, reqs, cu)[0] for layer in range(2)] + + def step(): + positions = ssm_backend.prepare_decode(reqs, cu) + out = forward() + ssm_backend.accept_updates(reqs, accepted) + 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() + old_cursors = [phases[req] + pending[req] for req in ids] + [0] + for req, width in zip(ids, widths): + if pending[req] + width > 8: + expected[:, req] = expected[:, req].to(dtype).float() + pending[req] = 0 + phases[req] ^= 16 + if not capture_metadata: + positions.copy_(ssm_backend.prepare_decode(reqs, cu)) + graph.replay() + if not capture_metadata: + ssm_backend.accept_updates(reqs, accepted) + 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[0, t], k[0, t], v[0, t], a[t], b[t], alog, alog, current) + torch.testing.assert_close(outputs[layer][0, 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() == [phases[i] + pending[i] for i in range(2)] + [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 and phases[0] == 0 + torch.testing.assert_close(state[:, 0], before[:, 0], atol=0, rtol=0) + ssm_backend.merge_accepted_updates(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 = ReplaySSMState(state, 4, 3) + snapshot = ssm_backend.get_accepted_state(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]) +def test_compact_preserves_accepted_prefix(dtype, run_config=None, kda=False): + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState + + torch.manual_seed(17) + hv, kd, vd, width = 4, 128 if kda else 64, 128 if kda else 64, 3 + state = torch.randn(2, 4, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 + ssm_backend = CompactSSMState(state, width, torch.bfloat16, num_key_heads=hv, kda=kda, run_config=run_config) + reqs = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, 3, 5, 5], device="cuda", dtype=torch.int32) + accepted = torch.tensor([1, 0, 0, 0], device="cuda", dtype=torch.int32) + shape = (1, 5, hv, kd) + q = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.randn((5, hv * kd if kda else hv), device="cuda", dtype=torch.bfloat16) + b = torch.randn(5, hv, device="cuda", dtype=torch.bfloat16) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv * kd if kda else hv, device="cuda") * 0.1 + for _ in range(15): + before = state.clone() + outputs = [ssm_backend.forward(layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + assert torch.equal(state, before), "verify must not mutate the checkpoint" + ssm_backend.accept_updates(reqs, accepted) + for layer in range(2): + for req, start, length in [(2, 0, 3), (0, 3, 2)]: + cur = before[layer, req].float() + for j in range(length): + t = start + j + out, cur = reference(q[0, t], k[0, t], v[0, t], a[t], b[t], alog, bias, cur, kda=kda) + torch.testing.assert_close(outputs[layer][0, t].float(), out, atol=0.004, rtol=0.02) + cur = cur.to(dtype).float() + if j == int(accepted[req]): + torch.testing.assert_close( + state[layer, req].float(), cur, atol=0.008 if dtype == torch.bfloat16 else 2e-5, rtol=0.02 + ) + torch.testing.assert_close(state[:, 1], before[:, 1], rtol=0, atol=0) + torch.testing.assert_close(state[:, 3], before[:, 3], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("width", [3, 4]) +def test_compact_matches_native_mtp(dtype, width, run_config=None): + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState + from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, + ) + + torch.manual_seed(77) + batch, h, hv, kd, vd = 3, 2, 4, 128, 128 + q = torch.randn(1, batch * width, h, kd, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, batch * width, hv, vd, device="cuda", dtype=torch.bfloat16) + a = torch.randn(batch * width, hv, device="cuda", dtype=torch.bfloat16) + b = torch.randn(batch * width, hv, device="cuda", dtype=torch.bfloat16) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv, device="cuda") * 0.1 + state = torch.randn(1, batch + 1, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 + reference_state = state[0].repeat_interleave(width, 0) + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) + cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * width + idx = torch.arange(batch * width, device="cuda", dtype=torch.int32).view(batch, width) + counts = torch.ones(batch, device="cuda", dtype=torch.int32) + accepted = torch.tensor([0, width // 2, width - 1, 0], device="cuda", dtype=torch.int32) + ssm_backend = CompactSSMState(state, width, torch.bfloat16, num_key_heads=h, run_config=run_config) + for _ in range(32): + q.normal_() + k.normal_() + v.normal_() + a.normal_().sub_(3) + b.normal_() + out = ssm_backend.forward(0, q, k, v, a, b, alog, bias, reqs, cu)[0] + ref, _ = mtp_fused_recurrent_gated_delta_rule( + q, + k, + v, + reference_state, + cu, + idx, + idx, + counts, + alog, + bias, + a, + b, + run_config={"num_stages": 1, **(run_config or {"BV": 8, "num_warps": 1})}, + ) + ssm_backend.accept_updates(reqs, accepted) + # BF16 casts can straddle a rounding midpoint across compiler/layouts. + rtol = 2 * torch.finfo(dtype).eps if dtype == torch.bfloat16 else 0 + torch.testing.assert_close(out, ref, rtol=rtol, atol=2e-5 if dtype == torch.bfloat16 else 0) + torch.testing.assert_close( + state[0, :batch], + reference_state[reqs * width + accepted[:batch]], + rtol=rtol, + atol=5e-4 if dtype == torch.bfloat16 else 0, + ) + counts.copy_(accepted[:batch] + 1) + accepted[:batch].add_(1).remainder_(width) + + +@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']}") +@pytest.mark.parametrize("mode", ["gdn", "replay"]) +def test_autotune_candidates_match_recurrence(dtype, run_config, mode): + if mode == "gdn": + test_compact_matches_native_mtp(dtype, 4, run_config) + else: + 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_() + backend = state._ssm + backend.raw_keys.normal_() + backend.raw_values.normal_() + backend.gates.uniform_(-0.1, -0.01) + backend.betas.uniform_(0.1, 0.9) + active_before = state.ssm_state[:, 1].clone() + conv_expected = state.conv_state[:, 1, :, 1:4].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): + backend.cursors[1] = count + expected = active_before.float().clone() + for i in range(count): + key = backend.raw_keys[:, 1, :, i].float() + key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() + expected *= backend.gates[:, 1, :, i].exp()[..., None, None] + delta = backend.betas[:, 1, :, i, None] * ( + backend.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) + ) + expected += key[..., None] * delta[..., None, :] + materialized = backend.get_accepted_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], materialized.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 backend.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), + ): + backend.cursors[2] = 3 + state.mtp_state_index[2] = 2 + restore() + assert backend.cursors[2].item() == state.mtp_state_index[2].item() == 0 + torch.testing.assert_close(state.ssm_state[:, 2], materialized, 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("dtype", [torch.float32, torch.bfloat16]) +def test_kda_compact_reference_and_partial_acceptance(dtype): + test_compact_preserves_accepted_prefix(dtype, kda=True) + + +@pytest.mark.parametrize("config", get_configs()) +@pytest.mark.parametrize("mode", ["compact", "inline", "precompute"]) +def test_kda_tuning_candidates_match_reference(config, mode): + if mode == "compact": + test_compact_preserves_accepted_prefix(torch.bfloat16, run_config=config, kda=True) + else: + test_replay_acceptance_flush_and_graph( + 3, + (128, 128), + 4, + torch.bfloat16, + run_config=config, + compute_mode=mode, + head_geometry=(2, 2), + kda=True, + gate_shift=-3, + ) 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..6da95389ce --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -0,0 +1,333 @@ +import pytest +import torch + +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +from types import SimpleNamespace +from lightllm.common.attention_state.commit import MtpCommitContext + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@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", + run_config=config, + ) + for cls in [ReplaySSMState, ReplaySSMState] + ] + manager = SimpleNamespace() + manager._ssm = 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(1, 12, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 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 = build_dynamic_mtp_linear_att_state_params(rows, mtp, 4) + return cu, reqs, ssm_backend.prepare_decode(reqs, cu) + + def forward_accept(index, cu, reqs): + ssm_backend = ssm_backends[index] + out = [ssm_backend.forward(layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(3)] + if index == 0: + ssm_backend.accept_updates(reqs, accepted) + else: + manager._ssm.commit( + rows[starts.long()], + None, + manager.mtp_state_index, + MtpCommitContext(starts, rows, mtp, flags, 4), + "main", + ) + 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.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.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_() + 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) + positions = captured[0][2].cpu().tolist() + folds += sum(positions[j] % 16 + length > 8 for j, length in enumerate(lengths)) + 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].state.float(), ssm_backends[0].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 = [ + ReplaySSMState( + initial.clone(), + 8, + 1, + dtype, + num_key_heads=2, + kda=kda, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + 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") + + def forward(ssm_backend): + return [ssm_backend.forward(layer, q, k, v, a, beta, alog, bias, reqs)[0] for layer in range(3)] + + def warm(ssm_backend): + ssm_backend.prepare_decode(reqs) + out = forward(ssm_backend) + ssm_backend.accept_updates(reqs) + return out + + graphs, positions, outputs = [], [], [] + for early, ssm_backend in enumerate(ssm_backends): + warm(ssm_backend) + ssm_backend.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.prepare_decode(reqs) + if early: + ssm_backend.accept_updates(reqs) + out = forward(ssm_backend) + if capture_metadata and not early: + ssm_backend.accept_updates(reqs) + graphs.append(graph) + positions.append(pos) + outputs.append(out) + saw_fold = saw_nonfold = False + 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_() + old = ssm_backends[0].cursors[reqs.long()].clone() + for early, (ssm_backend, graph, pos) in enumerate(zip(ssm_backends, graphs, positions)): + if not capture_metadata: + pos.copy_(ssm_backend.prepare_decode(reqs)) + if early: + ssm_backend.accept_updates(reqs) + graph.replay() + if not capture_metadata and not early: + ssm_backend.accept_updates(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].state.float(), ssm_backends[0].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps + ) + folded = any(n % 16 + 1 > 8 for n in old[:2].cpu().tolist()) + saw_fold |= folded + saw_nonfold |= not folded + assert saw_fold and saw_nonfold + + +def test_cross_layer_fold_with_state_offset_above_int32(): + 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 = [ + ReplaySSMState(state, 8, 4, torch.bfloat16, num_key_heads=4, run_config=config) for state in [dense, sparse] + ] + q = torch.randn(1, 8, 4, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 8, 12, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((8, 12), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn_like(a) + alog = torch.zeros(12, device="cuda") + bias = torch.zeros_like(alog) + reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) + cu = reqs.new_tensor([0, 4, 8]) + + def step(ssm_backend, accepted): + pos = ssm_backend.prepare_decode(reqs, cu) + out = [ssm_backend.forward(layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(48)] + ssm_backend.accept_updates(reqs, accepted) + 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 = [ + ReplaySSMState(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute", run_config=config) + 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(1, batch * 4, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 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(): + ssm_backend.prepare_decode(ids, offsets) + out = [ssm_backend.forward(layer, q, k, v, a, beta, alog, bias, ids, offsets)[0] for layer in range(2)] + ssm_backend.accept_updates(ids, accepted) + return out + + step() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = step() + graphs.append(graph) + outputs.append(out) + ssm_backend.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_() + 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].state.float(), ssm_backends[1].state.float(), atol=0.002, rtol=torch.finfo(dtype).eps + ) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py new file mode 100644 index 0000000000..78f80de800 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_grouped_keys.py @@ -0,0 +1,81 @@ +"""The storage optimization must be bit-identical to the redundant-key layout.""" + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("mode", ["replay", "compact"]) +@pytest.mark.parametrize("group", [1, 2, 3]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_grouped_keys_graph_fold_partial_accept(mode, group, dtype): + torch.manual_seed(192) + layers, slots, h, kd, vd, width = 2, 5, 2, 128, 128, 4 + hv = h * group + state = torch.randn(layers, slots, hv, kd, vd, device="cuda", dtype=dtype) * 0.01 + + def make_ssm_backend(key_heads): + if mode == "replay": + return ReplaySSMState(state.clone(), 16, width, num_key_heads=key_heads) + return CompactSSMState(state.clone(), width, torch.bfloat16, num_key_heads=key_heads) + + ssm_backends = [make_ssm_backend(hv), make_ssm_backend(h)] + assert ssm_backends[0].keys.numel() == ssm_backends[1].keys.numel() * group + # Packed strides, out-of-order requests, an empty real request, nonempty HOLD. + packed = torch.randn(1, 8, (2 * h + hv) * kd, device="cuda", dtype=torch.bfloat16) + q, k, v = [x.view(1, 8, heads, kd) for x, heads in zip(packed.split([h * kd, h * kd, hv * kd], -1), [h, h, hv])] + a = torch.randn(8, hv, device="cuda", dtype=torch.bfloat16) + b = torch.randn(8, hv, device="cuda", dtype=torch.bfloat16) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(a.shape[-1], device="cuda") * 0.1 + reqs = torch.tensor([2, 0, 1, 4], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, 4, 7, 7, 8], device="cuda", dtype=torch.int32) + accepted = torch.tensor([0, -1, 0, -1, -1], device="cuda", dtype=torch.int32) + + def step(ssm_backend): + ssm_backend.prepare_decode(reqs, cu) + out = [ssm_backend.forward(layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(layers)] + ssm_backend.accept_updates(reqs, accepted) + return out + + graphs, outputs = [], [] + for ssm_backend in ssm_backends: + step(ssm_backend) + ssm_backend.state.copy_(state) + if mode == "replay": + ssm_backend.cursors.zero_() + reqs.fill_(4) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs.append(step(ssm_backend)) + graphs.append(graph) + reqs.copy_(torch.tensor([2, 0, 1, 4], device="cuda", dtype=torch.int32)) + + for iteration in range(67): + packed.normal_() + a.normal_().sub_(3) + b.normal_() + accepted[0] = iteration % 3 + accepted[2] = (iteration + 1) % 4 + for graph in graphs: + graph.replay() + for left, right in zip(*outputs): + torch.testing.assert_close(left, right, atol=0, rtol=0) + assert torch.count_nonzero(right[:, -1]) == 0 + if iteration in (18, 39, 66): + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(reqs) + torch.testing.assert_close(ssm_backends[0].state, ssm_backends[1].state, atol=0, rtol=0) + torch.testing.assert_close(ssm_backends[1].state[:, [1, 3, 4]], state[:, [1, 3, 4]], atol=0, rtol=0) + if mode == "replay": + torch.testing.assert_close(ssm_backends[0].cursors, ssm_backends[1].cursors, atol=0, rtol=0) + if iteration == 40: + # Request-slot reuse / checkpoint restore must not resurrect history. + for ssm_backend in ssm_backends: + ssm_backend.clear_history(2) + ssm_backend.state[:, 2].copy_(state[:, 2]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py new file mode 100644 index 0000000000..25bcb15126 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_parallel.py @@ -0,0 +1,172 @@ +import pytest +import torch + +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("capacity", [8, 16, 64]) +def test_derived_fold_commits_only_accepted_history(capacity): + torch.manual_seed(731) + state = torch.randn(2, 5, 4, 128, 128, device="cuda", dtype=torch.bfloat16) * 0.02 + initial = state.clone() + ssm_backend = ReplaySSMState(state, capacity, 5, num_key_heads=2, compute_mode="parallel", fold_method="derived") + ssm_backend.keys.fill_(float("nan")) + ssm_backend.deltas.fill_(float("nan")) + ssm_backend.gates.fill_(float("nan")) + ssm_backend.raw_keys.fill_(float("nan")) + ssm_backend.raw_values.fill_(float("nan")) + ssm_backend.betas.fill_(float("nan")) + # Fold request 2 from the second ring half; request 0 has an empty span. + count = capacity - 2 + ssm_backend.cursors[2] = 2 * capacity + count + ssm_backend.cursors[0] = count + ssm_backend.keys[:, 2, :, capacity : capacity + count].normal_(0, 0.1) + ssm_backend.deltas[:, 2, :, capacity : capacity + count].normal_(0, 0.1) + ssm_backend.gates[:, 2, :, capacity : capacity + count].fill_(-0.1) + expected = initial[:, 2].float() + for j in range(count): + keys = ssm_backend.keys[:, 2, :, capacity + j].repeat_interleave(2, dim=1) + delta = ssm_backend.deltas[:, 2, :, capacity + j] + decay = ssm_backend.gates[:, 2, :, capacity + j].exp() + expected = expected * decay[:, :, None, None] + keys[:, :, :, None] * delta[:, :, None, :] + reqs = torch.tensor([2, 0, 4], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 5, 5, 10], dtype=torch.int32, device="cuda") + ssm_backend.prepare_decode(reqs, cu) + assert torch.isfinite(state).all() + torch.testing.assert_close(state[:, 2], expected.bfloat16(), atol=2e-4, rtol=0.008) + torch.testing.assert_close(state[:, [0, 1, 3, 4]], initial[:, [0, 1, 3, 4]], atol=0, rtol=0) + assert ssm_backend.cursors[2] == 0 + assert ssm_backend.cursors[0] == count + assert ssm_backend.cursors[4] == 0 + + +@pytest.mark.parametrize("fold", ["raw", "derived"]) +def test_parallel_rejects_unsupported_shapes(fold): + state = torch.zeros(1, 3, 4, 32, 64, device="cuda", dtype=torch.bfloat16) + with pytest.raises(AssertionError, match="128x128"): + ReplaySSMState(state, 16, 5, num_key_heads=2, compute_mode="parallel", fold_method=fold) + + +def test_derived_fold_requires_parallel_solver(): + state = torch.zeros(1, 3, 4, 128, 128, device="cuda", dtype=torch.bfloat16) + with pytest.raises(AssertionError): + ReplaySSMState(state, 16, 5, num_key_heads=2, fold_method="derived") + + +@pytest.mark.parametrize("width", [2, 3, 5, 8, 16]) +def test_parallel_raw_against_sequential_reference(width): + from unit_tests.common.basemodel.triton_kernel.linear_att.test_replayssm import ( + test_replay_acceptance_flush_and_graph as check_replay, + ) + + check_replay( + width=width, + dims=(128, 128), + capacity=16, + dtype=torch.bfloat16, + compute_mode="parallel", + run_config={"BV": 128, "num_warps": 2, "num_stages": 1}, + ) + + +@pytest.mark.parametrize("capacity,width", [(8, 3), (16, 5), (64, 16)]) +@pytest.mark.parametrize("capture", [False, True], ids=["eager", "graph"]) +def test_parallel_derived_continuous_state(capacity, width, capture): + from unit_tests.common.basemodel.triton_kernel.linear_att.test_replayssm import reference + + torch.manual_seed(731) + state = torch.randn(2, 4, 4, 128, 128, device="cuda", dtype=torch.bfloat16) * 0.02 + initial = state.clone() + ssm_backend = ReplaySSMState( + state, capacity, width, num_key_heads=2, compute_mode="parallel", fold_method="derived" + ) + for history in ( + ssm_backend.keys, + ssm_backend.deltas, + ssm_backend.gates, + ssm_backend.raw_keys, + ssm_backend.raw_values, + ssm_backend.betas, + ): + history.fill_(float("nan")) + reqs = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) + lengths = [width, width - 1] + cu = torch.tensor([0, width, 2 * width - 1, 2 * width], device="cuda", dtype=torch.int32) + q = torch.randn(1, 2 * width, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 2 * width, 4, 128, device="cuda", dtype=torch.bfloat16) + a = torch.randn(2 * width, 4, device="cuda", dtype=torch.bfloat16) - 3 + b = torch.randn_like(a) + alog = torch.randn(4, device="cuda") * 0.1 + bias = torch.randn_like(alog) * 0.1 + accepted = torch.zeros(4, device="cuda", dtype=torch.int32) + + def step(): + ssm_backend.prepare_decode(reqs, cu) + outputs = [ssm_backend.forward(layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + ssm_backend.accept_updates(reqs, accepted) + return outputs + + # Compile before capture, then reset the live state and capture only HOLD rows. + step() + state.copy_(initial) + ssm_backend.cursors.zero_() + if capture: + reqs.fill_(ssm_backend.hold_req_id) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = step() + reqs.copy_(torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32)) + + expected = initial.float() + pending = [0, 0] + folds = [0, 0] + for iteration in range(3 * capacity + 5): + for tensor in (q, k, v, b): + tensor.normal_() + a.normal_(-3, 1) + counts = [iteration % width + 1, (iteration + 1) % (width - 1) + 1] + for seq, req in enumerate((2, 0)): + accepted[req] = counts[seq] - 1 + if pending[seq] + lengths[seq] > capacity: + expected[:, req] = expected[:, req].bfloat16().float() + pending[seq] = 0 + folds[seq] += 1 + pending[seq] += counts[seq] + if capture: + graph.replay() + else: + outputs = step() + for layer in range(2): + for seq, req in enumerate((2, 0)): + current = expected[layer, req].clone() + for j in range(lengths[seq]): + token = seq * width + j + output, current = reference( + q[0, token], k[0, token], v[0, token], a[token], b[token], alog, bias, current + ) + torch.testing.assert_close(outputs[layer][0, token].float(), output, atol=0.006, rtol=0.02) + if j + 1 == counts[seq]: + expected[layer, req].copy_(current) + assert torch.count_nonzero(outputs[layer][0, -1]) == 0 + + if iteration == capacity: + saved_state, saved_cursors = state.clone(), ssm_backend.cursors.clone() + snapshot = ssm_backend.get_accepted_state(2) + torch.testing.assert_close(snapshot.float(), expected[:, 2], atol=0.002, rtol=0.02) + torch.testing.assert_close(state, saved_state, atol=0, rtol=0) + torch.testing.assert_close(ssm_backend.cursors, saved_cursors, atol=0, rtol=0) + # Restore the canonical checkpoint while another request keeps its history. + ssm_backend.clear_history(2) + state[:, 2].copy_(snapshot) + expected[:, 2].copy_(snapshot.float()) + pending[0] = 0 + + assert min(folds) >= 2 # Exercise both ring halves, including after restore. + ssm_backend.merge_accepted_updates(reqs[:2]) + torch.testing.assert_close(state.float(), expected.bfloat16().float(), atol=0.002, rtol=0.02) + torch.testing.assert_close(state[:, [1, 3]], initial[:, [1, 3]], atol=0, rtol=0) + assert torch.count_nonzero(ssm_backend.cursors) == 0 diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py new file mode 100644 index 0000000000..22862ea29f --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune.py @@ -0,0 +1,422 @@ +import collections +import math +from dataclasses import asdict + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState +from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState +from lightllm.common.attention_state.commit import MtpCommitContext, StateCommitBatch +from lightllm.common.attention_state.bounded_att.ssm.forward import SSMForward +from lightllm.common.kernel_config import KernelConfigs +from lightllm.common import kernel_config as kernel_config_module +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + + +def make_inputs(mode, dtype=torch.bfloat16, device="cpu", width=4, compute_mode="inline", kda=False): + layers, slots, batch, h, hv, kd, vd = 2, 9, 3, 2, 4, 32, 64 + state = torch.randn(layers, slots, hv, kd, vd, dtype=dtype, device=device) * 0.01 + if mode == "replay": + ssm_backend = ReplaySSMState(state, 8, width, num_key_heads=h, compute_mode=compute_mode, kda=kda) + else: + ssm_backend = CompactSSMState(state, width, torch.bfloat16, num_key_heads=h, kda=kda) + tokens = batch * width + packed = torch.randn(1, tokens, 2 * h * kd + hv * vd, dtype=torch.bfloat16, device=device) + q, k, v = [ + x.view(1, tokens, heads, dim) + for x, heads, dim in zip(packed.split([h * kd, h * kd, hv * vd], -1), [h, h, hv], [kd, kd, vd]) + ] + gates = torch.randn(tokens, hv * (kd + 1) if kda else hv * 2, device=device, dtype=q.dtype) + a, b = gates[:, :-hv], gates[:, -hv:] + return dict( + ssm_backend=ssm_backend, + mode=mode, + q=q, + k=k, + v=v, + a=a, + b=b, + a_log=torch.randn(hv, device=device) * 0.1, + bias=torch.randn(a.shape[-1], device=device) * 0.1, + cu_seqlens=torch.arange(batch + 1, device=device, dtype=torch.int32) * width, + ) + + +def tensors(ssm_backend): + return {name: x for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} + + +def test_replay_projection_has_separate_autotune_cache(): + inline = make_inputs("replay") + precompute = make_inputs("replay", compute_mode="precompute") + assert tuning.select_config._static_key(**inline) != tuning.select_config._static_key(**precompute) + _, rebuilt = tuning.rebuild_inputs(**precompute) + assert rebuilt["ssm_backend"].compute_mode == "precompute" + assert rebuilt["ssm_backend"].keys.dtype == precompute["ssm_backend"].keys.dtype == torch.bfloat16 + assert tuning.select_config._static_key(**precompute) == tuning.select_config._static_key(**rebuilt) + key = tuning.select_config._static_key(**precompute) + legacy = dict(key, v=2) + assert KernelConfigs.get_config_file_name(key) != KernelConfigs.get_config_file_name(legacy) + assert len(KernelConfigs.get_config_file_name(key).encode()) <= 255 + precompute["ssm_backend"].keys = precompute["ssm_backend"].keys.float() + assert tuning.select_config._static_key(**precompute) != key + + +@pytest.fixture(autouse=True) +def isolate_tuner(monkeypatch, tmp_path): + torch.manual_seed(42) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: 1) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(kernel_config_module, "get_current_device_name", lambda: "NVIDIA H200") + tuner = tuning.select_config + monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(tuner, "warmuped_configs_set", set()) + monkeypatch.setattr(tuner, "_oom_keys", set()) + + +@pytest.mark.parametrize("mode", ["compact", "replay"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_rebuild_has_real_requests_and_independent_packed_inputs(mode, dtype, kda): + inputs = make_inputs(mode, dtype, kda=kda) + ssm_backend = inputs["ssm_backend"] + for x in tensors(ssm_backend).values(): + x.fill_(1) + args, rebuilt = tuning.rebuild_inputs(**inputs) + assert not args + scratch = rebuilt["ssm_backend"] + assert scratch.state.shape[:2] == (1, 4) # One layer, active batch + HOLD. + assert scratch._config_is_fixed + assert rebuilt["workload"][0].tolist() == [0, 1, 2] + assert rebuilt["workload"][1].tolist() == [3, 3, 3, 0] + assert rebuilt["workload"][2].tolist() == [0, 1, 2, 0] + for name in ("q", "k", "v", "a", "b", "a_log", "bias"): + assert rebuilt[name].stride() == inputs[name].stride() + assert rebuilt[name].data_ptr() != inputs[name].data_ptr() + assert tuning.select_config._static_key(**inputs) == tuning.select_config._static_key(**rebuilt) + assert tuning.select_config._run_key(**inputs) == tuning.select_config._run_key(**rebuilt) + for x in tensors(scratch).values(): + x.zero_() + for x in tensors(ssm_backend).values(): + assert torch.all(x == 1) + key = tuning.select_config._static_key(**inputs) + assert len(KernelConfigs.get_config_file_name(key).encode()) <= 255 + other = make_inputs(mode, torch.float32 if dtype == torch.bfloat16 else torch.bfloat16, kda=kda) + assert tuning.select_config._static_key(**other) != key + other = make_inputs(mode, dtype, width=3, kda=kda) + assert tuning.select_config._static_key(**other) != key + inputs["q"] = inputs["q"].contiguous() + assert tuning.select_config._static_key(**inputs) != key + + +@pytest.mark.parametrize("level", [0, 1, 2, 3]) +def test_layout_is_fixed_before_graph_capture_and_cached_configs_do_not_execute(monkeypatch, level): + inputs = make_inputs("compact") + ssm_backend = inputs["ssm_backend"] + config = {"BV": 16, "num_warps": 1, "num_stages": 2} + tuner = tuning.select_config + key = autotuner_module.frozendict(tuner._static_key(**inputs)) + tuner.cached_configs[key] = {str(tuner._run_key(**inputs)): config} + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + monkeypatch.setattr(tuner, "_bench", lambda *args, **kwargs: pytest.fail("cached layout must not benchmark")) + default = ssm_backend.run_config.copy() + with Autotuner.autotune_warmup(AutotuneKernelType.GENERAL): + tuning.resolve_run_config(**inputs) + assert ssm_backend.run_config == default and not ssm_backend._config_is_fixed + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + selected = tuning.resolve_run_config(**inputs) + assert selected == (default if level == 3 else config) + assert ssm_backend.run_config == default and not ssm_backend._config_is_fixed + # Reusing the same bucket must not tune again. + monkeypatch.setattr(tuning, "select_config", lambda *args, **kwargs: pytest.fail("layout changed after capture")) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + tuning.resolve_run_config(**inputs) + assert not tuner.warmuped_configs_set + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "mode,width,compute_mode", + [ + ("compact", 4, "inline"), + ("replay", 4, "inline"), + ("replay", 1, "inline"), + ("replay", 4, "precompute"), + ], +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_joint_tuning_isolated_state_and_graph_replay(monkeypatch, mode, dtype, width, compute_mode, kda): + inputs = make_inputs(mode, dtype, "cuda", width=width, compute_mode=compute_mode, kda=kda) + if width == 1: + for name in ("q", "k", "v"): + inputs[name] = inputs[name].transpose(0, 1) + inputs["cu_seqlens"] = None + ssm_backend = inputs["ssm_backend"] + for name, x in tensors(ssm_backend).items(): + if name != "state": + x.zero_() + before = {name: x.clone() for name, x in tensors(ssm_backend).items()} + tuner = tuning.select_config + configs = [{"BV": 16, "num_warps": 1, "num_stages": 1}, {"BV": 32, "num_warps": 2, "num_stages": 2}] + monkeypatch.setattr(tuner, "configs_gen_func", lambda: configs) + bench = tuner._bench + timings = [] + + def checked_bench(*args, **kwargs): + assert kwargs["ssm_backend"] is not ssm_backend + assert kwargs["ssm_backend"].state.shape[0] == 1 + elapsed = bench(*args, **kwargs) + assert math.isfinite(elapsed), f"failed {mode} {kwargs['run_config']}" + if mode == "replay": + # The measured cycle must reach a fold, not just empty-history updates. + scratch = kwargs["ssm_backend"] + assert scratch.cursors[0].item() >= 2 * scratch.capacity + timings.append(elapsed) + for name, x in tensors(ssm_backend).items(): + torch.testing.assert_close(x, before[name], atol=0, rtol=0) + return elapsed + + monkeypatch.setattr(tuner, "_bench", checked_bench) + args = [inputs[name] for name in ("q", "k", "v", "a", "b", "a_log", "bias")] + reqs = torch.full((3,), ssm_backend.hold_req_id, device="cuda", dtype=torch.int32) + ssm_backend.prepare_decode(reqs, inputs["cu_seqlens"]) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + # Exercise the real forward hook with the HOLD-only graph warmup input. + ssm_backend.forward(0, *args, reqs, inputs["cu_seqlens"]) + assert len(timings) == len(configs) + selected = ssm_backend.get_run_config(*args, inputs["cu_seqlens"]) + assert asdict(selected) in configs + # An explicit identical layout supplies a reference for capture/replay and + # alternating real/HOLD rows. It must never trigger its own tuning. + expected = make_inputs(mode, dtype, "cuda", width=width, compute_mode=compute_mode, kda=kda)["ssm_backend"] + expected.run_config = ssm_backend.run_config + expected._config_is_fixed = True + expected.state.copy_(before["state"]) + accepted = torch.zeros(ssm_backend.state.shape[1], device="cuda", dtype=torch.int32) + + def step(target): + target.prepare_decode(reqs, inputs["cu_seqlens"]) + out = [ + target.forward(layer, *args, reqs, inputs["cu_seqlens"], run_config=selected)[0] + for layer in range(target.state.shape[0]) + ] + if mode == "compact": + target.accept_updates(reqs, accepted, run_config=selected) + else: + target.accept_updates(reqs, accepted) + return out + + reqs.fill_(ssm_backend.hold_req_id) + step(ssm_backend) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual_out = step(ssm_backend) + reqs.copy_(torch.tensor([2, 0, ssm_backend.hold_req_id], device="cuda", dtype=torch.int32)) + for iteration in range(12): + accepted[0] = iteration % width + accepted[2] = (iteration + 1) % width + graph.replay() + expected_out = step(expected) + for actual, ref in zip(actual_out, expected_out): + torch.testing.assert_close(actual, ref, atol=0, rtol=0) + torch.testing.assert_close(ssm_backend.state, expected.state, atol=0, rtol=0) + if mode == "replay": + cursor = ssm_backend.cursors.clone() + snapshot = ssm_backend.get_accepted_state(0) + torch.testing.assert_close(ssm_backend.cursors, cursor, atol=0, rtol=0) + torch.testing.assert_close(snapshot, expected.get_accepted_state(0), atol=0, rtol=0) + ssm_backend.merge_accepted_updates(reqs[:2]) + expected.merge_accepted_updates(reqs[:2]) + torch.testing.assert_close(ssm_backend.state, expected.state, atol=0, rtol=0) + torch.testing.assert_close( + ssm_backend.state[:, ssm_backend.hold_req_id], before["state"][:, ssm_backend.hold_req_id], atol=0, rtol=0 + ) + assert len(timings) == len(configs) + + +@pytest.mark.parametrize("mode", ["compact", "replay"]) +def test_kda_rule_and_gate_bound_separate_tuning_cache(mode): + inputs = make_inputs(mode, kda=True) + key = tuning.select_config._static_key(**inputs) + assert key != tuning.select_config._static_key(**make_inputs(mode)) + _, rebuilt = tuning.rebuild_inputs(**inputs) + assert rebuilt["ssm_backend"].kda + assert rebuilt["ssm_backend"].gate_lower_bound == inputs["ssm_backend"].gate_lower_bound + assert key == tuning.select_config._static_key(**rebuilt) + inputs["ssm_backend"].gate_lower_bound = -2.5 + assert key != tuning.select_config._static_key(**inputs) + + +@pytest.mark.parametrize("tp", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_production_tuning_cache_can_be_written(monkeypatch, tmp_path, tp, dtype, kda): + from types import SimpleNamespace + + monkeypatch.setattr(kernel_config_module, "get_current_device_name", lambda: "NVIDIA H100 80GB HBM3 MIG 3g.40gb") + heads = 64 // tp + ssm_backend = SimpleNamespace( + state=SimpleNamespace(shape=(34, 1, heads, 128, 128), dtype=dtype), + keys=SimpleNamespace(dtype=torch.bfloat16), + verify_width=3, + num_key_heads=heads, + capacity=8, + kda=kda, + gate_lower_bound=-5.0, + compute_mode="precompute", + ) + packed = torch.empty(1, 3, heads * 128 * 3, dtype=torch.bfloat16) + q, k, v = [x.view(1, 3, heads, 128) for x in packed.chunk(3, -1)] + a = torch.empty(3, heads * 128 if kda else heads, dtype=q.dtype) + b = torch.empty(3, heads, dtype=q.dtype) + key = tuning.static_key(ssm_backend, "replay", q, k, v, a, b, torch.tensor([0, 3])) + path = tmp_path / KernelConfigs.get_config_file_name(key) + path.write_text("{}") + assert path.is_file() + + +def test_batch_configs_survive_other_warmups_and_preserve_padding(monkeypatch): + inputs = make_inputs("compact") + ssm_backend = inputs["ssm_backend"] + small = dict(inputs, **{name: inputs[name][:, :4] for name in ("q", "k", "v")}) + small["a"], small["b"] = inputs["a"][:4], inputs["b"][:4] + small["cu_seqlens"] = torch.tensor([0, 4, 4], dtype=torch.int32) + selected = [] + + def choose(ssm_backend, mode, q, *args): + config = {"BV": 16 if q.shape[1] == 4 else 32, "num_warps": 2} + selected.append(config) + return config + + monkeypatch.setattr(tuning, "select_config", choose) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + large_config = tuning.resolve_run_config(**inputs) + small_config = tuning.resolve_run_config(**small) + assert tuning.resolve_run_config(**inputs) is large_config + assert len(selected) == 2 and large_config != small_config + assert tuning.resolve_run_config(**small) is small_config + assert ssm_backend.run_config == {"BV": 8, "num_warps": 1} + _, rebuilt = tuning.rebuild_inputs(**small) + assert rebuilt["workload"][0].tolist() == [0, rebuilt["ssm_backend"].hold_req_id] + assert rebuilt["cu_seqlens"].tolist() == [0, 4, 4] + assert tuning.run_key(**{k: rebuilt[k] for k in ("ssm_backend", "q", "cu_seqlens")}) == (4, 2) + # A graph token bucket can also exceed the request-pool capacity. + _, limited = tuning.rebuild_inputs(**dict(inputs, cu_seqlens=torch.tensor([0, 4, 8], dtype=torch.int32))) + assert limited["workload"][0].tolist() == [0, 1] + assert limited["cu_seqlens"].tolist() == [0, 4, 8] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("mode,projection", [("compact", "inline"), ("replay", "inline"), ("replay", "precompute")]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_two_verify_graphs_accept_with_their_own_configs(monkeypatch, mode, projection, dtype, kda, dynamic): + from types import SimpleNamespace + import weakref + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttDecodeAttState + from lightllm.common.basemodel.attention.linear.kda import KDALinearAttDecodeAttState + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.batch_objs import PostLayerOutput, ModelMtpOutputCollector + + inputs = make_inputs(mode, dtype, "cuda", compute_mode=projection, kda=kda) + ssm_backend = inputs["ssm_backend"] + initial = ssm_backend.state.clone() + configs = [{"BV": 8, "num_warps": 1, "num_stages": 1}, {"BV": 32, "num_warps": 2, "num_stages": 1}] + monkeypatch.setattr(tuning, "select_config", lambda *args: configs[args[2].shape[1] > 8]) + from lightllm.common.attention_state.bounded_att import KDAState, GDNState + + manager = object.__new__(KDAState if kda else GDNState) + manager._ssm = ssm_backend + manager.state_name = "main_kda" if kda else "main_gdn" + manager.mtp_state_index = torch.zeros(9, dtype=torch.int32, device="cuda") + buckets = [] + for index, ids in enumerate(([0], [1, 3])): + # The physical verify buckets include HOLD, but acceptance is unpadded. + count = len(ids) + 1 + args = [inputs[name] for name in ("q", "k", "v", "a", "b", "a_log", "bias")] + args[:3] = [x[:, : count * 4].contiguous() for x in args[:3]] + args[3:5] = [x[: count * 4].contiguous() for x in args[3:5]] + capacity = min(count * 4, ssm_backend.hold_req_id) if dynamic else count + reqs = torch.tensor(ids + [ssm_backend.hold_req_id] * (capacity - len(ids)), dtype=torch.int32, device="cuda") + cu = (torch.arange(capacity + 1, dtype=torch.int32, device="cuda") * 4).clamp_max(count * 4) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + config = ssm_backend.get_run_config(*args, cu) + assert asdict(config) == configs[index] + ref = make_inputs(mode, dtype, "cuda", compute_mode=projection, kda=kda)["ssm_backend"] + ref.state.copy_(initial) + ref.run_config = asdict(config) + ref._config_is_fixed = True + + def forward(target): + target.prepare_decode(reqs, cu) + outputs = [] + selected = None + for layer in range(2): + result, actual_config = target.forward(layer, *args, reqs, cu, run_config=selected) + if selected is not None: + assert actual_config is selected + selected = actual_config + outputs.append(result) + return outputs, selected + + forward(ssm_backend) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs, config = forward(ssm_backend) + state = InferStateInfo() + state.att_state_manager = SimpleNamespace(get_bounded_attention_states=lambda role: (manager,)) + state.hidden_collector = SimpleNamespace(finish_output=lambda **kw: ModelMtpOutputCollector()) + att_class = KDALinearAttDecodeAttState if kda else GDNLinearAttDecodeAttState + state.decode_att_state1 = att_class( + _infer_state_weak_ref=weakref.ref(state), + forwards={manager.state_name: SSMForward(manager, reqs, reqs, cu, None, config)}, + ) + output = TpPartBaseModel._create_model_output( + None, PostLayerOutput(logits=outputs[0].view(count * 4, -1)), state + ) + commits = output.state_commits + assert bool(commits) == (mode == "compact") + output = TpPartBaseModel._create_unpad_decode_model_output(None, output, len(ids) * 4) + assert output.state_commits is commits + buckets.append((graph, outputs, output, ref, args, reqs, cu, ids)) + # Published contexts and captured launches must survive later changes to + # the tuner's dictionaries, including when two graphs replay alternately. + for config in configs: + config.update(BV=128, num_warps=4, num_stages=3) + ssm_backend.state.copy_(initial) + if mode == "replay": + ssm_backend.cursors.zero_() + req_order = [0, 1, 3] + rows = torch.tensor([r for r in req_order for _ in range(4)], dtype=torch.int32, device="cuda") + mtp = torch.arange(4, dtype=torch.int32, device="cuda").repeat(3) + starts = torch.arange(3, dtype=torch.int32, device="cuda") * 4 + for iteration in range(24): + # Alternate replay order; neither Graph replay executes Python config selection. + for bucket_index in [0, 1] if iteration % 2 else [1, 0]: + graph, outputs, output, ref, args, reqs, cu, ids = buckets[bucket_index] + graph.replay() + ref.prepare_decode(reqs, cu) + for layer in range(2): + expected = ref.forward(layer, *args, reqs, cu)[0] + torch.testing.assert_close(outputs[layer], expected, atol=0, rtol=0) + counts = [1 + (iteration + i) % 4 for i in range(3)] + flags = torch.tensor([int(j < n) for n in counts for j in range(4)], dtype=torch.int32, device="cuda") + context = MtpCommitContext( + starts, rows, mtp, flags, 4, tuple(StateCommitBatch(len(b[-1]), b[2].state_commits) for b in buckets) + ) + manager.commit_bounded_att_state(rows[starts.long()], torch.tensor(counts, device="cuda"), context=context) + for _, _, _, ref, _, reqs, _, ids in buckets: + ref.accept_updates(reqs[:-1], manager.mtp_state_index) + torch.testing.assert_close(ssm_backend.state[:, ids], ref.state[:, ids], atol=0, rtol=0) + if mode == "replay": + assert torch.equal(ssm_backend.cursors[ids], ref.cursors[ids]) + assert len(ssm_backend.batch_configs) == 2 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..8ba764390e --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -0,0 +1,69 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning +from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState + + +@pytest.mark.parametrize("width", [1, 5]) +def test_scratch_budget_is_independent_of_model_layer_count(width): + state = torch.empty((48, 257, 48, 128, 128), dtype=torch.bfloat16, device="meta") + ssm_backend = ReplaySSMState(state, 16, width, num_key_heads=16) + tokens = 256 * width + inputs = [ + torch.empty(shape, dtype=torch.bfloat16, device="meta") + for shape in [(tokens, 16, 128), (tokens, 16, 128), (tokens, 48, 128), (tokens, 48), (tokens, 48), (48,), (48,)] + ] + required = tuning.scratch_size_bytes(ssm_backend, 256, 256, inputs) + assert required < 1024 ** 3 + single_layer = ReplaySSMState(state[:1], 16, width, num_key_heads=16) + assert required == tuning.scratch_size_bytes(single_layer, 256, 256, inputs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("width", [1, 5]) +def test_tuning_budget_falls_back_without_touching_serving_state(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 = ReplaySSMState(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 if width > 1 else None + if cu is not None: + q = q.unsqueeze(0) + inputs = dict( + ssm_backend=ssm_backend, mode="replay", q=q, k=q, v=q, a=gates, b=gates, a_log=a_log, bias=bias, cu_seqlens=cu + ) + before = {name: x.clone() for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device: (0, 1024 ** 3)) + monkeypatch.setattr(torch.cuda, "memory_reserved", lambda device: 0) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda device: 0) + tuner = tuning.select_config + monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", {}) + monkeypatch.setattr(tuner, "_oom_keys", set()) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: 1) + monkeypatch.setattr( + tuner, "_bench", lambda *args, **kwargs: pytest.fail("Insufficient memory must skip benchmarks") + ) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + config = tuning.resolve_run_config(**inputs) + assert tuning.resolve_run_config(**inputs) is config + assert config == {"BV": 32, "num_warps": 1} + assert not list(tmp_path.iterdir()) + for name, expected in before.items(): + torch.testing.assert_close(getattr(ssm_backend, name), expected, atol=0, rtol=0) + + reqs = torch.arange(3, device="cuda", dtype=torch.int32) + ssm_backend.prepare_decode(reqs, cu) + ssm_backend.forward(0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) + torch.cuda.synchronize() diff --git a/unit_tests/common/test_autotuner.py b/unit_tests/common/test_autotuner.py index 547246cfe5..5eb12f9e57 100644 --- a/unit_tests/common/test_autotuner.py +++ b/unit_tests/common/test_autotuner.py @@ -484,3 +484,61 @@ def test_default_api_and_nested_phase_restore_after_exception(): assert not Autotuner.is_autotune_warmup() assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + + +@pytest.mark.parametrize("world_size,failed_rank", [(1, 0), (2, 0), (2, 1)]) +@pytest.mark.parametrize("cached", [False, True]) +def test_rebuild_oom_skips_search_on_all_ranks_without_saving_fallback( + tmp_path, monkeypatch, world_size, failed_rank, cached +): + kernel, calls, benchmarks, cache_file = make_kernel( + tmp_path, monkeypatch, "scratch_oom", AutotuneKernelType.DECODE_ATTENTION, warmup_all_exist_config=False + ) + if cached: + cache_file.write_text(json.dumps({"8": {"block": 1}})) + before = cache_file.read_bytes() if cached else None + rebuilds = [] + + def rebuild(size, run_config=None): + rebuilds.append(size) + if failed_rank == 0: + raise torch.cuda.OutOfMemoryError("scratch allocation") + return (size,), {} + + monkeypatch.setattr(kernel, "rebuild_input_func", rebuild) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: world_size > 1) + monkeypatch.setattr(autotuner_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(autotuner_module, "get_global_world_size", lambda: world_size) + monkeypatch.setattr(kernel, "_get_autotune_group", lambda: None) + collectives = [] + + def all_gather_object(outputs, obj, group): + collectives.append(obj) + if len(collectives) == 3 and isinstance(obj, bool): + outputs[:] = [rank != failed_rank for rank in range(world_size)] + else: + outputs[:] = [obj] * world_size + + monkeypatch.setattr(autotuner_module.dist, "all_gather_object", all_gather_object) + fallback = {"block": 1} if cached else None + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == fallback + assert kernel(16) == fallback + assert rebuilds == [16] + assert calls == [(16, fallback), (16, fallback)] + assert not benchmarks + assert (cache_file.read_bytes() if cache_file.exists() else None) == before + assert "16" not in kernel.cached_configs.get(frozendict({"dtype": "test"}), {}) + if world_size > 1: + assert len(collectives) == 4 # need, key, allocation status, repeated need + + +def test_rebuild_programming_error_is_not_treated_as_oom(tmp_path, monkeypatch): + kernel, _, _, _ = make_kernel(tmp_path, monkeypatch, "rebuild_error") + + def rebuild(size, run_config=None): + raise ValueError("invalid scratch layout") + + monkeypatch.setattr(kernel, "rebuild_input_func", rebuild) + with Autotuner.autotune_warmup(), pytest.raises(ValueError, match="invalid scratch layout"): + kernel(16) 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..cd0006dd45 --- /dev/null +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -0,0 +1,132 @@ +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 + if state._ssm is not None: + assert state._ssm.state is state.ssm_state + 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) + actual = state.get_accepted_state(2) + 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._ssm.cursors[2] == 0 + assert state._ssm.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", "compact", "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 + assert copy._ssm is state._ssm + 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( + 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._ssm.raw_keys.zero_() + state._ssm.raw_values.zero_() + state._ssm.gates.fill_(-0.5) + state._ssm.betas.zero_() + state._ssm.cursors[1], state._ssm.cursors[2] = 1, 3 + expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] + before = vars(state).copy() + source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, state.ssm_slots_per_req].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[:, state.ssm_slots_per_req], 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..8eb25c073d --- /dev/null +++ b/unit_tests/common/test_linear_state.py @@ -0,0 +1,155 @@ +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, + ) + state = cls(config, 4, mtp_step=step, device=device, ssm_state_mode=mode, ssm_replay_config={"cache_len": 4}) + 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), ("compact", 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_state_index is not None: + state.mtp_state_index[src] = offset + expected_conv = state.conv_state[:, src, ..., offset : offset + 3].clone() + expected_ssm = state.ssm_state[:, src * state.ssm_slots_per_req + (offset if mode == "native" else 0)].clone() + if mode == "replay": + state._ssm.raw_keys.zero_() + state._ssm.raw_values.zero_() + state._ssm.gates.fill_(-0.5) + state._ssm.betas.zero_() + state._ssm.cursors[src] = 1 + expected_ssm = (expected_ssm.float() * torch.exp(torch.tensor(-0.5))).to(dtype) + 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._ssm.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), + ): + if state.mtp_state_index is not None: + state.mtp_state_index[dst] = step + if mode == "replay": + state._ssm.cursors[dst] = 3 + restore() + actual = state.get_accepted_state(dst) + 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_state_index is not None: + assert state.mtp_state_index[dst].item() == 0 + if mode == "replay": + assert state._ssm.cursors[dst].item() == 0 + state.reset_bounded_att_state([dst]) + assert not state.conv_state[:, dst].any() + start = dst * state.ssm_slots_per_req + assert not state.ssm_state[:, start : start + state.ssm_slots_per_req].any() + + +def copy_on_cpu(*, copies): + for source, destination in copies: + destination.copy_(source) + + +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.conv_ssm_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, 1:4].fill_(10 + rank) + state.conv_state[:, src, 2:4, 1:4].fill_(20 + rank) + state.conv_state[:, src, 4:, 1:4].fill_(30 + rank) + state.ssm_state[:, src * state.ssm_slots_per_req + 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) 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..391a0aa598 --- /dev/null +++ b/unit_tests/common/test_linear_state_contract.py @@ -0,0 +1,329 @@ +import pickle +from dataclasses import FrozenInstanceError +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.ssm_types import SSMRunConfig +from lightllm.common.attention_state.commit import MtpCommitContext, StateCommit, StateCommitBatch +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.models.glm5_next.config import Glm5NextResourceConfig +from lightllm.server.core.objs.start_args_type import StartArgs + + +def test_published_config_is_detached_immutable_and_pickleable(): + tuning = {"BV": 8, "num_warps": 1} + config = SSMRunConfig(**tuning) + tuning.update(BV=128, num_warps=4, num_stages=1) + assert config == SSMRunConfig(BV=8, num_warps=1, num_stages=3) + with pytest.raises(FrozenInstanceError): + config.BV = 128 + restored = pickle.loads(pickle.dumps(config)) + assert restored == config + + +def test_config_rejects_unrepresented_kernel_options(): + with pytest.raises(TypeError, match="unexpected keyword argument"): + SSMRunConfig(**{"BV": 32, "num_warps": 1, "future_layout_option": 2}) + with pytest.raises(TypeError, match="must be an integer"): + SSMRunConfig(BV=[32], num_warps=1) + + +@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 + + +@pytest.mark.parametrize("mode", ["compact", "replay"]) +def test_kda_backend_receives_gate_bound_at_construction(linear_config, mode): + owner = KDAState(linear_config, 3, mtp_step=2, ssm_state_mode=mode, gate_lower_bound=-2.5, device="cpu") + assert owner._ssm.gate_lower_bound == -2.5 + + +@pytest.mark.parametrize("mode", ["compact", "replay", "parallel"]) +def test_ssm_backend_publishes_detached_run_config(mode): + from lightllm.common.attention_state.bounded_att.ssm.replay import ReplaySSMState + from lightllm.common.attention_state.bounded_att.ssm.compact import CompactSSMState + + ssm_backend = object.__new__(CompactSSMState if mode == "compact" else ReplaySSMState) + tuning = {"BV": 32, "num_warps": 2} + ssm_backend.run_config = tuning + ssm_backend._config_is_fixed = True + ssm_backend.compute_mode = "parallel" if mode == "parallel" else "inline" + config = ssm_backend.get_run_config(*([None] * 7)) + assert isinstance(config, SSMRunConfig) + tuning["BV"] = 128 + assert config.BV == 32 + assert config.num_stages == 3 + with pytest.raises(FrozenInstanceError): + config.BV = 128 + + +def make_commit_context(batches): + return MtpCommitContext( + torch.tensor([0, 4, 8]), + torch.tensor([2] * 4 + [0] * 4 + [1] * 4), + torch.arange(4).repeat(3), + torch.ones(12, dtype=torch.int32), + 4, + batches, + ) + + +def test_compact_commits_each_unpadded_partition_with_its_verify_config(monkeypatch): + from lightllm.common.attention_state.bounded_att.ssm import base + + owner, _, _ = make_state(mode="compact", step=3, device="cpu") + reqs = torch.tensor([2, 0, 1], dtype=torch.int32) + configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] + commits = [] + for subset, config in zip((reqs[:1], reqs[1:]), configs): + forward = owner.prepare_decode(subset, torch.arange(subset.numel() + 1) * 4) + forward.run_config = config + commits.append(forward.finish()) + owner._ssm.run_config = {"BV": 128, "num_warps": 4} + calls = [] + monkeypatch.setattr(base, "linear_att_mtp_state_index_update", lambda **kwargs: None) + monkeypatch.setattr( + owner._ssm, + "accept_updates", + lambda reqs, accepted, run_config: calls.append((reqs.tolist(), accepted, run_config)), + ) + context = make_commit_context( + (StateCommitBatch(1, (commits[0],)), StateCommitBatch(0), StateCommitBatch(2, (commits[1],))) + ) + owner.commit_bounded_att_state(reqs, torch.tensor([1, 2, 3]), context=context) + assert [reqs for reqs, _, _ in calls] == [[2], [0, 1]] + assert all(accepted is owner.mtp_state_index for _, accepted, _ in calls) + assert [config for _, _, config in calls] == configs + + +@pytest.mark.parametrize("failure", ["missing", "wrong_owner", "mutable_config", "no_context"]) +def test_invalid_compact_receipt_rejected_before_state_changes(monkeypatch, failure): + owner, _, _ = make_state(mode="compact", step=3, device="cpu") + owner.mtp_state_index.fill_(2) + name = "draft0_other_state" if failure == "wrong_owner" else owner.state_name + config = {"BV": 8, "num_warps": 1} if failure == "mutable_config" else SSMRunConfig(8, 1) + commits = () if failure == "missing" else (StateCommit(name, config),) + context = None if failure == "no_context" else make_commit_context((StateCommitBatch(3, commits),)) + with pytest.raises(ValueError, match="compact"): + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), context=context) + assert owner.mtp_state_index.tolist() == [2] * 4 + + +def test_commit_partitions_and_resource_identity_are_validated(): + with pytest.raises(ValueError, match="nonnegative"): + StateCommitBatch(-1) + commit = StateCommit("main", SSMRunConfig(8, 1)) + with pytest.raises(ValueError, match="duplicate"): + StateCommitBatch(1, (commit, commit)) + context = make_commit_context((StateCommitBatch(3, (commit,)),)) + with pytest.raises(ValueError, match="unknown state"): + context.validate(3, {"draft0"}) + with pytest.raises(ValueError, match="counts differ"): + context.validate(2, {"main"}) + with pytest.raises(ValueError, match="partition"): + make_commit_context((StateCommitBatch(2),)).validate(3, {"main"}) + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_native_accept_restore_and_reuse_preserve_other_requests(cls): + owner, _, _ = make_state(cls, step=2, device="cpu") + owner.conv_state.copy_(torch.arange(5)) + for slot in range(owner.ssm_state.shape[1]): + owner.ssm_state[:, slot].fill_(slot) + before = owner.ssm_state.clone() + owner.commit_bounded_att_state(torch.tensor([1]), torch.tensor([3], dtype=torch.int32)) + accepted = owner.get_accepted_state(1) + assert torch.equal(accepted.conv[0, 0], torch.tensor([2, 3, 4], dtype=accepted.conv.dtype)) + assert torch.all(accepted.ssm == 5) + owner.restore(2, *accepted) + restored = owner.get_accepted_state(2) + assert torch.equal(restored.conv, accepted.conv) + assert torch.equal(restored.ssm, accepted.ssm) + owner.reset_bounded_att_state([2]) + assert torch.equal(owner.ssm_state[:, :6], before[:, :6]) + assert torch.equal(owner.ssm_state[:, 9:], before[:, 9:]) + assert not owner.ssm_state[:, 6:9].any() + 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): + from lightllm.common.attention_state.bounded_att.ssm import base + + owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") + calls = [] + monkeypatch.setattr(base, "linear_att_mtp_state_index_update", lambda **kwargs: calls.append(kwargs)) + monkeypatch.setattr(owner._ssm, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), context=make_commit_context(())) + assert len(calls) == 1 + assert calls[0]["replay_cursors"] is owner._ssm.cursors + assert calls[0]["req_to_mtp_state_index"] is owner.mtp_state_index + + +@pytest.mark.parametrize("change", ["owner", "shape", "optional", "dtype"]) +def test_graph_layout_mismatch_rejected_before_copy(change): + from copy import copy + + owner, _, _ = make_state(mode="native", device="cpu") + target = owner.prepare_decode(torch.tensor([0, 1], dtype=torch.int32), torch.tensor([0, 3, 6])) + source = owner.prepare_decode(torch.tensor([1, 2], dtype=torch.int32), torch.tensor([0, 2, 5])) + if change == "owner": + source.resource = copy(owner) + elif change == "shape": + source.state_indices = source.state_indices[:1] + elif change == "optional": + source.accepted_token_counts = None + else: + source.cu_seqlens = source.cu_seqlens.int() + with pytest.raises(ValueError, match="resource|layout"): + target.copy_inputs(source) + assert target.req_indices.tolist() == [0, 1] + assert target.cu_seqlens.tolist() == [0, 3, 6] + + +def test_forward_contexts_select_each_layer_resource_and_publish_both_commits(): + import weakref + from dataclasses import replace + from lightllm.common.attention_state import AttentionStateManager + from lightllm.common.basemodel.attention.linear.gdn import GDNLinearAttDecodeAttState + from lightllm.common.basemodel.batch_objs import PostLayerOutput + from lightllm.common.basemodel.basemodel import TpPartBaseModel + from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector + from lightllm.common.basemodel.infer_struct import InferStateInfo + + first, _, _ = make_state(mode="compact", device="cpu") + second = GDNState( + replace(first.config, logical_layer_ids=(3,), ssm_dtype=torch.bfloat16), + 4, + mtp_step=2, + device="cpu", + ssm_state_mode="compact", + ) + resources = AttentionStateManager(request_capacity=4) + resources.register_state("main", first) + resources.register_state("main", second) + resources.register_state("draft0", GDNState(first.config, 4, mtp_step=2, device="cpu")) + resources.finish_state_registration() + infer = InferStateInfo() + infer.batch_size = 6 + infer.b_req_idx = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.int32) + infer.att_state_manager = resources + infer.hidden_collector = NoopHiddenCollector() + backend = SimpleNamespace( + model=SimpleNamespace(is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2)), + uses_dynamic_spec_verify_layout=lambda: False, + ) + att = GDNLinearAttDecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + infer.decode_att_state1 = att + att.init_state() + assert len(att.forwards) == 2 + assert att.get_forward(0) is att.get_forward(1) + assert att.get_forward(3).resource is second + configs = [SSMRunConfig(8, 1), SSMRunConfig(32, 2)] + att.get_forward(0).run_config, att.get_forward(3).run_config = configs + output = TpPartBaseModel._create_model_output(None, PostLayerOutput(torch.ones(6, 1)), infer) + assert {c.state_name: c.payload for c in output.state_commits} == { + first.state_name: configs[0], + second.state_name: configs[1], + } + + +@pytest.mark.parametrize("kda", [False, True]) +@pytest.mark.parametrize("speculative", [False, True]) +def test_native_forward_preserves_algorithm_inputs(monkeypatch, kda, speculative): + from importlib import import_module + + owner, _, _ = make_state(KDAState if kda else GDNState, step=2 if speculative else 0, device="cpu") + module, function = ( + ("fla.ops.kda_decode", "fused_recurrent_kda") + if kda + else ("mtp_fused_recurrent", "mtp_fused_recurrent_gated_delta_rule") + if speculative + else ("fla.ops", "fused_recurrent_gated_delta_rule") + ) + calls = [] + + def recurrent(**kwargs): + calls.append(kwargs) + return kwargs["v"], None + + monkeypatch.setattr( + import_module("lightllm.common.basemodel.triton_kernel.linear_att." + module), function, recurrent + ) + reqs = torch.tensor([1, 2], dtype=torch.int32) + cu = torch.tensor([0, 3, 6], dtype=torch.int32) if speculative else None + forward = owner.prepare_decode(reqs, cu) + shape = (1, 6) if speculative else (2, 1) + c = owner.config + q = torch.zeros(*shape, c.num_k_heads, c.key_dim) + v = torch.zeros(*shape, c.num_v_heads, c.value_dim) + gate = torch.zeros(q.numel() // c.key_dim, c.num_v_heads * (c.key_dim if kda else 1)) + beta = torch.zeros(gate.shape[0], c.num_v_heads) + a_log, bias = torch.zeros(c.num_v_heads), torch.zeros(gate.shape[-1]) + assert forward.run_layer(1, q, q, v, gate, beta, a_log, bias) is v + call = calls[0] + assert call["initial_state"].data_ptr() == owner.ssm_state[1].data_ptr() + assert call["ssm_state_indices"] is forward.state_indices + if speculative: + assert torch.equal(call["cu_seqlens"], cu) + assert call["num_accepted_tokens"] is forward.accepted_token_counts + if kda: + assert call["raw_gate"].shape == (*shape, c.num_v_heads * c.key_dim) + assert call["raw_beta"].shape == (*shape, c.num_v_heads) + assert call["lower_bound"] == -5.0 + else: + assert call["a_raw"] is gate and call["b_raw"] is beta + assert forward.finish() is None diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index c07fd123f0..20d59199bc 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -9,6 +9,19 @@ import pytest +@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..fa39331c26 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -48,12 +48,16 @@ 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", "compact", "replay"]) def test_hybrid_checkpoint_restore_and_packed_kv_copy( monkeypatch, small_page, tp_world_size, mtp_step, + mode, ): + if mode == "compact" and mtp_step == 0: + pytest.skip("compact requires MTP") monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) @@ -106,6 +110,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 +148,16 @@ 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._ssm.raw_keys.zero_() + linear._ssm.raw_values.zero_() + linear._ssm.gates.fill_(-0.5) + linear._ssm.betas.zero_() + linear._ssm.cursors[0] = 1 + ssm = linear._ssm.get_accepted_state(0).clone() + linear._ssm.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 +177,16 @@ 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) + assert torch.equal(linear.ssm_state[:, 2 * linear.ssm_slots_per_req], ssm) + if mode == "replay": + assert linear._ssm.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() + assert not linear.ssm_state[:, 2 * linear.ssm_slots_per_req : 3 * linear.ssm_slots_per_req].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_indexer_topk.py b/unit_tests/models/glm5_next/test_indexer_topk.py index 76646c5224..2c0e20675f 100644 --- a/unit_tests/models/glm5_next/test_indexer_topk.py +++ b/unit_tests/models/glm5_next/test_indexer_topk.py @@ -98,20 +98,19 @@ def test_topk_without_vllm(indexer, monkeypatch): @pytest.mark.parametrize("mtp_size", [1, 3, 6]) -@pytest.mark.parametrize("max_pools", [1024, 262144]) -def test_paged_decode_matches_nonpaged_with_graph_and_mtp(indexer, monkeypatch, mtp_size, max_pools): +@pytest.mark.parametrize("requests,max_pools", [(5, 640), (5, 1024), (5, 262144), (64, 640)]) +def test_paged_decode_matches_nonpaged_with_graph_and_mtp(indexer, monkeypatch, mtp_size, requests, max_pools): pytest.importorskip("deep_gemm") - _require_vllm_topk() torch.manual_seed(27) - rows = 5 * mtp_size + rows = requests * mtp_size storage = torch.empty(1024, 1, 584, device="cuda", dtype=torch.bfloat16) packed = storage.view(torch.uint8)[:, :, -132:] packed[:, 0, :128] = torch.randn(1024, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) scales = torch.pow(2.0, torch.randint(-3, 2, (1024,), device="cuda").float()) packed[:, 0, 128:] = scales.view(torch.uint8).view(-1, 4) - table = torch.full((5, max_pools * 4), -1, device="cuda", dtype=torch.int32) - table[:, 3::4] = torch.randint(0, 1024, (5, max_pools), device="cuda", dtype=torch.int32) - req_idx = torch.tensor([1, 3, 0, 2, 4], device="cuda", dtype=torch.int32).repeat_interleave(mtp_size) + table = torch.full((requests, max_pools * 4), -1, device="cuda", dtype=torch.int32) + table[:, 3::4] = torch.randint(0, 1024, (requests, max_pools), device="cuda", dtype=torch.int32) + req_idx = torch.arange(requests, device="cuda", dtype=torch.int32).roll(2).repeat_interleave(mtp_size) lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) q = torch.randn(rows, 32, 128, device="cuda").to(torch.float8_e4m3fn) weights = torch.randn(rows, 32, device="cuda") * 0.05 @@ -142,8 +141,10 @@ def run(): with torch.cuda.graph(graph): actual = run() graph_logits = logits_outputs.pop() + graph.replay() + assert (actual == -1).all() for bases in ([0, 2047, 2051, max_pools * 4 - mtp_size + 1, 0], [7, 65, 511, 2053, 0]): - seqs = torch.tensor(bases, device="cuda", dtype=torch.int32)[:, None] + seqs = torch.tensor(bases, device="cuda", dtype=torch.int32).repeat((requests + 4) // 5)[:requests, None] seqs = seqs + torch.arange(mtp_size, device="cuda", dtype=torch.int32)[None, :] seqs[-1].zero_() lengths.copy_(seqs.flatten()) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 82e029c5dd..01c899c497 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -127,17 +127,18 @@ 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), ("compact", 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.attention.linear.kda import KDALinearAttDecodeAttState 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)] @@ -145,10 +146,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) @@ -186,7 +187,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, @@ -205,11 +206,19 @@ 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) + linear = KDAState(config, request_capacity=4, mtp_step=mtp_step, device="cuda", ssm_state_mode=mode) + 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) @@ -225,7 +234,7 @@ 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)) + state = KDALinearAttDecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) state.init_state() actual = state.decode_att( None, @@ -239,6 +248,23 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): ), ) 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 mode == "native": + torch.testing.assert_close(selected.ssm_state[0], expected_ssm, atol=2e-3, rtol=1e-2) + else: + ids = torch.tensor(requests, dtype=torch.int32, device="cuda") + if mtp_step: + accepted = torch.zeros(4, dtype=torch.int32, device="cuda") + accepted[ids] = torch.tensor(lengths, dtype=torch.int32, device="cuda") - 1 + if mode == "compact": + selected._ssm.accept_updates(ids, accepted, run_config=state.get_forward(0).run_config) + else: + selected._ssm.accept_updates(ids, accepted) + for req, length in zip(requests, lengths): + torch.testing.assert_close( + selected._ssm.get_accepted_state(req)[0], + expected_ssm[req * width + length - 1], + atol=2e-3, + rtol=1e-2, + ) if dynamic: - assert torch.equal(conv[3], old_conv[3]) + 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 ff93312516..40f64b2639 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,23 @@ 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", "compact"]) 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) + if dest_mode == "compact" and decode_mtp_step == 0: + pytest.skip("compact requires MTP") + 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 = dest[0].linear.ssm_slots_per_req 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 +167,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]) + mem.linear.ssm_state[:, src_req * mem.linear.ssm_slots_per_req].copy_(ssm[:, heads]) + if source_mode == "replay": + backend = mem.linear._ssm + backend.raw_keys.zero_() + backend.raw_values.zero_() + backend.gates.fill_(-0.5) + backend.betas.zero_() + backend.cursors[src_req] = 1 + ssm[:, heads].copy_(backend.get_accepted_state(src_req)) offset = 0 for state in mem.tails: end = offset + len(state.config.logical_layer_ids) @@ -173,6 +189,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._ssm.cursors[dst_req] = 3 + if mem.linear.mtp_state_index is not None: + 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 +201,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 +255,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 dest_mode == "replay": + assert mem.linear._ssm.cursors[dst_req] == 0 if mem.linear.mtp_state_index is not None: assert mem.linear.mtp_state_index[dst_req] == 0 for (state, stride), expected in zip( ( (mem.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..a31c0acaf8 --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -0,0 +1,150 @@ +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() + assert resource.get_storage_layer_index(logical_layer) == storage_layer + 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.triton_kernel.linear_att.fla.ops.fused_recurrent_gated_delta_rule", recurrent + ) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent." + "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/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index e020c6709b..a2af828623 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -290,7 +290,12 @@ def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: model = TpPartBaseModel.__new__(TpPartBaseModel) logits = torch.randn(3, 2) ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) - state = SimpleNamespace(hidden_collector=NoopHiddenCollector(), prompt_logics=None) + state = SimpleNamespace( + hidden_collector=NoopHiddenCollector(), + prompt_logics=None, + decode_att_state=None, + decode_att_state1=None, + ) output = model._create_model_output(PostLayerOutput(logits, ids), state) unpadded = model._create_unpad_decode_model_output(output, 2) torch.testing.assert_close(unpadded.logits, logits[:2]) diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py index 230e23b0e5..4959a4ac65 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -102,3 +102,49 @@ 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), ("compact", 4), ("replay", 0)]) +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": + assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "inline", "fold_method": "raw"} + + +def test_cold_s3_model_still_enforces_hybrid_constraints(launch_args, monkeypatch, tmp_path): + from pathlib import Path + + monkeypatch.chdir(tmp_path) + + def prepare(model_dir): + Path(model_dir).mkdir(parents=True) + Path(model_dir, "config.json").write_text(json.dumps({"model_type": "qwen3_5_text"})) + + monkeypatch.setattr("lightllm.utils.petrel_helper.s3_model_prepare", prepare) + args = launch_args("qwen3_5_text", disable_dynamic_prompt_cache=True) + args.model_dir = "s3://test-bucket/qwen" + with pytest.raises(AssertionError, match="hybrid attention models require dynamic prompt cache"): + _launch_subprocesses(args) 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..57151db0ae --- /dev/null +++ b/unit_tests/server/test_ssm_state_mode.py @@ -0,0 +1,203 @@ +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), + ("compact", "bfloat16", 2, "qwen3_5", 16, None), + ("compact", "float32", 2, "qwen3_5", 16, None), + ("compact", "float32", 15, "qwen3_5", 16, None), + ("compact", "float32", 16, "qwen3_5", 16, "mtp_step <= 15"), + ("compact", "float32", 32, "qwen3_5", 16, "mtp_step <= 15"), + ("compact", "bfloat16", 15, "glm5_next", 16, None), + ("compact", "bfloat16", 16, "glm5_next", 16, "mtp_step <= 15"), + ("compact", "bfloat16", 2, "glm5_next", 16, 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), + ("compact", "bfloat16", 0, "qwen3_5", 16, "requires MTP"), + ("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"), + ("compact", "float32", 2, "llama", 16, "require a GDN"), + ("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", "compact", "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", "parallel"]) +@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) + args.ssm_state_mode = "compact" + with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("fold", ["raw", "derived"]) +def test_parallel_replay_configuration(monkeypatch, fold): + from lightllm.server.core.objs.start_args_type import StartArgs + + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = StartArgs( + ssm_state_mode="replay", + ssm_replay_config={"compute_mode": "parallel", "fold_method": fold}, + mtp_step=4, + mtp_mode="eagle_with_att", + linear_att_ssm_data_type="bfloat16", + ) + _validate_and_normalize_ssm_config(args) + args.mtp_step = 0 + with pytest.raises(AssertionError, match="requires MTP"): + _validate_and_normalize_ssm_config(args) + + +def test_derived_replay_requires_parallel_configuration(): + from lightllm.server.core.objs.start_args_type import StartArgs + + args = StartArgs(ssm_state_mode="replay", ssm_replay_config={"fold_method": "derived"}) + with pytest.raises(AssertionError, match="requires parallel"): + _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": "auto"}, "fold_method"), + ], +) +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", "compact"]) +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", "fold_method": "raw"} + 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"), + ("--replayssm_fold_method", "raw"), + ("--ssm_replay_config", "{invalid}"), + ], +) +def test_removed_flags_and_malformed_json(option, value): + with pytest.raises(SystemExit): + make_argument_parser().parse_args([option, value])